Day0 Inference Optimization on Ascend CANN: A Deep Dive into Hybrid Architecture Acceleration

As large model parameter counts scale exponentially and new architectures such as mixture-of-experts (MoE) and state space models continue to emerge, the ability of the underlying computing framework to adapt to and optimize for these new architectures has become a decisive factor in model deployment efficiency. This article sets aside specific product attributes and focuses on Day0 inference optimization practices for next-generation hybrid architecture models running on heterogeneous computing platforms, using Ascend CANN as the background. It provides developers with a general-purpose guide to low-level optimization techniques across architecture analysis, operator fusion, layer-wise quantization, and long-sequence parallelism.

1. Inference Challenges and Core Mechanisms of Hybrid Architecture Models

Next-generation hybrid architecture models, such as those interleaving KDA (Kernelized Dynamic Attention) and DSA (Dynamic Sparse Attention) mechanisms, face several pain points when processing long sequences: high computational complexity, large memory footprint, and frequent cross-layer communication.

Interleaved State Recurrence and Sparse Attention

• KDA mechanism: Processes sequences through state recurrence, offering linear complexity. However, during the prefill phase, computation time grows significantly with sequence length, which directly lengthens time-to-first-token (TTFT).

• DSA mechanism: Introduces index selection on top of Multi-head Latent Attention (MLA). Specifically, a Lightning Indexer (LI) selects the Top-K historical positions, and Sparse Flash Attention (SFA) then reads the corresponding KV entries to perform sparse attention computation.

• Residual aggregation: The model aggregates cross-layer representations through Attention Residuals mechanisms, preserving information continuity but also placing higher demands on memory bandwidth and operator scheduling.

Reuse of Core Optimization Strategies

For these architectures, the inference engine reuses optimization accumulations from earlier stages, including FlashKDA/MegaKDA operator fusion, two-stage AttnRes fusion, and MegaMoE/MC2 MoE operator optimizations. Combined with hybrid parallelism, communication optimization, and multi-stream orchestration, the underlying framework introduces general acceleration techniques such as SuperKernel, graph mode, compilation caching, and NZ weight layout. Meanwhile, speculative inference and state replay mechanisms are retained to comprehensively improve inference efficiency across computation, communication, scheduling, and storage.

2. Computation Graph Fusion and Multi-Stream Orchestration: Breaking the Serial Decoding Bottleneck

In DSA (sparse attention) computation, index preparation and attention computation often have strict serial dependencies, leaving hardware resources idle. The optimization approach reduces decoding wait time through operator-level fusion and multi-stream coordination.

Deep Fusion of LI and SFA

The Lightning Indexer fuses relevance scoring, head weight weighting, and Top-K selection into a single large operator. At the low level, this fused operator directly reads paged KV cache, avoiding repeated writes of intermediate tensors to device memory and substantially reducing memory access overhead and operator launch overhead. Subsequently, SFA reads the main attention KV entries based on the selected indices to complete sparse computation. Notably, this fused path supports both BF16 and quantized computation in both the prefill and decode phases, ensuring efficient execution throughout the entire pipeline.

Three-Stream Coordination for Indexer Prolog

During input preparation, Q projection, K normalization, Q quantization, and head weight projection are traditionally executed serially. By introducing three-stream coordination, the system assigns Q projection and K normalization to stream A, and Q quantization and head weight projection to stream B, executing them in parallel at the hardware level. This multi-stream orchestration significantly reduces waiting time during input preparation, pushing compute resource utilization to its limits.

# Pseudocode: multi-stream orchestration logic
import torch
stream_a = torch.cuda.Stream()
stream_b = torch.cuda.Stream()
with torch.cuda.stream(stream_a):
    q_proj = compute_q_proj(input)
    k_norm = normalize_k(k_cache)
with torch.cuda.stream(stream_b):
    q_quant = quantize(q_proj)
    head_weight_proj = compute_head_weights(input)
torch.cuda.synchronize()
# Subsequent fused operator computation

3. Layer-Wise Mixed Quantization: Precisely Reducing Cache and Data Movement Overhead

