Inference Acceleration for Next-Gen Attention Architectures: From Operator Reuse to Parallel Orchestration

As large model architectures continue to evolve, the sparsification and stateful transformation of attention mechanisms have become the mainstream direction for reducing long-sequence inference costs. A recently released next-generation model introduces an interleaved hybrid attention structure based on its predecessor and has been fully open-sourced. The accompanying acceleration framework completed adaptation and optimization on launch day, enabling developers to quickly deploy and validate on domestic acceleration platforms. From an engineering perspective, this article outlines the technical workflow of this adaptation: how to quickly support new structures while leveraging existing optimization assets, and how to systematically improve inference efficiency across four dimensions—computation, communication, scheduling, and storage.

1. Core Architectural Changes: Hybrid Attention and Cross-Layer Aggregation

To understand the optimization approach, we must first clarify the structural changes in the new model. The model interleaves two types of mechanisms in its attention layers:

• State-Recursive Attention (denoted as KDA): Processes sequences through state recursion, avoiding explicit modeling of all historical positions. This offers better computational and storage characteristics for long sequences.
• Sparse Index Attention (denoted as DSA): Introduces an index selection mechanism based on Multi-head Latent Attention (MLA). A lightweight indexer (Lightning Indexer, LI) first scores historical positions for relevance and selects the Top-K. Then, a Sparse Flash Attention module (SFA) reads the corresponding KV based on these indices to complete the attention calculation.

Additionally, the model continues to aggregate cross-layer representations through an Attention Residuals mechanism to enhance deep information传递传递传递 transmission传递传递传递.

The engineering challenges brought by this structure are clear: the index selection phase introduces additional scoring, weighting, and sorting operations, while the sparse reading phase requires accessing the KV cache based on non-contiguous indices. Both cause operator fragmentation, frequent write-backs of intermediate results, and non-contiguous memory access. If implemented operator by operator in a conventional manner, the latency during the decoding phase will significantly increase.

2. Underlying Logic of Fast Adaptation: Reusing Existing Optimization Assets

The new model was not adapted from scratch but inherited the optimization assets accumulated on the acceleration framework by its predecessor. This is the key to achieving “day-zero support.” The reusable capabilities fall into four main categories:

1. State-Recursive Operator Optimization: Fused operator implementations for KDA-like structures that merge multiple steps in recursive computations, reducing kernel launch overhead and intermediate tensor materialization.
2. Two-Stage Fusion of Attention Residuals: Merges read and write operations during cross-layer aggregation into a two-stage pipeline, reducing memory bandwidth pressure.
3. MoE and Communication Fusion Operators: Fused operators and communication-computation overlap schemes tailored for expert parallelism scenarios, reducing the exposed time of All-to-All communications.
4. Inference Scheduling and Storage Mechanisms: Includes speculative decoding, state replay mechanisms, as well as general techniques like SuperKernel fusion scheduling, graph-mode execution, compilation caching, and weight layout optimization.

On this foundation, the new architecture only required incremental adaptation: index fusion for the DSA path, hierarchical quantization strategies, and sequence parallel splitting for long sequences. This “asset reuse + incremental gap-filling” strategy significantly compressed the adaptation cycle and mitigated the risk of performance regression when launching the new model.

3. DSA Fusion and Multi-Stream Orchestration: Reducing Decoding Latency

3.1 Three-Stage Indexer Fusion

The indexer’s original computation can be split into three steps: relevance scoring, head weight weighting, and Top-K selection. If implemented sequentially, each step requires writing intermediate results back to global memory before launching the next operator. The accumulated launch and memory access overheads are particularly prominent during the decoding phase.

The optimization strategy fuses these three steps into a single execution unit and directly reads the paged KV cache, avoiding repeated write-backs of intermediate results. The benefits are twofold: first, the number of operator launches decreases; second, index data no longer needs to be moved back and forth in memory. The indexer supports both BF16 and quantized paths in both Prefill and Decode stages, allowing flexible selection based on precision requirements.

3.2 Index-Based Reading in Sparse Attention

Based on the Top-K indices output by the indexer, the SFA reads the corresponding positions from the main attention KV cache to perform sparse computation. The key here lies in the memory access pattern: indices are non-contiguous, requiring efficient paged addressing and gather-read strategies to prevent bandwidth wastage.

3.3 Three-Stream Coordination in the Pre-Index Phase

