Large Language Model Partitioning for Low-Latency Inference at the Edge

TL;DR

This work addresses low-latency autoregressive inference for large language models at the edge by partitioning a single-layer decoder across multiple devices at the level of attention heads, co-locating each head with its K/V cache. A myopic, interval-based algorithm dynamically assigns and migrates attention-head blocks to balance memory, compute, and communication constraints, reducing latency and mitigating cache growth. Across small-scale tests, the method achieves within 15–20% of the exact optimal latency; in larger, more realistic settings it delivers substantial speedups (up to 9–10x) over static layer-based partitioning baselines like EdgeShard and Galaxy. The approach enables parallelized head-level execution and scalable edge inference for autoregressive decoding, with future directions including multi-layer extensions, foresight-based decision-making, and energy-aware deployments.

Abstract

Large Language Models (LLMs) based on autoregressive, decoder-only Transformers generate text one token at a time, where a token represents a discrete unit of text. As each newly produced token is appended to the partial output sequence, the length grows and so does the memory and compute load, due to the expanding key-value caches, which store intermediate representations of all previously generated tokens in the multi-head attention (MHA) layer. As this iterative process steadily increases memory and compute demands, layer-based partitioning in resource-constrained edge environments often results in memory overload or high inference latency. To address this and reduce inference latency, we propose a resource-aware Transformer architecture partitioning algorithm, where the partitioning decision is updated at regular intervals during token generation. The approach is myopic in that it is based on instantaneous information about device resource availability and network link bandwidths. When first executed, the algorithm places blocks on devices, and in later executions, it migrates these blocks among devices so that the sum of migration delay and inference delay remains low. Our approach partitions the decoder at the attention head level, co-locating each attention head with its key-value cache and allowing dynamic migrations whenever resources become tight. By allocating different attention heads to different devices, we exploit parallel execution of attention heads and thus achieve substantial reductions in inference delays. Our experiments show that in small-scale settings (3-5 devices), the proposed method achieves within 15 to 20 percent of an exact optimal solver's latency, while in larger-scale tests it achieves notable improvements in inference speed and memory usage compared to state-of-the-art layer-based partitioning approaches.

Figures (4)

• Figure 1: Illustration of our approach, showing how attention head-level partitions (red lines) and additional cuts for feed-forward network (FFN) and linear projection (proj) blocks (green lines) enable flexible assignment of blocks across multiple edge devices. The controller node makes inference requests and handles partitioning decisions, while the attention heads (with their K/V caches), FFN, and proj blocks are allocated among devices D1, D2, D3, and D4.
• Figure 2: A single-layer decoder-only Transformer at intervals $\tau$ and $\tau + 1$, continuing autoregressively up to the $N$-th token generation. Red lines show attention head-level partitioning in MHA; green lines indicate projection and feed-forward partitioning. Cyan vs. green blocks highlight new tokens building on the K/V caches of the attention heads. At the bottom, a controller allocates these blocks across edge devices.
• Figure 3: Inference latency vs. generated token step $n$ on 25 devices. Our approach (Resource-Aware) avoids steep growth.
• Figure 4: Total memory usage vs. generated token step $n$ in the 25-device setup. Our resource-aware approach mitigates memory growth more effectively.