Logit-Aware MIMO AirComp for Distributed Mixture-of-Experts LLM Inference over Wireless Edge Networks

Lyutianyang Zhang; Yunjian Jia; Liu Cao; Dengke Wang; Jinke Ren; Shuguang Cui

cs.ITSep 22, 2026

Logit-Aware MIMO AirComp for Distributed Mixture-of-Experts LLM Inference over Wireless Edge Networks

Authors: Lyutianyang Zhang, Yunjian Jia, Liu Cao, Dengke Wang, Jinke Ren, Shuguang Cui

Organizations: School of Microelectronics and Communication Engineering, Chongqing University, Chongqing, China · School of Science and Engineering, the Shenzhen Future Network of Intelligence Institute, and the Guangdong Provincial Key Laboratory of Future Networks of Intelligence, the Chinese University of Hong Kong (Shenzhen), Shenzhen, Guangdong 518172, China

Abstract

Distributed mixture-of-experts (MoE) inference is a promising architecture for deploying large language models (LLMs) at wireless edge networks because sparse experts can be placed across coordinated base stations (BSs), while the anchor node and user equipment (UE) can offload LLM inference tasks to BSs. The communication bottleneck is the MoE aggregation, where the anchor BS must recover a weighted sum of selected expert outputs before each decoding step. Over-the-air computation (AirComp) is well matched to this operation because the wireless multiple-access channel naturally superposes simultaneous transmissions. However, conventional AirComp minimizes communication distortion, whereas MoE aggregation errors have unequal impact on LLM outputs. We propose a logit-aware MIMO AirComp framework that estimates local logit sensitivity as block-level weights and jointly optimizes receive combiners and BS precoders under per-BS power constraints. We also develop an alternating algorithm that protects decoding decisions from aggregation perturbations. Using OpenCompass, we evaluate Qwen3-30B-A3B-Instruct-2507-FP8 on GSM8K and ARC-Challenge. At 30 dB aggregation SNR, perturbed Qwen3 retains 99.3% of clean GSM8K accuracy and 94.8% of clean ARC-Challenge accuracy. In a direct ARC-Challenge closed-loop audit over 295 examples, SW-AirComp achieves 43.73% accuracy at 30 dB and 26.78% at 20 dB, outperforming unweighted, matched-filter, and zero-forcing AirComp. In the wireless simulator, SW-AirComp reduces decision-relevant aggregation distortion; at 20 dB, its final weighted-sum mean-squared error is 53.5% lower than unweighted AirComp under the same channels, samples, and sensitivity weights. Logit-RMSE and logit-gap diagnostics further show that the gain comes from reducing output-logit perturbation and lowering the risk of top-token changes.

Figures & tables

Fig. 1: System architecture of distributed MoE inference.

Panel of Figure 2 — Fig. 2: Construction of scalar AirComp weights from block-level PHY error and local next-token logit sensitivity.

Benchmark	Item	Prompt excerpt and choices	Gold answer
GSM8K	#0	Janet’s ducks lay 16 eggs per day. She eats three for breakfast and uses four for muffins. She sells the remainder at the farmers’ market for $2 each.	18
GSM8K	#1	A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?	3
ARC-c	#0	Juan and LaKeisha roll a few objects down a ramp. They want to see which object rolls the farthest. What should they do so they can repeat their investigation? Choices: A) Put the objects in groups; B) Change the height of the ramp; C) Choose different objects to roll; D) Record the details of the investigation.	D
ARC-c	#1	High-pressure systems stop air from rising into the colder regions of the atmosphere where water can condense. What will most likely result if a high-pressure system remains in an area for a long period of time? Choices: A) fog; B) rain; C) drought; D) tornado.	C

TABLE I: Examples of OpenCompass Subsets

Fig. 3: Qwen3 OpenCompass accuracy under controlled MoE aggregation perturbations.

Parameter	Value
Number of BSs	$M=4$
Active BSs per aggregation	$\|\mathcal{M}_{b}\|=2$
Transmit antennas per BS	$N_{\rm T}\in\{32,64,128\}$
Receive antennas	$N_{\rm R}=2$
System bandwidth	$20$ MHz
Maximum BS transmit power	$P_{m}^{\rm tot}=1$ W ( $30$ dBm)

