Day0 Inference Optimization for New Architecture LLMs Based on CANN

Introduction

With the continuous evolution of open-source large language model (LLM) architectures, novel hybrid attention mechanisms and sparse computing impose greater challenges on underlying inference frameworks. How to provide efficient inference support on Day0—the first day a model goes open source—while fully reusing past optimization accumulations for rapid adaptation, has become a key issue in the AI infrastructure domain. Using a cutting-edge new architecture LLM as an example, this article deeply analyzes the core technical paths to achieving Day0 inference support and performance optimization based on CANN (Compute Architecture for Neural Networks), providing a reference for developers deploying models on the Ascend platform.

1. New Architecture Features and Reuse Optimization Strategies

1.1 Interleaved Architecture and State Recursion Mechanism

The execution logic of this new architecture LLM interleaves two core modules: the State Recursion Module (KDA) and the Dynamic Sparse Attention Module (DSA).

• KDA Module: Processes sequence information through state recursion to maintain the coherence of historical context.
• DSA Module: Introduces an index selection mechanism based on Multi-head Latent Attention (MLA). Its core workflow consists of two parts: first, a lightweight indexer filters the Top-K historical positions most relevant to the current token; second, a Sparse Flash Attention (SFA) precisely reads the corresponding KV cache based on the filtered indices to complete the sparse attention computation. Additionally, the model aggregates cross-layer representations through an attention residual mechanism to enhance feature transmission efficiency.

1.2 Historical Operator Accumulation and New Architecture Adaptation

The core of achieving Day0 rapid adaptation lies in “maximizing reuse.” Leveraging previous optimization experience, the inference for this new architecture directly reuses and upgrades efficient state recursion operators like FlashKDA/MegaKDA, the AttnRes two-stage fusion operator, and expert network optimization operators like MegaMoE/MC2.
At the overall scheduling level, the platform combines hybrid parallel strategies, communication link optimization, and multi-stream orchestration technologies. It introduces SuperKernel fusion, graph mode execution, compilation cache mechanisms, and NZ (Narrow-Z) weight layout, while continuing to utilize speculative decoding and state replay mechanisms. This combination comprehensively improves inference efficiency across computation, communication, scheduling, and storage dimensions, significantly reducing the R&D cycle for new architecture adaptation.

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

During the execution of the DSA module, the degree of fusion between index computation and attention computation directly determines the waiting latency in the decoding phase.

2.1 Operator-Level Deep Fusion

In traditional implementations, relevance scoring, head weight summation, and Top-K selection often exist as serial operators, generating significant write-backs of intermediate results to memory and operator launch overhead. In CANN’s optimized implementation, these three steps are fused into a single operator, supporting direct reading of the paged K cache, effectively eliminating intermediate memory搬运overhead. The subsequent SFA operator directly reads the main attention KV based on the selected indices to complete sparse computation. To accommodate different deployment scenarios, both the Prefill and Decode phases support BF16 and quantized computation paths for the aforementioned indexer.

2.2 Indexer Prolog Three-Stream Coordination

In the input preparation phase, CANN further implements Indexer Prolog optimization. Through a three-stream coordination mechanism, independent computing tasks such as Q projection and K normalization, as well as Q quantization and head weight projection, are executed in parallel. This multi-stream orchestration breaks the original serial waiting chain, pushing computing resource utilization toward its peak.

3. Hierarchical Quantization Strategy: Reducing Cache and Memory Access Overhead

The memory bottleneck in LLM inference is typically concentrated on the KV cache and intermediate index data. Instead of adopting a one-size-fits-all quantization scheme, this new architecture model employs a finely tuned hierarchical quantization strategy based on the distinct characteristics of DSA and KDA.

3.1 Preserving High-Precision Projection Computation

To ensure the accuracy of the model’s foundational feature extraction, the projection computation layer in DSA retains BF16 precision, avoiding irreversible accuracy loss from quantization.

