Communication-Efficient Distributed On-Device LLM Inference Over Wireless Networks

The main contributions of this paper are summarized as follows.

1) We propose a novel distributed on-device LLM inference framework by employing tensor parallelism and AirComp. While tensor parallelism effectively distributes computational workloads across edge devices, its frequent all-reduce operations incur significant communication overhead, which offsets the computational benefits and becomes a major bottleneck for inference performance. To address this challenge, we develop a communication-efficient AirComp all-reduce approach by exploiting the signal superposition property of wireless multiple-access channels.

2) To utilize the heterogeneous computational capabilities of edge devices and mitigate communication distortions, we investigate a joint model assignment and transceiver optimization problem to minimize the average transmission MSE. The formulated mixed-timescale stochastic non-convex optimization problem is inherently intractable. Thus, we develop an efficient two-stage algorithm that decomposes the original problem into short-term transceiver optimization and long-term model assignment optimization subproblems. The resulting subproblems are further solved by employing SDR and stochastic SCA, respectively. The proposed algorithm requires no prior knowledge of channel statistics, and it converges almost surely to a stationary point of the original problem.

3) We validate the effectiveness of the proposed framework through simulations with two state-of-the-art open-source LLMs and a real-world text dataset. Simulation results demonstrate that the proposed algorithm outperforms benchmark schemes across various network settings, achieving up to 5x inference speed acceleration and improving inference accuracy.

The rest of this paper is organized as follows. In Section II, we elaborate on the system model and present the problem formulation. In Section III, we develop a two-stage algorithm and prove its convergence. In Section IV, we extend the algorithm for multi-antenna edge devices. Simulation results are presented in Section V, and we conclude the paper in Section VI.

Notations: Column vectors and matrices are denoted by boldface lowercase and boldface capital letters, respectively. The symbol $\mathbb{R}$ denotes the set of real numbers. $\mathbb{C}^{M\times N}$ represents the space of the $M\times N$ complex-valued matrices. $(\cdot)^{\mathsf{T}}$ and $(\cdot)^{\mathsf{H}}$ stand for the transpose and the conjugate transpose of their arguments, respectively. $\textup{tr}(\mathbf{A})$ denote the trace of matrix $\mathbf{A}$. $\mathbb{E}[\cdot]$ denotes the expectation operation. $\nabla$ represents the gradient operator. $|\cdot|$ and $\|\cdot\|$ stand for the $\ell_{1}$ and $\ell_{2}$ norm of vectors.

To deploy LLMs on resource-limited edge devices, distributed on-device inference with tensor parallelism has been proposed. This method involves partitioning large neural network tensors (e.g., weight matrices) of LLMs into smaller segments and distributing them across multiple edge devices for simultaneous processing. The complete workflow of the proposed distributed on-device LLM inference system is illustrated in Fig. 1. When a device initiates an inference request, the edge server dynamically identifies available local devices and partitions the model parameters. Then, each device loads its assigned model segment into memory and performs forward computation. After each layer of the LLM is computed, an all-reduce operation aggregates the intermediate layer outputs from all devices, ensuring synchronization and consistency across devices during inference. In the proposed distributed inference framework, the device shares its input (typically token embeddings rather than raw data) with other participating devices. For scenarios demanding strict confidentiality, encryption schemes (e.g., homomorphic encryption) or secure enclaves can be adopted to mitigate privacy leakage. Furthermore, we highlight two typical scenarios illustrating real-world, trusted environments particularly suitable for our distributed inference framework, as shown in the following.

• •

  Organizational or HPC Clusters: Large institutions (e.g., corporate data centers, national labs, or university HPC centers) often host massive LLMs that exceed the capacity of a single node. In these clusters, multiple servers within the same security domain can distribute model segments or layers among them, securely exchanging raw input data via internal networks. Since all compute nodes reside in the same trusted infrastructure (with well-defined access control, encryption, and compliance policies), they can fully leverage parallelization to reduce per-inference latency and alleviate memory bottlenecks, without risking data exposure to external environments.

• •

  Single-User or Local Edge Scenarios: Individual users or small teams may possess multiple personal devices or localized edge servers (e.g., the home server or on-premises GPU node). These devices operate within a single-user network or closed local environment, allowing them to share raw inputs without breaching privacy. By splitting the LLM’s parameters or layers across these trusted devices, users can achieve faster response times and reduced memory load per device. These benefits are especially valuable for real-time applications (e.g., smart home assistants or AR/VR), where offloading data to external clouds may be undesirable or impractical.

