LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention

LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention

Speakers: Shang Yang (MIT); Junxian Guo (MIT); Haotian Tang (MIT); Guangxuan Xiao (MIT); Song Han (Professor, MIT)

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=mx4lHNfwbM0

Overview

In the rapidly evolving landscape of artificial intelligence, long context Large Language Models (LLMs) have emerged as a pivotal technology, unlocking new frontiers in applications ranging from comprehensive document and video understanding to complex multi-turn reasoning tasks. However, the efficient deployment and serving of these models present significant challenges, primarily due to the quadratic computational complexity of the attention mechanism with respect to sequence length. The talk "LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention" introduces a novel system designed to address these bottlenecks, offering a unified framework for optimizing both the prefilling stage (processing the initial long prompt) and the decoding stage (generating subsequent tokens).

Presented by Shang Yang, a second-year PhD student at MIT, alongside his colleagues and advisor Professor Song Han, LServe proposes a comprehensive solution centered on the intelligent application of sparsity. The core innovation lies in integrating various sparse attention patterns—both static and dynamic—into a single, efficient system, while simultaneously tackling system-level challenges like memory management and overhead reduction. By doing so, LServe aims to not only accelerate inference but also reduce the memory footprint of the Key-Value (KB) cache, making long-context LLM serving more accessible and cost-effective for real-world deployments.

The significance of LServe cannot be overstated in an era where the demand for LLMs capable of processing ever-longer contexts is continually growing. The system's ability to achieve substantial speedups—up to 2.1 times faster than existing state-of-the-art frameworks like vLLM—while meticulously preserving model accuracy, marks a crucial step forward. This work is critical for enabling the next generation of AI applications that rely on deep contextual understanding, ensuring that the computational demands of these powerful models do not become an insurmountable barrier to innovation and widespread adoption.

Background

▶ Watch: LServe: Efficient long-sequence LLM serving introduction (0:00)

The advent of long context LLMs has revolutionized the capabilities of AI, allowing models to process and generate information spanning thousands or even hundreds of thousands of tokens. This expanded context window is essential for applications like summarizing entire books, analyzing lengthy legal documents, transcribing and understanding long videos, and performing intricate, multi-step reasoning. However, this power comes at a steep computational cost. The primary bottleneck in serving these models is the attention computation, which exhibits a quadratic complexity (O(N^2)) with respect to the input sequence length (N). As N grows, the attention operation quickly dominates the overall inference time, surpassing the cost of linear layers and other components. This issue is particularly acute in both the prefilling stage, where a long input prompt is processed, and the decoding stage, where subsequent tokens are generated, often requiring attention over the entire accumulated context.

Traditional approaches to serving LLMs, while optimized for shorter sequences, struggle to scale efficiently to very long contexts. The KB cache, which stores the key and value states for all previously processed tokens, also grows linearly with sequence length, consuming significant GPU memory. This memory pressure further exacerbates the serving challenge, limiting batch sizes and the maximum context length that can be practically handled.

To mitigate the quadratic complexity of attention, researchers have explored sparse attention mechanisms. The fundamental idea behind sparse attention is to restrict the attention operation to only a subset of relevant tokens, rather than attending to every token in the sequence. While fine-grained sparsity—where individual tokens are selectively attended to—can offer high theoretical efficiency, its irregular memory access patterns and computational overhead often make it impractical on modern hardware like GPUs. Block sparse attention, which operates on blocks of tokens, offers a more hardware-friendly alternative by maintaining better memory locality and reducing overhead.

Prior work has explored different facets of sparse attention. For instance, methods like query-aware selective attention, exemplified by the Quest paper, attempt to dynamically identify and attend to the most important KB pages based on the current query token. Another approach, seen in the Due Attention work from the same research group, introduces static sparsity during the prefilling stage by categorizing attention heads into those that require full context and those that can operate with a limited, "streaming" view. Despite these advancements, a unified system that seamlessly integrates various sparse attention patterns across both prefilling and decoding stages, while addressing system-level challenges like the overhead of sparsity and the practicalities of KB cache management, has remained an open problem. Specifically, the "page size dilemma"—where smaller page sizes offer better accuracy but larger page sizes provide better hardware efficiency—and the overhead of dynamic page selection during decoding, have presented significant hurdles to widespread adoption of sparse attention in long-context LLM serving.

Key Findings

▶ Watch: Introducing block sparse attention for acceleration (1:58)

LServe introduces a groundbreaking approach to efficient long-sequence LLM serving, delivering a suite of innovations that collectively address the core challenges of attention computation and KB cache management. The key findings and contributions of this work are:

• Unified Hybrid Sparse Attention Framework: LServe successfully unifies diverse sparse attention patterns—encompassing both static sparsity (primarily for the prefilling stage) and dynamic sparsity (for the decoding stage)—within a cohesive block sparse attention framework. This integration allows for comprehensive optimization across the entire LLM inference pipeline.
• Overhead-Free Static Sparsity with Fused Kernel: For the prefilling stage, LServe leverages an intelligent distinction between streaming heads (attending to a fixed, local context) and dense attention heads. Crucially, it introduces a fused attention kernel that processes both types of heads concurrently within a single kernel launch. This design completely eliminates the overhead typically associated with managing and switching between different attention patterns, ensuring that sparsity is introduced without performance penalties.
• Hierarchical Paging System to Resolve the Page Size Dilemma: LServe proposes a novel hierarchical paging system that decouples the logical granularity of attention pruning from the physical memory layout of the KB cache. By introducing smaller logical pages within larger physical pages, LServe enables fine-grained accuracy control (via logical page importance scoring) while simultaneously benefiting from the hardware efficiency of larger, contiguous physical memory blocks. This innovation ensures that important tokens are not inadvertently ignored due to increased physical page size, thereby preserving model accuracy.
• Reusable Page Selector for Reduced Dynamic Sparsity Overhead: During the decoding stage, dynamic sparsity relies on a page selector to identify relevant KB pages. Recognizing that this selection process can become a bottleneck when attention computation itself is highly optimized, LServe introduces a reusable page selector. This mechanism exploits the temporal locality inherent in natural language, where consecutive query tokens often attend to similar KB pages. By reusing selection results for a chunk of query tokens, LServe significantly amortizes the computational cost of page selection.
• Significant Performance Gains and Accuracy Preservation: Through rigorous evaluation, LServe demonstrates substantial efficiency improvements. It achieves up to 2.1 times speedup over vLLM, a leading state-of-the-art LLM serving system, across various models, sequence lengths, and hardware configurations. Critically, these performance gains are realized without compromising model accuracy, validated on challenging long-context benchmarks such as Needle in Haystack and Ruler benchmark, as well as reasoning tasks like AIME 2024 and MMLU 500.
• KB Cache Memory Savings: As a direct byproduct of its sparse attention optimizations, particularly the fixed-length streaming head pages and the selective attention in the decoding stage, LServe inherently reduces the memory consumption of the KB cache. This memory efficiency further enhances the system's ability to serve longer sequences and larger batch sizes.

Technical Deep Dive

▶ Watch: LServe's unified hybrid sparse attention patterns (2:19)

LServe's architecture is a sophisticated integration of static and dynamic sparse attention techniques, underpinned by hardware-aware optimizations and a novel memory management system. The core principle is to introduce sparsity at the block level to maximize hardware efficiency while meticulously preserving the model's contextual understanding.

Unified Sparse Attention Framework

The foundation of LServe is its unified sparse attention framework, which orchestrates different sparsity patterns. The system employs block sparse attention, a technique that skips entire blocks of less important KB cache pages. This approach is favored over fine-grained sparsity due to its superior hardware efficiency, as it maintains better memory access continuity and incurs less overhead. LServe categorizes attention into two primary types:

1. Static Sparsity for Prefilling: Applied during the initial prompt processing.
2. Dynamic Sparsity for Decoding: Applied during token generation.

This hybrid approach ensures that optimizations are tailored to the distinct computational characteristics of each stage.

Prefilling Stage: Static Sparsity with Fused Attention Kernel

In the prefilling stage, LServe draws inspiration from Due Attention, a previous work from the same research group. The key idea here is to differentiate between various attention heads within the LLM. Some heads, termed streaming heads, are identified as only needing to attend to a limited context—specifically, the very first few "sync tokens" and a small window of "local tokens." Other heads require full attention over the entire input sequence and are treated as dense attention heads.

To process these heterogeneous attention patterns without introducing overhead, LServe introduces a fused attention kernel. This specialized kernel is designed to handle both streaming and dense attention heads simultaneously within a single GPU kernel launch. By avoiding separate kernel invocations for different head types, LServe eliminates the typical overhead associated with managing and switching between diverse attention mechanisms. This ensures that the benefits of static sparsity are realized without any performance penalty.

Furthermore, the KB cache design for streaming heads is optimized. Since streaming heads only require a fixed-length context (sync tokens + local tokens), their corresponding KB cache pages can be allocated with a fixed, smaller size. This contrasts with dense heads, which require dynamic allocation for the full context, leading to direct memory savings in the KB cache.

Decoding Stage: Dynamic Sparsity with Hierarchical Paging and Reusable Page Selector

The decoding stage is optimized through dynamic sparsity, where the system selectively attends to the most relevant KB pages for each newly generated query token. This process relies on a page selector that prunes out less important KB pages.

