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

arXiv: https://arxiv.org/abs/2505.02533 · HTML: https://mlsanity.com/paper/2505.02533/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving.
- Builds on 6 of its 20 references that are in the map.
- First posted on arXiv 5 May 2025.

## Placed in

- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 6 of 6)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference](https://mlsanity.com/paper/2405.17245/) (2024)
- [SplitPlace: AI Augmented Splitting and Placement of Large-Scale Neural Networks in Mobile Edge Environments](https://mlsanity.com/paper/2205.10635/) (2022)
- [SplitLLM: Collaborative Inference of LLMs for Model Placement and Throughput Optimization](https://mlsanity.com/paper/2410.10759/) (2024)

## Built on by (most cited first, 5 of 5)

- [Splitwise: Collaborative Edge-Cloud Inference for LLMs via Lyapunov-Assisted DRL](https://mlsanity.com/paper/2512.23310/) (2025)
- [Federated Attention: A Distributed Paradigm for Collaborative LLM Inference over Edge Networks](https://mlsanity.com/paper/2511.02647/) (2025)
- [Rethinking Inference Placement for Deep Learning across Edge and Cloud Platforms: A Multi-Objective Optimization Perspective and Future Directions](https://mlsanity.com/paper/2510.22909/) (2025)
- [Living the Novel: A System for Generating Self-Training Timeline-Aware Conversational Agents from Novels](https://mlsanity.com/paper/2512.07474/) (2025)
- [BALANCE: Hybrid Autoregressive-Speculative LLM Inference at the Network Edge](https://mlsanity.com/paper/2608.05926/) (2026)

## Abstract (arXiv)

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.