BALANCE: Hybrid Autoregressive-Speculative LLM Inference at the Network Edge

Index Terms: Large language models, edge inference, autoregressive decoding, speculative decoding, user scheduling, task throughput maximization.

The paper addresses edge LLM inference in next-generation mobile networks: In next-generation mobile networks, large language model (LLM) inference at the network edge is emerging as an important paradigm for delivering LLM services to users. The two representative inference approaches are autoregressive decoding (AD) and speculative decoding (SD). In AD, an LLM generates the response token by token, producing one new token in each forward pass based on the input and previously generated tokens. However, the sequential decoding process of AD incurs high response latency. To address this, SD has been proposed. In SD, a small language model (SLM) first generates multiple draft tokens via AD, and the target LLM verifies them in parallel within a single forward pass. SD maintains output quality since the verification preserves the target LLM's output distribution, but SD introduces higher memory overhead due to the coexistence of the SLM and LLM as well as the memory for parallel draft verification.

The central challenge is that "Neither AD nor SD alone can efficiently serve users with heterogeneous demands on resource-constrained edge servers. Serving users with AD alone is memory-efficient but may fail to satisfy stringent latency requirements, whereas relying solely on SD reduces latency but may exceed available memory resources. This is quantified experimentally: when serving a single user, SD reduces the average end-to-end (E2E) latency by 39.1% compared with AD, but increases the GPU memory consumption by 18.3%" (using Llama-2-7B as the LLM and TinyLlama-1.1B as the SLM on an NVIDIA GeForce RTX 4090 GPU).

The authors propose a hybrid autoregressive-speculative inference (BALANCE) framework. In BALANCE, an edge server hosts both an LLM and an SLM, enabling the simultaneous support of AD and SD for different users. The server jointly selects users for service, assigns each selected user to AD or SD, and allocates computing resources between the two modes. Users scheduled to the same mode are batched for inference. The procedure involves input uploading, LLM prefilling for AD, SLM/LLM prefilling for SD, iterative decoding/draft-verification, and output downloading.

Communication model: We consider orthogonal frequency-division multiple access for both uplink and downlink transmissions. The uplink time is Δτkup = γLIk/(BkRk), and the uploading times for AD and SD modes are given by τad,up = max Δτkup over scheduled users and similarly for τsd,up. Downlink time is τkdown = (xk + yk)Δτkdown.

Computing model: Workloads are modeled using Transformer matrix multiplications (QKV projections, MHA, output projection, FFN). For AD, the LLM prefilling workload is W ad,p = ΓL,p(Lad,Imax) Σxk, and decoding-step workloads are Wlad,a = Σ xk I l≤LOk −1 ΓL,a l (Lad,Imax). For SD, both LLM and SLM prefill the batched prompts, then iterate: the SLM autoregressively generates draft tokens for unfinished requests, and the LLM verifies the batched draft sequences in one forward pass. The number of iterations for user k is Nk = ⌈LOk/LAk + 1⌉ with expected acceptance rate Ak.

Memory model: Memory for weights is qL and qS (equations 13–14). Peak KV cache memory is qad = qad,a for AD and qsd = qsd,a + qsd,v for SD.

E2E latency: tEk = xk(τad,up + τadLad,Omax−1) + yk(τsd,up + τsdNsdmax) + τkdown. Latencies under GPU resource allocation z are modeled with affine functions: τad,p = αL W ad,p/z + βL, etc.

P1: max r(X + Y) = Σ(xk + yk) subject to xk + yk ≤ 1 (each user in at most one mode), tEk ≤ Tk (latency requirements), I z<1 qS + qL + qad + qsd ≤ Q (memory constraint), binary constraints, and z ∈ Z. The objective is to maximize the number of served users in edge LLM inference while satisfying users' E2E latency requirements and the server's memory constraint. The problem is combinatorial and NP-hard: "Consider a special case of P1 with fixed z. The problem then reduces to user scheduling under latency requirements and memory resource constraints, which can be mapped to a multi-dimensional multiple-choice knapsack problem (MMKP). Since the MMKP is NP-hard, P1 is NP-hard."

The authors transform P1 into P2 by introducing an auxiliary memory-split variable λ (with Q1 = λQ − qL − I z<1 qS and Q2 = (1−λ)Q − qL − I z<1 qS), proving equivalence (Proposition 2). Fixing z and λ reduces P2 to P3 (Proposition 3). P3 is then decoupled into two sub-problems, P3.1 and P3.2, which schedule users in the AD and SD modes, respectively.

Algorithm 1 (AD mode): Enumerates feasible values of Lad,Imax, Lad,Omax, t1, and τ1. For each tuple, user k's latency is t̃adk = (αL W̃ad,pk + Σl αL W̃ad,alk)/z, and the number of served users comes from a sorted permutation where the count d*1 is the maximum d1 satisfying sum of t̃ad ≤ Δt1 and d1 ≤ ⌊Q1/q̃ad⌋. Proposition 4: Algorithm 1 obtains an optimal solution to P3.1 under given z and λ.