Addressing the Page Size Dilemma with Hierarchical Paging

A significant challenge in dynamic sparsity is the page size dilemma. Smaller logical page sizes for pruning generally lead to better accuracy because they allow for finer-grained selection of important tokens. However, from a hardware efficiency perspective, larger physical page sizes are preferable. Larger pages mean better memory continuity in the KB cache, which translates to more efficient data loading and processing on GPUs. Standard dynamic sparsity methods often have to choose between accuracy (smaller pages) and efficiency (larger pages).

LServe resolves this by introducing a hierarchical paging system. This system decouples the logical pruning granularity from the physical memory layout using a two-level abstraction:

1. Logical Pages: Smaller, conceptual units within a physical page, used for importance estimation.
2. Physical Pages: Larger, contiguous blocks of memory in the KB cache, optimized for hardware efficiency.

For each logical page, LServe computes channel-wise statistics, specifically Kmax and Kmin (maximum and minimum values across channels). These statistics serve as compact representatives of the logical page's features. When a new query token arrives, its importance score with respect to each logical page is calculated by multiplying the query with the logical page's Kmax and Kmin values element-wise, then summing them up under the guidance of channel-wise maximums.

The importance score of a physical page is then defined as the maximum of the importance scores of all its constituent logical pages. This design ensures that if even a single token within a large physical page is highly important, its corresponding logical page will have a high score, causing the entire physical page to be selected by the pruning algorithm. Consequently, no important tokens are "amortized" or overlooked due due to the larger physical page size, thereby preserving model accuracy while still leveraging hardware-efficient large physical pages.

Reducing Overhead with Reusable Page Selector

As LServe significantly reduces the computational complexity of attention during decoding (potentially to a constant number of operations by attending to a fixed number of sparse pages), the overhead of the page selector itself can become the new bottleneck. Existing literature often overlooks this.

LServe tackles this with a reusable page selector. This optimization leverages the intrinsic temporal locality observed in natural language processing: consecutive query tokens often exhibit similar attention patterns and tend to attend to the same or very similar KB pages. Instead of performing the page selection process independently for every single query token, LServe reuses the page selection results from previous query tokens for a small "chunk" of subsequent query tokens. This strategy effectively amortizes the computational cost of page selection across multiple tokens, significantly reducing its overall overhead and ensuring that the page selector does not become a system bottleneck.

Experimental Setup & Results

▶ Watch: The 'page size dilemma' in dynamic sparsity (5:57)

LServe underwent comprehensive evaluation against state-of-the-art baselines, demonstrating both its efficiency gains and its ability to preserve model accuracy across diverse long-context and reasoning tasks. The experiments covered various models, sequence lengths, and hardware configurations, though specific GPU types were not explicitly detailed beyond "different hardwares" (implying standard GPU setups).

Accuracy Benchmarks

To validate the accuracy preservation of its sparse attention mechanisms, LServe was evaluated on several challenging benchmarks:

• Needle in Haystack: This benchmark tests a model's ability to retrieve a specific piece of information (the "needle") embedded within a very long document (the "haystack"). LServe demonstrated that its methods preserve model accuracy across varying context lengths, maintaining performance even as the sequence length increased. The effectiveness of the hierarchical paging system was specifically highlighted here; experiments showed that LServe maintains accuracy even with increased physical page sizes, with models failing at the same token budget (e.g., 3K tokens) regardless of the physical page size, indicating that important tokens were not missed.
• Ruler Benchmark (Nvidia): A highly challenging benchmark designed to assess long-context model capabilities, further confirming LServe's ability to maintain high accuracy under demanding conditions.
• Reasoning Models (AIME 2024, MMLU 500): LServe was also evaluated on benchmarks requiring complex reasoning, such as the AIME 2024 (American Invitational Mathematics Examination) and MMLU (Massive Multitask Language Understanding) 500 datasets. The results indicated that the accuracy of models using LServe remained very good, showcasing that the sparsity introduced does not detrimentally affect the model's reasoning capabilities.

Overall, the accuracy evaluations consistently showed that LServe's pruning methods are effective in preserving model fidelity across a wide range of tasks and context lengths.

Efficiency Benchmarks

LServe's efficiency was benchmarked against prominent existing LLM serving systems, including:

• M inference (likely referring to optimized inference frameworks)
• Due Attention (a prior work from the same group focusing on static sparsity)
• QKV / vLLM (the speaker mentions "VM" and "vLLM" in the context of comparison, indicating vLLM as a primary baseline).

The key metric for efficiency was speedup. The results were compelling:

• Overall Speedup: LServe consistently achieved much better performance than existing state-of-the-art frameworks across different models, sequence lengths, and hardware.
• Headline Number: In its best-case scenarios, LServe delivered up to 2.1 times speedup over vLLM. This significant improvement highlights the effectiveness of LServe's unified sparse attention, fused kernels, hierarchical paging, and reusable page selector in reducing computational load and overhead.
• Overhead Reduction: Specific figures demonstrated that the reusable page selector greatly reduced the overhead of the page selection process, which would otherwise become a system bottleneck in a vanilla dynamic sparsity design. This confirms the efficacy of LServe's hardware-aware optimizations.

The experiments collectively underscore LServe's capability to deliver substantial performance enhancements for long-sequence LLM serving without compromising the critical aspect of model accuracy.

Practical Implications

▶ Watch: Solving page size dilemma with hierarchical paging (6:24)

LServe's innovations carry significant practical implications for various stakeholders involved in the development and deployment of LLMs, particularly those dealing with extensive context windows.

For practitioners and model deployers, LServe offers a tangible solution to the computational and memory bottlenecks that currently limit the real-world application of long-context LLMs. The ability to achieve up to 2.1x speedup over vLLM means that applications requiring deep contextual understanding—such as advanced chatbots capable of understanding entire conversations, legal review tools processing lengthy contracts, or scientific discovery platforms analyzing vast research papers—can now be deployed more efficiently and cost-effectively. This enables higher throughput for a given hardware budget or the ability to serve even longer sequences than previously feasible, unlocking new use cases for AI.

Infrastructure teams responsible for maintaining and scaling LLM serving systems will find LServe particularly valuable. The system's inherent KB cache memory savings, a byproduct of its sparse attention design, allows for more efficient utilization of GPU memory. This translates to increased batch sizes, which can significantly improve overall system throughput, or the capacity to support longer maximum context lengths without requiring additional expensive hardware. The reduction in computational overhead, especially from the fused attention kernel and reusable page selector, simplifies resource planning and ensures more predictable performance characteristics under heavy load.

For model builders and researchers, LServe provides a robust framework for incorporating sparse attention without the typical concerns about accuracy degradation or complex integration. The hierarchical paging system elegantly solves the "page size dilemma," demonstrating that hardware efficiency can be achieved without sacrificing the fine-grained accuracy often associated with smaller logical attention units. This offers a blueprint for designing future sparse attention mechanisms that are both theoretically sound and practically deployable. The clarity that tokens are selected for attention rather than discarded from the KB cache is also crucial; it ensures that potentially important tokens are always available for future query tokens, maintaining flexibility and robustness.

However, like any advanced system, LServe involves certain tradeoffs and limitations. While LServe rigorously demonstrates accuracy preservation, the optimal balance between sparsity level, performance gains, and model fidelity might still require careful tuning for specific downstream tasks and model architectures. The effectiveness of the reusable page selector relies on the assumption of temporal locality in natural language; while generally true, edge cases with highly unpredictable attention patterns might see less benefit. The system's performance is also tied to the efficiency of its custom kernels, implying a potential dependency on specific hardware architectures (e.g., NVIDIA GPUs, given the mention of the Ruler benchmark). While LServe addresses many overheads, the complexity of implementing and maintaining such a unified, hardware-aware sparse attention system could be higher than simpler dense attention serving frameworks, requiring specialized expertise. Nonetheless, the substantial benefits in efficiency and scalability offered by LServe make it a compelling solution for the ongoing challenge of serving long-context LLMs.

Key Takeaways

• LServe addresses the critical attention bottleneck in long-sequence LLM serving by unifying static and dynamic sparse attention patterns within a block-sparse framework.
• A novel fused attention kernel eliminates overhead for hybrid sparsity during the prefilling stage, processing both streaming and dense attention heads concurrently.
• The hierarchical paging system resolves the "page size dilemma," decoupling logical pruning granularity from physical memory layout to ensure accuracy while leveraging hardware-efficient large physical pages.
• A reusable page selector significantly reduces the overhead of dynamic page selection during decoding by exploiting temporal locality in natural language, amortizing costs across multiple query tokens.
• LServe achieves substantial efficiency gains, offering up to 2.1x speedup over vLLM, while rigorously preserving model accuracy across challenging long-context and reasoning benchmarks.
• The system also contributes to KB cache memory savings, enhancing the capacity for serving longer sequences and larger batch sizes.

About the Speaker(s)

The primary presenter for this talk was Shang Yang, a second-year PhD student at MIT. He is advised by Professor Song Han, a distinguished faculty member at MIT. The research presented in LServe is a collaborative effort, with contributions from fellow MIT researchers Junxian Guo, Haotian Tang, and Guangxuan Xiao, alongside Professor Song Han. Their collective expertise from MIT underpins the innovative approaches developed in LServe for efficient LLM serving.