LLMs are primarily built on the Transformer architecture, which typically consists of dozens of Transformer layers [45]. Each Transformer layer includes a self-attention mechanism and a multi-layer perceptron (MLP). To achieve efficient distributed inference, tensor parallelism partitions both the self-attention and MLP layers within each Transformer block into smaller tensor segments, as shown in Fig. 1. We note that both pipeline parallelism and tensor parallelism are two prevalent model partitioning strategies widely adopted in distributed inference frameworks. While pipeline parallelism partitions the model across layers, tensor parallelism partitions computations within each layer across multiple devices. Tensor parallelism is particularly attractive for on-device inference due to its inherent advantages in significantly reducing idle times (pipeline bubbles), achieving finer-grained memory allocation, and, when combined with AirComp-based aggregation, greatly minimizing communication overhead. These properties highlight the practical benefits and superior suitability of tensor parallelism for resource-constrained and latency-sensitive inference scenarios considered in this work.

Employing tensor parallelism for distributed LLM inference requires frequent all-reduce operations, which cause significant communication overhead in practical wireless networks. To address this issue, we propose a communication-efficient AirComp all-reduce approach. The AirComp aggregates distributed data efficiently by leveraging the signal superposition property of wireless multiple-access channels, allowing simultaneous transmissions to compute nomographic functions (e.g., arithmetic mean) [46]. In the proposed distributed LLM inference system, the aggregation of intermediate layer outputs in the all-reduce step aligns with this operation, making the AirComp suitable to mitigate communication overhead. Note that edge devices performing AirComp must achieve symbol-level synchronization to ensure their transmitted signals arrive concurrently at the receiver, minimizing aggregation errors due to timing offsets. In our framework, synchronization among edge devices can be practically realized through the well-established timing advance (TA) mechanism. Specifically, the edge server estimates each device’s timing offset and instructs each device to adjust its signal transmission timing via dedicated TA commands. By aligning transmissions precisely, edge devices can ensure simultaneous arrival and accurate signal aggregation at the receiver.

We consider a wireless network consisting of an edge server with $N_{r}$ antennas and $N$ single-antenna edge devices. We further extend the proposed framework to a more general scenario invloving multi-antenna edge devices in Section IV. The uplink channels from edge devices to the server are block-fading, where channel statistics remain constant throughout the inference process, with channel states varying independently across different time intervals. Let $z_{n}$ denote the per-round transmitted entry of device $n$’s intermediate layer output $\mathbf{Z}_{n}$, which has a complete dimensionality of $L_{0}$. To reduce transmission power, the transmitted symbols are normalized to have zero mean and unit variance, i.e., $\mathbb{E}[\|z_{n}\|^{2}]=1$, where the normalization factor is uniform for all devices and can be inverted at the server. Given synchronized symbol boundaries, all devices transmit their intermediate layer outputs simultaneously. To mitigate the distortion of received signals caused by channel noise, aggregation beamforming is adopted. Let $\mathbf{a}\in\mathbb{C}^{N_{r}\times 1}$ denote the aggregation beamforming vector at the edge server. After the AirComp, the received signal at the server is given by,

where $\mathbf{h}_{n}\in\mathbb{C}^{N_{r}\times 1}$ denotes the uplink channel from device $n$ to the server, $b_{n}$ is the transmit power of device $n$, and $\mathbf{n}\sim\mathcal{CN}\left(0,\sigma^{2}\mathbf{I}\right)$ denotes the additive white Gaussian noise vector with $\sigma^{2}$ being the noise power. In the single-antenna setting, each device employs only a scalar transmit-power coefficient $b_{n}$ (instead of a beamforming vector) to scale its transmitted scalar entry $z_{n}$. The distortion of $\hat{z}$ with respect to the desired target summation $z=\sum_{n=1}^{N}z_{n}$ is measured by the MSE, which is defined as

The MSE serves as a metric to evaluate the performance of the AirComp all-reduce operations. As shown in the simulations later, the inference accuracy of the distributed on-device LLM inference system is greatly influenced by the transmission error during the AirComp phase. By substituting (10) into (11), the MSE can be explicitly represented as a function of aggregation beamforming vector $\mathbf{a}$ and transmitter scalars $\left\{b_{n}\right\}_{n=1}^{N}$ as follows,

Edge devices involved in inference tasks typically have limited energy supply. Thus, we assume that for each device $n$, the energy consumption for both the forward computation of each LLM layer and the transmission of the intermediate output cannot exceed the maximum power budget $P_{n}^{\textup{max}}$. To model the computation energy consumption, we first introduce a model assignment vector $\mathbf{m}=[m_{1},\ldots,m_{N}]$ with its entry $m_{n}\in[0,1]$ representing the proportion of model allocated to device $n$. Consequently, the computation energy consumption for device $n$ is given by $e_{n}m_{n}s^{\textup{tot}}$, where $e_{n}$ denotes the device-specific energy coefficient that reflects the energy cost associated with accessing and processing each weight during computation, and $s^{\textup{tot}}$ is the number of parameters (weights) for each layer. The communication energy consumption of device $n$ can be derived as $L_{0}\|b_{n}\|^{2}$. Accordingly, the power constraint is given by