TABLE II: Wireless Parameters

Fig. 4: ARC-Challenge Accuracy: AirComp Perturbations.

Fig. 6: Sensitivity-Weighted MSE vs SNR for Antenna Scaling.

Panel of Figure 9 — Fig. 7: Massive-MIMO aggregation study over 16 wireless aggregation blocks.

Fig. 8: Normalized WSMSE convergence.

Panel of Figure 11 — Fig. 9: (a) reports each baseline’s weighted-logit RMSE normalized by SW-AirComp. (b) reports the percentage of aggregation blocks whose estimated logit perturbation exceeds the clean top-token logit gap; lower is better.

LLM Inference

Large language model inference studies how models generate outputs efficiently, reliably, and securely across hardware and workloads.

See all 285 papers

AIR-LLM: Broadcasting AI Weights over Radio for Memory-Free Edge LLM Inference via RF Computing
Sep 30, 2026Zhihui Gao, Tingjun Chen, Dirk EnglundOn-Device Language Model Inference LLM Inference Acceleration
Massachusetts Institute of Technology Cambridge, MA, USA · Duke University Durham, NC, USA
OrderMoE: An expert similarity driven distributed edge MoE inference
Jul 19, 2026Xin Yuan, Ning Li, Quan Chen +2Edge Inference Mixture-of-Experts Inference
School of Ocean Engineering, Harbin Institute of Technology, Heilongjiang, China · School of Computer Science and Technology, Harbin Institute of Technology, Heilongjiang, China · School of Computer Science and Technology, Guangdong University of Technology, Guangdong, China +2
Scaling Multi-Node Mixture-of-Experts Inference Using Expert Activation Patterns
Apr 25, 2026Abhimanyu Bambhaniya, Geonhwa Jeong, Jason Park +6Expert Load Balancing Mixture-of-Experts Inference
1Georgia Institute of Technology, USA · 2Meta Platforms, Inc.

Sep 30, 2026J/K move · Enter open · S save

AIR-LLM: Broadcasting AI Weights over Radio for Memory-Free Edge LLM Inference via RF Computing

Zhihui Gao, Tingjun Chen, Dirk Englund

Massachusetts Institute of Technology Cambridge, MA, USA · Duke University Durham, NC, USA

On-Device Language Model Inference LLM Inference Acceleration Wireless Communications

Open paper AskPDF arXiv

Next-generation large language models (LLMs) are expanding from the cloud to ubiquitous edge devices. However, edge devices typically either lack the memory to store increasingly large LLM weights or, even with enough memory, spend unaffordable energy on loading the weights. This raises our question: can an edge device run an LLM without storing or loading its weights, but receive them over the air and consume them on the fly? Inspired by wireless broadcasting, we present AIR-LLM, an LLM inference architecture for edge devices, which is composed of: (i) a central radio (e.g., 5G base stations) that broadcasts the LLM weights into the air, and (ii) the edge user that receives the weights and completes the general matrix-vector multiplication (GEMV) of LLM inference directly in the radio frequency (RF) domain using RF mixers. To further shorten the airtime, AIR-LLM exploits MIMO spatial multiplexing and proposes an energy-efficient precoder-postcoder pair on the edge to calibrate its own wireless channel. Since the central radio stays user-unaware, AIR-LLM is user-scalable so that one broadcast serves unlimited users within its coverage. We implement AIR-LLM on the NVIDIA Sionna ray-traced channels of two real-world urban scenes and the profiling of a real RF mixer. With a WikiText-2 perplexity degradation of 4.0% on LLaMA-3.1-8B, AIR-LLM saves the energy by 157.7x/40.4x against the FP16 and weight-only quantization baselines; with 20 users, its airtime is 104.1x/26.0x shorter, respectively.

Abstract

LLM Inference

Explore similar work

AIR-LLM: Broadcasting AI Weights over Radio for Memory-Free Edge LLM Inference via RF Computing

OrderMoE: An expert similarity driven distributed edge MoE inference

Scaling Multi-Node Mixture-of-Experts Inference Using Expert Activation Patterns