Day0 Inference Optimization for Next-Generation Hybrid-Architecture LLMs Based on CANN

As large language models continue to scale in parameter size and context length, a single Transformer architecture struggles to balance computational complexity and inference quality for ultra-long sequences. Recently, a new hybrid architecture has emerged. The model alternately uses a state recurrence processing module (referred to as KDA below) and a sparse attention module with index selection (referred to as DSA below). Built on Multi-head Latent Attention (MLA), this architecture uses a lightning indexer to select the Top-K historical positions, and then a sparse flash attention module reads the corresponding KV cache to complete computation. It also aggregates cross-layer representations through attention residuals.

For such a complex new architecture, how the underlying AI infrastructure achieves Day0 adaptation and high-performance deployment has become a core challenge for developers. This article discusses how to reuse existing optimization efforts within the CANN heterogeneous computing architecture ecosystem to improve inference efficiency across computation, communication, scheduling, and storage.

1. Reusing Existing Optimizations for Rapid Adaptation

Day0 support for a new model is not built from scratch. The key is efficient reuse of common technologies. Based on earlier operator optimization work on similar architectures, this adaptation directly reused FlashKDA and MegaKDA operator optimizations, two-stage AttnRes fusion, and core operators such as MegaMoE and MC2.

For the overall execution graph, the solution combines hybrid parallelism, communication optimization, and multi-stream orchestration. It also introduces the SuperKernel fusion mechanism, graph-mode execution, compilation caching, and the NZ (Tile-N-Z) weight layout. In addition, it reuses the DSpark speculative inference and ReplaySSM state replay mechanisms. These accumulated techniques allow the underlying computation graph to complete efficient memory planning and operator fusion at compilation time, turning the adaptation cost of the new architecture into parameter configuration and graph structure fine-tuning.

1. DSA Operator Fusion and Three-Stream Parallelism to Break the Decode Serial Bottleneck

In the DSA computation flow, the lightning indexer needs to perform relevance scoring, head weight aggregation, and Top-K selection. Traditional implementations often split these steps into multiple independent operators, causing frequent intermediate result writes to memory and high kernel launch overhead.

2.1 Deep Operator Fusion

The CANN-based solution fuses the above three steps into a single kernel. The fused operator directly reads the underlying paged K cache, avoiding allocation and read/write of intermediate tensors. At the same time, SFA reads the corresponding KV from the main attention cache according to the selected indices to complete sparse computation. Both the prefill and decode phases support BF16 and quantized computation paths, ensuring flexibility across the entire pipeline.

2.2 Three-Stream Coordination for Indexer Prolog

During input preparation, Q projection and K normalization, Q quantization and head weight projection often have data dependencies. To reduce serial waiting, the solution introduces a three-stream coordination mechanism:

• Stream 1: Executes Q projection and K normalization.
• Stream 2: Executes Q quantization and head weight projection in parallel.
• Stream 3: Handles dependency triggering and synchronization.

Through a well-designed inter-stream dependency graph, computation-intensive tasks without dependencies are fully overlapped across multiple streams, significantly shortening decode wait time.

1. Hierarchical Dynamic Quantization to Reduce Memory and Compute Pressure

In LLM inference, the memory footprint and memory access bandwidth of the KV cache are often key bottlenecks. DSA retains BF16 high-precision projection computation to ensure model accuracy, while adopting a hierarchical quantization strategy for indexing and cache reads.

3.1 Full Quantization for the Indexer

In the lightning indexer stage, MXFP8 is used for dynamic quantization. The quantization granularity is per-group, with a group size of 32, meaning every 32 elements share one scaling factor. This fine-grained dynamic quantization greatly reduces computation while preserving numerical sensitivity during index scoring.

3.2 Block-wise Quantization for SFA Cache

For the KV cache read by SFA, FP8 is used, with quantization granularity refined to per-token-head-tile and a tile size of 128. This means quantization is applied in blocks of 128 elements within each token and each head.

# Pseudocode: hierarchical quantization configuration
quant_config = {
    "indexer_quant": {
        "type": "MXFP8",
        "granularity": "per_group",
        "group_size": 32
    },
    "sfa_kv_cache_quant": {
        "type": "FP8",
        "granularity": "per_token_head_tile",
        "tile_size": 128
    },
    "projection_keep": "BF16"
}

This strategy effectively reduces storage overhead for index data and KV cache while alleviating memory bandwidth pressure, allowing memory-bound operators to access data faster.

1. Long-Sequence Sharding and Load Balancing: Prefill CP and Zigzag Layout

For ultra-long sequences, KDA prefill time grows linearly with sequence length, significantly affecting time-to-first-token (TTFT). When only tensor parallelism is used to shard by attention heads, each device still needs to process the full sequence.

4.1 FlashKDA CP and Combined CPxTP Deployment

To address per-device memory and computation bottlenecks, the solution implements FlashKDA context parallelism (CP) based on CANNBot-DSL. This strategy shards input along the sequence dimension and distributes long-sequence computation across multiple devices. Through combined CPxTP deployment, together with sequence parallelism within TP groups, per-device activation values and intermediate tensor usage are greatly reduced, significantly lowering TTFT for ultra-long sequences.

4.2 Zigzag Layout for Causal Index Load Imbalance

In DSA causal index computation, positions later in the sequence have longer histories. With simple equal sharding, devices handling later sequence segments face heavy computation loads, while devices handling earlier segments remain idle.

To solve this load imbalance, the solution adopts a Zigzag layout: paired first-and-last sequence segments are assigned to the same device. The tail of a long sequence complements the head of a short sequence, balancing index computation across devices. After computation, communication operations restore the original data order to ensure logical correctness.

1. Summary and Developer Recommendations

The CANN-based inference optimization for next-generation hybrid-architecture LLMs demonstrates how underlying infrastructure can combine operator reuse, deep fusion, multi-stream orchestration, hierarchical quantization, and multi-dimensional parallelism to achieve efficient Day0 deployment for complex new architectures. For developers working on practical deployments, the following recommendations apply:

1. Fully leverage graph-mode execution and compilation caching in the underlying framework to avoid repeatedly building computation graphs at runtime and to shorten model startup time.
2. Pay attention to the match between memory and computation. For ultra-long-text scenarios, prioritize CP sharding and Zigzag layout, using TTFT as the primary optimization metric.
3. Use the general container platform and monitoring and alerting services provided by Baidu AI Cloud to observe GPU utilization and memory fluctuations under multi-stream orchestration in real time, and dynamically adjust speculative inference and state replay strategies for the best throughput-cost balance.

As LLM architectures continue to evolve, the generalization and forward-compatibility capabilities of the underlying computing framework will become critical factors in how quickly LLM-based applications can be brought to production.