Algorithm 2 (SD mode): Similarly enumerates Lsd,Imax, Lsd,Amax, Ñ1, t2, and τ2, computing candidate latency t̃sdk and memory q̃sdk. Since both t̃sdk and q̃sdk are monotonically non-decreasing with Nk, sorting users by t̃sdk also preserves the non-decreasing order of q̃sdk, the count d*2 is obtained via equation (32). Proposition 5: Algorithm 2 obtains an optimal solution to P3.2 under given z, λ, and K̂ad.

Algorithm 3 (Overall): Enumerates all (z, λ) pairs, runs Algorithms 1 and 2 for each, and keeps the best solution. Theorem 1: Algorithm 3 achieves a constant approximation guarantee. Specifically, the produced solution Ẋ and Ẏ satisfy r(Ẋ + Ẏ) ≥ 1/2 r(X* + Y*), where X* and Y* denote the optimal user scheduling decisions of P1. Theorem 2: Algorithm 3 has a polynomial time complexity of O(K 6 log K), where K is the total number of users.

Setup: Hybrid testing combining device measurements (AMD Ryzen Threadripper PRO 5975WX CPU, NVIDIA GeForce RTX 4090 GPU with 24 GB memory, 256 GB DDR4 RAM) and numerical simulations. Models: TinyLlama-1.1B (SLM) and Llama-2-7B (LLM) with FP16 precision. K ranges 20–70; Tk ∈ [0.5 s, 13 s]; input/output lengths up to 1000 tokens; LD = 1; Ak ∈ [0.6, 0.8]. Baselines: AD-only (Algorithm 1), SD-only (Algorithm 2), and exhaustive search.

Performance gains: BALANCE consistently outperforms both AD-only and SD-only baselines. In Fig. 4, BALANCE improves throughput over AD-only and SD-only baselines by 37.9% and 33.4%, respectively, under varying K, and by 31.5% and 27.3%, respectively, under varying Tmax. In Fig. 5, BALANCE achieves average gains over AD-only and SD-only baselines of 38.8% and 23.9%, respectively, under varying LImax, and 30.4% and 27.3%, respectively, under varying LOmax.

Running time: Compared with exhaustive search (K from 10 to 16), "the running time of the proposed algorithm grows nearly linearly with K, while exhaustive search increases nearly exponentially. The proposed algorithm achieves an average acceleration of about 5,100 times and up to 21,800 times at K = 16. Throughput-wise, the proposed algorithm achieves near-optimal task throughput, with only about 0.9% average degradation."

Ablation study: Compared against ECRA (z = 0.5), EMA (λ = 0.5), and ECRMA (both 0.5). "In Fig. 7(a), the proposed algorithm improves the task throughput by 5.6%, 18.7%, and 23.7% on average compared with ECRA, EMA, and ECRMA, respectively. The corresponding average performance improvements are 9.0%, 4.9%, and 18.4% in Fig. 7(b)."

The paper investigates a hybrid autoregressive-speculative inference (BALANCE) framework for edge LLM inference to balance the latency-memory tradeoff between conventional autoregressive decoding (AD) and speculative decoding (SD). It formulates an NP-hard joint user scheduling and resource allocation problem, transforms it into two sub-problems, and develops a polynomial-time algorithm with a constant approximation guarantee. Experiments show BALANCE "consistently outperforms the conventional AD and SD schemes and achieves significant improvements in task throughput. Overall, BALANCE provides an effective and practical solution for resource-constrained edge LLM inference systems."

Here are the specific improvements to AI systems based on this paper, along with what the improved system can do:

• Hybrid inference mode selection. The AI system can dynamically assign each user to either autoregressive decoding (AD) or speculative decoding (SD) instead of using a single fixed decoding strategy. It can simultaneously serve both modes on one edge server, adapting to per-user latency deadlines and available memory.

• Joint user scheduling and resource allocation. The system can jointly decide which users to serve, which inference mode each served user gets, and how to split GPU compute and memory between AD and SD. This maximizes the number of users served under end-to-end latency and memory constraints.

• Memory-constrained admission control. Using explicit models for weight memory, KV cache memory, and draft-verification memory, the system can reject or defer users when peak memory would be exceeded, preventing out-of-memory failures while still maximizing throughput.

• Latency-aware batching. The system can batch users assigned to the same inference mode and compute the exact number of batches, draft-token lengths, and verification iterations needed to meet each user’s deadline. It can enumerate feasible batch configurations and select the set of users with the highest total count.

• Polynomial-time near-optimal decisions. The system can make scheduling and resource-split decisions in O(K 6 log K) time with a constant-factor approximation guarantee (at least half of the optimal served-user count), making it practical for large user populations. Compared with exhaustive search, it runs thousands of times faster while suffering only 0.9% throughput degradation.

• Adaptive compute/memory split. The system can optimize the memory split parameter λ and the compute fraction z between AD and SD in real time, rather than using fixed 50/50 splits. This yields 5–24% throughput gains over fixed-split baselines.

• Speculative decoding parameter tuning. The system can tune the number of draft tokens and expected acceptance rates per user based on output length requirements, balancing the latency reduction of SD against its extra memory overhead.

• Higher overall task throughput. With all the above, the improved AI system can serve 30–40% more users than AD-only or SD-only approaches under representative edge-server memory and heterogeneous latency requirements, while maintaining the target LLM’s output distribution via speculative decoding verification.