Efficient Heterogeneous Large Language Model Decoding with Model-Attention Disaggregation

TL;DR

This work identifies decoding as a memory-bound bottleneck in transformer-based LLMs within disaggregated serving. It introduces model-attention disaggregation, separating KV-cache/attention to memory-optimized devices from other operators on computation-optimized devices, and implements Lamina with a fully host-bypassed network stack and an automated model converter. Through real-world traces and a heterogeneous GPU cluster, Lamina achieves substantial throughput gains (16.1% to 90.1%) at similar hardware costs, while maintaining latency within service-level objectives thanks to staggered pipelining and effective overlap. The approach generalizes to other operator-level disaggregations and presents a practical pathway to more efficient large-scale LLM serving in production environments.

Abstract

Transformer-based large language models (LLMs) exhibit impressive performance in generative tasks but also introduce significant challenges in real-world serving due to inefficient use of the expensive, computation-optimized accelerators. Although disaggregated serving architectures have been proposed to split different phases of LLM inference, the efficiency of decoding phase is still low. This is caused by the varying resource demands of different operators in the transformer-based LLMs. Specifically, the attention operator is memory-intensive, exhibiting a memory access pattern that clashes with the strengths of modern accelerators, especially for long context requests. To enhance the efficiency of LLM decoding, we introduce model-attention disaggregation. This approach leverages a collection of cheap, memory-optimized devices for the attention operator while still utilizing high-end accelerators for other parts of the model. This heterogeneous setup ensures that each component is tailored to its specific workload, maximizing overall performance and cost efficiency. Our comprehensive analysis and experiments confirm the viability of splitting the attention computation over multiple devices. Also, the communication bandwidth required between heterogeneous devices proves to be manageable with prevalent networking technologies. To further validate our theory, we develop and deploy Lamina, an LLM inference system that incorporates model-attention disaggregation in a distributed heterogeneous cluster. Experimental results indicate that Lamina can provide 16.1 ~ 90.1% higher estimated throughput than existing solutions with similar costs.

Figures (14)

• Figure 1: The disaggregated architecture of LLM serving.
• Figure 2: Measured time consumption and MFU of non-attention operators in LLaMA3-70B during one decode iteration. Results with different tensor parallelisms are presented. The dotted lines indicate the projected values using the roofline model.
• Figure 3: Measured time consumption and model bandwidth utilization (MBU) of attention operators in LLaMA3-70B during one decode iteration. Results with different sequence lengths and hardware configurations are presented.
• Figure 4: The required network bandwidth for decoding LLaMA3-70B using attention offloading with H100 and H20, with at most $20\%$ latency slow-down for network overhead.
• Figure 5: Diagram of WR submission with FHBN send and conventional GPU-aware send implementations.