In the proposed distributed LLM inference system, the overall performance is determined by the model assignment policy $\mathbf{m}$ and the transceiver design $\mathbf{a}$, $\left\{b_{n}\right\}$. Optimal model assignment ensures that each device processes a suitable portion of the model based on its computational capability (e.g., memory size and compute power). Meanwhile, efficient transceiver optimization can reduce signal misalignment error and suppress channel noise, thereby improving inference accuracy. Thus, to improve inference performance, we formulate a joint model assignment and transceiver optimization problem that aims to minimize the average MSE, subject to the per-device power constraints. Importantly, the transceiver design can adapt dynamically to instantaneous CSI. In contrast, adapting the model assignment policy to instantaneous CSI in a real-time manner is impractical due to the significant latency caused by loading different model segments. Thus, model assignment should be finished before inference based on the long-term channel statistics.

The resulting problem is therefore formulated as a mixed-timescale joint optimization of the short-term transceiver variables $\mathbf{a}$, $\left\{b_{n}\right\}$ and the long-term model assignment policy $\mathbf{m}$ as follows,

where the expectation $\mathbb{E}_{\mathbf{h}}\left[\cdot\right]$ is taken over all random channel realizations $\mathbf{h}=\left\{\mathbf{h}_{n}\right\}_{n=1}^{N}$. However, the problem $\mathcal{P}_{1}$ is challenging to be solved due to the following three reasons.

• •

  Non-convexity: The objective function is inherently non-convex due to the coupling between the receiver aggregation beamformer $\mathbf{a}$ and the transmitter scalars $\left\{b_{n}\right\}$.

• •

  Expectation over Random Channels: The objective involves an expectation over random CSI, which requires prior knowledge of channel statistics.

• •

  Interdependence of Timescales: The per-device power constraints link the short-term transceiver variables with the long-term model assignment policy, leading to a complex interplay between the two timescales.

To address these challenges, we develop a two-stage algorithm that separately solves the short-term transceiver optimization and the long-term model assignment optimization in the following section.

We start by decomposing problem $\mathcal{P}_{1}$ into a family of short-term transceiver optimization problems and a long-term model assignment optimization problem as follows.

The short-term problem $\mathcal{P}_{s}$ is challenging to be solved due to the inherent non-convexity caused by the coupling between receiver aggregation beamformer $\mathbf{a}$ and the transmitter scalers $\left\{b_{n}\right\}_{n=1}^{N}$. Thus, we first simplify problem $\mathcal{P}_{s}$ by demonstrating that the channel inversion precoding is optimal conditioned on the aggregation beamformer.

Let $\mathbf{g}$ represent the normalized aggregation beamformer that satisfies $\mathbf{g}^{\mathsf{H}}\mathbf{g}=1$, and consequently $\mathbf{a}=\sqrt{\alpha}\mathbf{g}$ where $\alpha$ is optimized to satisfy the power constraints of edge devices. By applying Lemma 1, problem $\mathcal{P}_{s}$ can be reformulated as follows,

Then, by employing the equation $\|\mathbf{g}^{\mathsf{H}}\mathbf{h}_{n}\|^{2}=\textup{tr}\left(\mathbf{h}_{n}\mathbf{h}_{n}^{\mathsf{H}}\mathbf{g}\mathbf{g}^{\mathsf{H}}\right)$, an equivalent formulation of problem (17) is obtained as follows,

The problem (18) remains intractable due to the non-convex norm constraint on $\mathbf{g}$. To address this issue, we apply the SDR approach that relaxes the non-convex norm constraint by employing its convex hull.

By applying Lemma 2, we can replace the non-convex norm constraint by its convex hull and reformulate a relaxed version of problem (18) as follows,

where ${\hat{\mathbf{G}}}=\mathbf{g}\mathbf{g}^{\mathsf{H}}$. The problem (19) can be proved to be convex, and the globally optimal solution $\hat{\mathbf{G}}^{*}$ can be obtained by using a convex solver (e.g., the CVX toolbox in MATLAB [49]).

We note that the optimal solution $\hat{\mathbf{G}}^{*}$ has a high probability to satisfy the rank-one constraint [47]. If a rank-one solution $\hat{\mathbf{G}}^{*}$ is obtained, the optimal solution $\mathbf{g}^{*}$ of the original problem (18) can be immediately achieved by extracting the dominant eigenvector of $\hat{\mathbf{G}}^{*}$ as $\mathbf{g}^{*}=[\mathbf{V}_{\hat{\mathbf{G}}^{*}}]_{:,1}$. Otherwise, if the rank of $\hat{\mathbf{G}}^{*}$ is larger than 1, we apply the Gaussian randomization algorithm [50] to map the solution to a feasible, near-optimal solution for the original non-convex problem.

