Optimizing LLM Inference: Sparse Attention and Multi-Level Caching Practices

Introduction: New Challenges in Inference Optimization from Architectural Evolution

As the parameter scale of large language models continues to grow, new architectural designs constantly emerge. While these designs enhance model capabilities, they also place significant pressure on inference deployment. Recently, a hybrid architecture LLM (referred to below as the “target model”) was introduced, integrating different attention mechanisms. Its innovative design imposes more complex combined optimization requirements on underlying computation, memory bandwidth, and parallel strategies. How to rapidly adapt these models to new hardware platforms while reusing optimization experience accumulated from previous models has become a critical challenge in AI engineering. From a technical perspective, this article analyzes the core architectural features of this model and explores a series of optimization practices, including operator fusion, quantization compression, and parallel scheduling, to provide a reference for building efficient and stable inference services.

Dual Evolution of Compute Infrastructure and Model Architecture

The Internal Logic of Hybrid Attention Mechanisms

Traditional Transformer architectures rely on global attention, where computational complexity grows quadratically with sequence length. To overcome this bottleneck, the target model adopts a hybrid architecture featuring “hierarchical processing and sparse focusing,” driven by two intertwined computational paths:

• State Deduction Path (SDP): Drawing inspiration from state space models, this path performs orderly deduction through a hidden state. It efficiently processes long-range dependencies with near-linear complexity, effectively mitigating “attention collapse” and capturing global contextual semantics.
• Sparse Attention Path (SAP): This path is a profound modification of the multi-head attention mechanism. It introduces a lightweight “Lightweight Indexer (LI)” module, which first uses an efficient scoring function to filter the Top-K most relevant positions from all historical key-value (KV) pairs. Subsequently, a “Sparse Flash Attention (SFA)” unit executes fine-grained attention computation exclusively on the selected KV blocks. Through online index selection, this mechanism achieves dynamic sparse interaction between tokens.

These two paths are not entirely independent. By introducing a “Cross-layer Residual Aggregation (CRA)” mechanism, the outputs of both paths are fused at each layer to achieve synergistic utilization of semantic information at different levels. This architecture requires inference engines to simultaneously support the efficient scheduling and coordination of two distinct computational flows.

Hardware Adaptation: Reuse and Innovation

Faced with this new architecture, the most efficient strategy is not building from scratch, but rather “standing on the shoulders of giants.” In past optimization efforts, a comprehensive fused operator library (such as highly efficient state deduction kernels) and a two-stage fusion strategy (for CRA computation) were developed for the SDP. These achievements can be directly reused in the target model’s SDP path. Meanwhile, for the pre-computation of sparse attention indices, mature index selection and lightweight scoring fusion operators already exist, requiring only parameter adjustments for the new model’s head dimensions and normalization methods.

The true innovation lies in decoupling the index computation from the sparse reads in the SAP path. Specifically, we need to build an independent “Indexer Agent” within the inference engine, allowing the index computation to be decoupled from the subsequent main attention computation. This means the LI module can act as an independent execution flow (Prolog), preparing selected KV indices in advance and passing them directly to the SFA unit via on-chip memory. This avoids writing intermediate results back to global memory, thereby eliminating unnecessary memory round-trips and kernel launch overhead.

Collaborative Optimization of Decoding Performance and Memory Bottlenecks

Operator Fusion and Multi-Stream Parallelism

The decoding phase of large models is limited by memory bandwidth and suffers from severe “serial dependencies.” To reduce decoding latency, optimization efforts focus heavily on execution flow reconstruction.

• Extreme Fusion of the Index Module: The three steps in the LI module—relevance scoring, head weight normalization, and Top-K filtering—are fused into a single operator. This operator directly reads paged key caches and utilizes register-level data reuse to complete all operations in one kernel call. This significantly reduces operator launch overhead on the device side and minimizes the allocation and destruction of intermediate tensors.
• Micro-Pipeline Parallel Scheduling (Indexer Prolog Three-Stream Collaboration): To eliminate pipeline bubbles caused by waiting for index information during decoding, a “Prolog” pre-execution mechanism was designed. By analyzing the data dependency graph, the query vector’s projection computation, the key vector’s normalization, the query vector’s quantization, and the head weight’s projection computation are assigned to three different compute streams for parallel execution. For example, while compute stream A executes the linear projection of the Q matrix, compute stream B can simultaneously execute the RMSNorm normalization of the K matrix. This fine-grained, instruction-level interleaving keeps hardware execution units (like vector and matrix units) highly utilized, hiding the decoding frontend preparation time within matrix multiplication operations and thereby reducing the absolute latency of a single decoding step.