Before index computation, the input preparation phase contains several parallelizable sub-tasks: there is no dependency between Q projection and K normalization, nor between Q quantization and head weight projection. The optimization uses three-stream coordination to assign these operation sets to different execution streams, transforming the originally sequential preparation phase into a pipelined execution and reducing sequential waiting in the first step of decoding.

4. Hierarchical Quantization: Reducing Cache Footprint and Transfer Overhead

Quantization strategies must match the precision sensitivity of each module rather than applying a one-size-fits-all approach. The solution adopts a hierarchical design:

Module	Precision Strategy	Quantization Granularity
DSA Projection Computation	Retain BF16	—
Indexer Q/K	Fully Quantized MXFP8	per-group, group size 32
Sparse Attention KV Cache	FP8	per-token-head-tile, tile size 128

Specifically:

• Projection computation retains BF16 to ensure the numerical stability of the main attention path.
• The indexer uses MXFP8 for full quantization, applying dynamic quantization to the index Q/K on a per-group basis, with a shared scaling factor for every 32 elements. This finer grouping granularity maintains accuracy in scoring scenarios that are sensitive to relative ranking.
• The KV cache uses FP8 with a quantization granularity of per-token-head-tile, meaning elements are quantized in blocks of 128 within each Token and each Head. This granularity balances memory access continuity with quantization error control.

Overall, this reduces the storage footprint and memory bandwidth requirements for both index data and the KV cache, thereby alleviating memory pressure during long-sequence decoding.

5. Sequence Parallelism and Load Balancing: Two Types of Splitting for Long Sequences

5.1 CP×TP Combined Deployment

In long-input scenarios, activations and intermediate tensors on a single card can easily become a bottleneck. The solution supports combined deployment of Context Parallelism (CP) and Tensor Parallelism (TP): splitting long inputs along the sequence dimension while performing sequence parallelism within the TP group, thereby reducing single-card activation footprint and intermediate tensor scale.

5.2 Zigzag Layout for Balancing Causal Index Load

Causal index computation in DSA has a natural load imbalance: the later a position is in the sequence, the longer its visible history, and thus the heavier its index computation load. If split by simple contiguous chunks, the later cards will bear a disproportionately heavy computation burden.

The solution adopts a Zigzag layout: pairing the beginning and ending segments of the sequence and assigning them to the same card. This ensures each card holds both an early segment with “short history” and a late segment with “long history,” balancing the cross-card index computation load. After computation, the layout is restored to its original order to ensure correctness.

5.3 Sequence Splitting for State-Recursive Modules

In ultra-long sequence scenarios, the Prefill computation time of KDA grows linearly with sequence length, directly impacting Time To First Token (TTFT). When splitting only by attention heads using TP, each card still processes the full sequence, failing to shorten the critical path.

To address this, the solution implements Context Parallelism for KDA with the support of a Domain-Specific Language (DSL): it further splits the input along the sequence dimension, distributing the long sequence computation across multiple cards. The goal is to reduce TTFT for ultra-long sequences while decreasing single-card activation and intermediate tensor usage. It is worth noting that state-recursive structures require handling cross-shard state transmission after sequence splitting, which is the most complex part of implementing this module.

6. Engineering Implementation Recommendations

Based on the strategies discussed above, developers should focus on the following points during actual deployment:

1. Prioritize reusing mature operator libraries. Apply incremental fusion only to new structures to avoid performance regressions caused by complete rewrites.
2. Maximize fusion for index-based computations. Merge scoring, weighting, and selection into a single execution and directly consume paged caches.
3. Apply hierarchical quantization granularity based on modules. Use finer grouping granularity for ranking-sensitive scoring paths.
4. Prioritize enabling sequence parallelism in long-sequence scenarios, paired with load-balancing layouts to handle computational unevenness in causal attention.
5. Leverage graph mode and compilation caching to reduce repeated compilation and scheduling overhead, shortening service cold start times.

Conclusion

Optimizing inference for new architectures is not a single-point breakthrough but the result of synergizing computation fusion, communication overlap, scheduling orchestration, and storage layout. Inheriting existing optimization assets enabled rapid adaptation for the new model, while fused execution for the sparse index path, hierarchical quantization, and sequence parallel splitting addressed the specific performance bottlenecks of the new structure. For inference services targeting long-sequence, high-concurrency scenarios, this approach offers excellent transferability: the more reusable optimization assets available, the lower the adaptation cost when the next shift in model architecture arrives.