In this subsection, we propose a stochastic SCA algorithm to solve the long-term model assignment problem $\mathcal{P}_{l}$. The proposed algorithm requires no prior knowledge of channel statistics. For clearer algorithmic description, we first reformulate the long-term problem $\mathcal{P}_{l}$ into an equivalent form as follows,

where $\mathbf{e}_{c}\!\left(\mathbf{m}\right)\!=\![\|b_{1}^{*}(\mathbf{m})\|^{2},\ldots,\|b_{N}^{*}(\mathbf{m})\|^{2}]^{\mathsf{T}}\!,$ $\mathbf{p}^{\textup{max}}=[{P}_{1}^{\textup{max}},\ldots,{P}_{N}^{\textup{max}}]^{\mathsf{T}}$, and $\mathbf{e}=[e_{1},\ldots,e_{N}]^{\mathsf{T}}$. The proposed stochastic SCA algorithm iteratively performs the following two steps: First, quadratic surrogate functions $\hat{f}_{0}(\mathbf{m})$, $\hat{f}_{1}(\mathbf{m})$ are constructed to approximate the non-convex components of the original objective and constraint functions $f_{0}(\mathbf{m})$, $f_{1}(\mathbf{m})$, respectively. Then, the resulting convex quadratic approximation problem is solved, and the long-term model assignment policy is updated based on the solution. The details of these two steps are illustrated as follows.

In this subsection, we analyze the asymptotic convergence performance of Algorithm 1 to a stationary point of the original problem $\mathcal{P}_{1}$.

We first show the convergence of the surrogate functions in the following lemma.

To elaborate the convergence result, we need to introduce the Slater’s condition for the converged surrogate function in the following.

The Slater’s condition is widely used in constrainted optimization algorithms (e.g., the majorization-minimization algorithm [53] and virtual queue-based online convex optimization algorithm [54]). With Lemma 3 and the Slater’s condition, we are ready to show the main convergence result of Algorithm 1 in the following theorem.

In Section V-B, we further verify the convergence of Algorithm 1 through simulations.

Building upon the single-antenna setting, we now consider a more generalized scenario where edge devices in the distributed LLM inference system are equipped with multiple antennas. Thus, the spatial diversity and spatial multiplexing are utilized to further improve communication efficiency. Specifically, we consider the server and each edge device are equipped with $N_{r}$ and $N_{t}$ antennas, respectively. Similar to the single-antenna case, all devices simultaneously upload their intermediate layer outputs through wireless multiple-access channels. Let $\mathbf{z}_{n}=[z_{n,1},\ldots,z_{n,L}]^{\mathsf{T}}$ denote the per-round transmitted $L$ entries of device $n$’s intermediate output $\mathbf{Z}_{n}$. Let $\mathbf{A}\in\mathbb{C}^{N_{r}\times L}$ and $\mathbf{B}_{n}\in\mathbb{C}^{N_{t}\times L}$ denote the aggregation beamforming matrix at the edge server and the data precoding matrix at device $n$, respectively. Then, the received signal at the server after the AirComp can be derived as follows,

where $\mathbf{H}_{n}\in\mathbb{C}^{N_{r}\times N_{t}}$ denotes the uplink MIMO channel from device $n$ to the edge server. In the multi-antenna setting, each device employs the precoding (beamforming) matrix to map its transmitted vector $\mathbf{z}_{n}$ onto multiple antennas for simultaneous transmission. The distortion of $\hat{\mathbf{z}}$ with respect to the desired target vector $\mathbf{z}=\sum_{n=1}^{N}\mathbf{z}_{n}$ is measured by the MSE, defined as

By substituting (31) into (32), the MSE can be explicitly represented as a function of transceiver beamforming matrices as follows,

To effectively utilize the heterogeneous computational capabilities of edge devices and mitigate communication distortions, we similarly investigate a joint model assignment and transceiver optimization problem. Specifically, the joint optimization problem in the multi-antenna scenario can be formulated as follows,

where the expectation $\mathbb{E}_{\mathbf{H}}\left[\cdot\right]$ is taken over all random channel realizations $\mathbf{H}=\left\{\mathbf{H}_{n}\right\}_{n=1}^{N}$.

In this subsection, we extend Algorithm 1 to a more general case involving multi-antenna edge devices. Similarly, we first decompose problem $\mathcal{P}_{2}$ into a family of short-term subproblems and a long-term subproblem as follows.