Day0 Adaptation and Performance Optimization for a New Long-Sequence Inference Model

1. Background: Adaptation Challenges of a New Hybrid Architecture

A new class of long-sequence inference models has recently been released as open source. The architecture introduces a significant departure from prior designs: it interleaves two attention mechanisms. The first processes sequential information through state recurrence. The second builds on multi-head latent attention with a learned indexer that selects the Top-K historical positions, after which a sparse flash attention kernel reads the corresponding KV pairs to perform sparse attention computation. Cross-layer representations are still aggregated via attention residuals.

This hybrid design creates fresh challenges for inference deployment: the sparse index computation exhibits causal load imbalance across positions, KV cache grows rapidly with sequence length, and time-to-first-token (TTFT) becomes hard to control for ultra-long sequences. To help developers deploy and experiment on the Ascend platform quickly, the team delivered Day0 support and deep optimizations across computation, communication, scheduling, and storage.

2. Reuse and Extension: From Prior-Generation Optimizations to Rapid Adaptation

Adapting the new model did not start from scratch. Multiple key operators and mechanisms from the previous generation were carried forward:

• State recurrence operator optimization: Reused operators such as FlashKDA and MegaKDA to accelerate state-space computation.
• Attention residual fusion: Kept the two-stage fusion approach to reduce intermediate overhead in cross-layer aggregation.
• Mixture-of-Experts optimization: Reused operators including MegaMoE and MC2 to lower expert routing and communication costs.
• Parallelism and scheduling: Leveraged hybrid parallelism, communication optimization, and multi-stream orchestration to improve pipeline efficiency.
• Compilation and layout: Applied SuperKernel, graph mode, compilation caching, and NZ weight layout to cut runtime overhead.
• Speculative inference and state replay: Carried over speculative decoding with state replay to accelerate generation while preserving state consistency.

These accumulated assets allowed the new architecture to be adapted in short order, yielding inference gains across compute, communication, scheduling, and storage.

3. Sparse Index Fusion and Multi-Stream Collaboration: Reducing Decode Wait Time

Sparse attention involves two distinct stages: index selection and sparse computation. This optimization fuses the indexer’s relevance scoring, head-weight weighting, and Top-K selection into a single kernel that reads directly from paged KV cache, eliminating intermediate write-backs and kernel launch overhead. The sparse flash attention then fetches the main attention KV using the selected indices to complete the computation. Both the Prefill and Decode phases support BF16 and quantized paths for the indexer.

Further, the indexer prolog uses a three-stream pipeline: Q projection with K normalization and Q quantization with head-weight projection run concurrently. This multi-stream orchestration increases overlap across compute units and reduces decode-stage latency.

4. Hierarchical Quantization: Reducing Cache and Transfer Overhead

KV cache storage and memory-access cost is a primary bottleneck in long-sequence inference. A hierarchical quantization strategy was adopted:

• Full indexer quantization: Uses MXFP8 format with per-group dynamic quantization on index Q/K (group size = 32, i.e., one scale factor per 32 elements), preserving index accuracy while sharply cutting storage and access cost.
• Sparse attention KV cache: Uses FP8 quantization at per-token-head-tile granularity (tile size = 128), meaning each block of 128 elements within a token and head is quantized independently. This maintains attention quality while reducing KV cache footprint and transfer volume.
• Projection computation kept in BF16: Sparse attention projections remain in BF16 to ensure numerical stability.

Together, these policies free up memory and bandwidth for longer sequences.

5. Prefill Context Parallelism: Sequence Partitioning and Load Balancing

For long-sequence Prefill, a CP×TP configuration is supported: the input is split along the sequence dimension, and intra-TP-group sequence parallelism further reduces per-device activation and intermediate tensor memory.

However, causal index computation in sparse attention is inherently imbalanced—later positions must attend to longer histories. A naive sequence split would overload later devices. To address this, a zigzag layout pairs head and tail segments on the same device, balancing index workload across devices. After computation, the original order is restored. This layout adds no extra communication while significantly improving parallel efficiency.

6. FlashKDA Context Parallelism: State-Space Parallelism for Ultra-Long Sequences

Under ultra-long sequences, Prefill cost in the state-recurrence component scales with sequence length, directly impacting TTFT. Tensor parallelism across attention heads alone leaves each device processing the full sequence, so sequence-dimension workload is not reduced.

FlashKDA CP extends partitioning along the sequence dimension, distributing the long-sequence computation across multiple devices to lower TTFT and per-device activation memory. Implemented via CANNBot-DSL, it coexists with existing TP strategies, offering finer-grained parallelism for ultra-long inference.

7. Practical Recommendations and Summary

Developers deploying this long-sequence inference model on the Ascend platform should consider the following:

1. Enable hierarchical quantization first: When memory or bandwidth is constrained, MXFP8 for the indexer and FP8 for the KV cache cut cache overhead substantially—verify that accuracy meets requirements.
2. Configure CP and TP wisely: Choose CP×TP based on sequence length and hardware topology; prefer the zigzag layout for long sequences to balance load.
3. Leverage multi-stream and fused operators: The indexer prolog’s three-stream pipeline and fused index execution effectively reduce decode latency—keep them enabled during Decode.
4. Activate FlashKDA CP for ultra-long sequences: When sequence length exceeds single-device capacity, FlashKDA CP splits state-recurrence work to reduce TTFT.

In summary, this Day0 adaptation and optimization effort—spanning operator fusion, hierarchical quantization, sequence partitioning, load balancing, and state-space parallelism—provides a complete path for efficient deployment of the new long-sequence inference model on the Ascend platform. Developers can build on these capabilities to achieve lower latency and higher throughput in long-sequence scenarios.