Multi-Level Quantized Cache: Finding the Optimal Balance Between Precision and Memory

The memory bottleneck for model inference primarily lies in the expansion of the KV Cache. The target model employs a “role-specific, precision-specific” multi-level quantization strategy that is highly noteworthy:

• Extreme Compression for the Index Path (MXFP8): The LI module is not particularly sensitive to the relevance measurement of KV pairs. Therefore, its inputs (Index Q and Index K) use the MXFP8 format for grouped dynamic quantization, with each group sharing a single scaling factor. This fine-grained dynamic quantization maximally preserves the relative precision required for sorting while cutting the storage overhead of index data by 50%.
• Balanced Approach for the Main Attention Path (FP8 Tile Quantization): The KV Cache for the main attention computation path (SFA) uses the FP8 data type, but with a finer quantization granularity. Elements within each token and each attention head are quantized in blocks of a specific Tile size. Compared to coarse-grained per-token or per-head quantization, Tile-level quantization better adapts to the numerical distribution differences across various blocks in the attention score matrix. This maintains high-precision output while halving both the KV cache footprint and memory transfer time.

Through this hierarchical quantization scheme, the model can process longer context windows within a limited memory space or accommodate larger batch sizes under the same memory capacity, directly boosting throughput.

Long Sequence Parallel Strategies and Load Balancing

The Combination of Context Parallelism (CP) and Tensor Parallelism (TP)

When dealing with ultra-long sequence inputs (e.g., book-length text), the memory and compute capacity of a single GPU are insufficient. The mainstream strategy is Context Parallelism (CP). CP splits the input sequence along the temporal dimension into multiple shards, distributing them across different devices for parallel processing.

However, employing CP without considering load balancing leads to a severe “bucket effect” (where the slowest device dictates overall performance). This is especially true for the sparse attention model mentioned above: its index computation workload is uneven. Data blocks located later in the sequence require evaluation of more historical KV pairs, resulting in a computational load significantly heavier than that of the sequence’s earlier blocks.

Zigzag Layout: An Efficient Load Balancing Solution

To address this load imbalance, a sequence sharding layout strategy known as “Zigzag” is highly effective. Its core idea is to cleverly alter the mapping relationship of how sequences are distributed across devices:

• Traditional Sharding Mode: Sequence shards are distributed sequentially, meaning shard 1 goes to device 0, shard 2 to device 1, and so on. This results in device 0 having a much smaller computational load than device N-1, which handles the tail end of the sequence.
• Zigzag Layout: The sequence is paired head-to-tail. For example, shard 1 and shard N are assigned to the same device, while shard 2 and shard N-1 are assigned to the next device. This ensures that each device holds both a lighter head-region shard and a heavier tail-region shard. The two compensate for each other in computation time, converging the index computation latency across all devices to a consistent level. Once computation is complete, an inter-node AllGather communication restores the original logical order of the sequence, ensuring correctness for subsequent layer computations.

DSL-Driven Operator Extension

To achieve ultimate performance under the parallel strategy outlined above, a Domain-Specific Language (DSL) was used to rewrite core kernels. The programming abstractions provided by the DSL allow for finer control over the overlap between computation and communication. Specifically, when implementing CP functionality within a FlashKDA-style operator, the input can be further split along the sequence dimension. After local computation is completed, compute devices only need to exchange intermediate states at the boundaries and apply minor processing to yield the final result. This avoids introducing additional activation storage overhead and naturally pipelines communication operations with the next computation phase, significantly reducing the Time-to-First-Token (TTFT) latency in ultra-long sequence scenarios.

Conclusion: Moving Toward Engineering Maturity in LLM Inference

Through the inference optimization practices applied to this novel hybrid attention model, it is evident that the engineering response to ever-changing model architectures lies in building a systematic optimization framework. This includes:

1. Abstraction and Reuse: Distilling underlying universal computational paradigms and establishing a rich library of optimized operators as a technical foundation.
2. Data Flow Reconstruction: Reducing HBM access overhead and improving device utilization through operator fusion and compute graph-level optimization.
3. Adaptive Quantization Strategies: Flexibly adopting mixed precision and multi-level quantization based on the sensitivity to data precision at different computational stages.
4. Divide and Conquer with Load Awareness: Designing intelligent task distribution algorithms combined with parallel strategies to resolve severe temporal and spatial imbalances under long sequences.

These practices not only improve the deployment efficiency of this specific model but, more importantly, validate the adaptability and responsiveness of the entire underlying software stack when facing diverse new architectures. For developers, understanding the underlying computational patterns of models and gaining insight into hardware characteristics are essential for the successful, efficient deployment of large model applications.