Quantization is an effective means of reducing memory footprint and data movement bandwidth pressure. However, globally uniform quantization often leads to accuracy loss. This approach adopts a fine-grained layer-wise quantization strategy, matching different data types to the sensitivity of different computation modules.

Preserving Precision in Projection Computation

DSA projection computation is relatively sensitive to precision, so BF16 is retained to ensure accurate feature extraction.

Full Quantization of LI with MXFP8

The Lightning Indexer uses MXFP8 format for full quantization, with per-group dynamic quantization. The group size is set to 32, meaning every 32 elements share a scaling factor. This fine-grained quantization substantially compresses index data volume while adapting to the distribution characteristics of different numerical ranges, maintaining index selection accuracy.

FP8 Block Quantization for SFA KV Cache

For SFA KV cache processing, FP8 quantization is adopted with per-token-head-tile granularity. The tile size is set to 128, meaning quantization is performed in blocks of 128 elements within each token and each head. This strategy greatly reduces KV cache storage volume, and since the amount of data accessed is halved, it directly improves memory bandwidth utilization when SFA reads KV data.

4. Long-Sequence Sharding and Load Balancing: Breaking Single-Card Memory Limits

When input sequences exceed hundreds of thousands of tokens, a single card cannot hold the complete activations and intermediate tensors. At this point, tensor parallelism (TP) that shards by attention heads is no longer sufficient, and context parallelism (CP) must be introduced.

Combined CP and TP Deployment

The solution supports combined CP and TP deployment, sharding long inputs along the sequence dimension and combining this with intra-TP-group sequence parallelism. This effectively reduces per-card activation and intermediate tensor usage, making ultra-long sequence inference feasible.

Zigzag Layout for Load Imbalance

A notable characteristic exists in DSA causal index computation: the later a sequence position is, the longer its history and the greater its computation load. If sequences are simply sharded evenly in order, the later compute nodes become severe bottlenecks. To address this, the solution adopts a Zigzag layout: two sequence segments paired from the beginning and end are assigned to the same compute node. For example, the first segment (short history) is paired with the last segment (long history), and the second segment is paired with the second-to-last segment. This long-short pairing strategy effectively balances cross-card index computation loads. After computation completes, the original order is restored to ensure logical correctness.

5. FlashKDA Context Parallelism with State Coordination via Low-Level DSL

In ultra-long sequence scenarios, KDA prefill computation time grows linearly with sequence length, severely impacting TTFT. To address this pain point, the framework implements FlashKDA CP based on a low-level domain-specific language (DSL).

Computation Distribution and Parallel State Precomputation

FlashKDA CP further shards input along the sequence dimension, distributing long-sequence computation across multiple cards. Before parallel computation, each card needs the state summary from preceding shards. The implementation logic is as follows: each card first precomputes its own shard’s state summary in parallel, then performs an AllGather communication operation to globally exchange state summaries.

Sequential Synthesis and Local Computation

After collecting all state summaries, each card synthesizes the state contributions from preceding shards in sequence order to derive the starting state for its own shard. With accurate starting states, each card can then complete its local computation fully in parallel. This mechanism not only substantially reduces TTFT for ultra-long sequences but also improves overall system concurrency by reducing per-card activation usage.

6. Summary and Engineering Practice Recommendations

The evolution of large model architectures is driving rapid iteration of the underlying AI infrastructure. Based on practical experience, successfully adapting and optimizing next-generation hybrid architecture models requires developers to focus on the following aspects:

• Operator fusion must go deep into memory access details: It is not just about merging operators, but also about direct reads of paged cache and elimination of intermediate results.

• Quantization strategies must be tailored to local conditions: From global BF16 to per-group MXFP8, and then to per-token-head-tile FP8, layer-wise quantization is the necessary path to balancing precision and performance.

• Parallelism strategies must break dimensional constraints: For ultra-long contexts, the combination of CP and TP, Zigzag load balancing, and DSL-based state shard synthesis are core means of breaking single-card bottlenecks.

In actual deployment, developers are advised to deeply understand both model architecture and low-level hardware characteristics, and to fully leverage mechanisms such as graph mode and compilation caching provided by the underlying framework to achieve efficient inference support for new architectures at the Day0 stage.