3.2 MXFP8 Quantization for Index Data

During the LI (Lightning Indexer) phase, full quantization adopts the MXFP8 data format. The system dynamically quantizes the index’s Q and K in a per-group manner, with a group size of 32, meaning every 32 elements share a single scaling factor. This fine-grained dynamic quantization significantly compresses the data bit-width of index computations while maintaining numerical distribution fitting accuracy.

3.3 FP8 Quantization for SFA KV Cache

For SFA’s KV cache management, FP8 quantization is applied with a granularity of per-token-head-tile and a tile size of 128. This means quantization is performed in blocks of 128 elements within the dimensions of each token and each head.
This hierarchical quantization strategy effectively reduces the storage footprint of index data and KV cache, while simultaneously alleviating DDR bandwidth pressure and on-chip transfer overhead through a more compact data layout.

4. Long Sequence Partitioning and Load Balancing Optimization

Facing the growing demand for long sequence processing in LLMs, relying solely on Tensor Parallelism (TP) is no longer sufficient to meet memory and latency constraints. CANN introduces a combined deployment scheme of Context Parallelism (CP) and Sequence Parallelism in this adaptation.

4.1 CPxTP Combined Deployment

The system supports a combined deployment mode of CP and TP. CP splits long inputs along the sequence dimension and, combined with the sequence parallelism mechanism within the TP group, distributes large segments of input across multiple accelerator cards for processing. This significantly reduces the memory footprint of activations and intermediate tensors on a single card.

4.2 Zigzag Layout for Causal Index Load Imbalance

A natural problem exists in DSA’s causal index computation: “the later a token is in the sequence, the longer the history it needs to traverse,” leading to extreme imbalances in computing load across different cards. To address this, CANN introduces a Zigzag layout: paired head and tail sequence segments are assigned to the same card for processing. This complementary head-to-tail allocation effectively balances the index computation volume across cards. After each card completes its computation, communication operations are used to restore the original order, ensuring logical correctness.

5. DSL-Based Optimization for Ultra-Long Sequences

In ultra-long sequence scenarios, the computation time of the KDA module during the Prefill phase grows linearly or even exponentially with sequence length, severely impacting the Time To First Token (TTFT). If only TP is used to split by attention heads, each card still needs to process the complete ultra-long sequence.

5.1 FlashKDA CP Mechanism

To solve this pain point, CANN implemented FlashKDA CP based on the underlying DSL (Domain Specific Language). This mechanism further splits the input along the sequence dimension, completely distributing the state recursion computation of long sequences across multiple cards. This not only significantly reduces TTFT in ultra-long sequence scenarios but also further cuts down the memory footprint of activations and intermediate tensors on a single card.

5.2 The Future of DSL and Operator Customization

Developing based on DSL is a critical path to achieving efficient underlying operator customization. Through DSL, developers can describe computation logic at a higher level of abstraction, and the compiler automatically generates instruction streams optimized for Ascend hardware characteristics. In this practice, each card first computes the necessary state variables in parallel, and then completes state synchronization through lightweight communication, perfectly supporting the implementation of FlashKDA CP. This is also a typical example of how future LLM inference frameworks will move toward “deep customization” and “software-hardware co-design.”

Conclusion

Day0 inference support for new architecture LLMs is not only a test of the adaptability of underlying computing frameworks but also an examination of the ability to reuse historical technical assets. By deeply fusing DSA operators, applying refined multi-stream orchestration and hierarchical quantization strategies, and combining long sequence optimization methods such as CP and Zigzag layouts, CANN has successfully achieved the efficient deployment of new architecture models. These practical experiences provide an actionable engineering paradigm and optimization insights for developers dealing with increasingly diverse LLM architectures in the future. With the support of infrastructure like Baidu AI Cloud, such underlying optimizations will further enhance the deployment efficiency of LLMs in real-world business scenarios.