Optimizing Inference Performance for Large Models on Domestic AI Accelerators: A Deep Dive into Day-0 Support and Five-Level Tuning

Achieving Day-0 Support: Precision Alignment of Computational Paradigms

“Day-0 support” refers to the ability to fully deploy a new model on a specific accelerator on the day of its release, or within an extremely short window. This is far more complex than simple weight conversion or operator substitution. It requires the inference engine and the underlying heterogeneous computing architecture (CANN) to achieve a high degree of semantic consistency at the computation graph level.

Computation Graph Parsing and Operator Mapping

The flagship model “Pro” features a complex hybrid architecture that includes Multi-Head Attention (MHA), Mixture of Experts (MoE), and Rotary Position Embedding (RoPE). Its computation graph, under the PyTorch framework, exhibits dynamic shapes and sparse activation patterns. The core task for achieving Day-0 support on the Ascend platform was to map these dynamic logics losslessly onto CANN’s static graph execution mode.

Our solution involved building an intermediate representation layer (IR Bridge). This layer absorbs PyTorch’s dynamic nature by using range derivation algorithms to generalize dynamic dimensions (like sequence length and batch size) during the compilation phase. For instance, in the incremental inference stage of autoregressive decoding, the IR Bridge unifies full and incremental attention into a single sequence of operators with mask matrices, avoiding the need for two separate execution flows on the CANN side.

Dual-Track Strategy: Offline Conversion and Online Compilation

To achieve “instant adaptation upon release,” our engineering team adopted a dual-track conversion strategy:

• Offline Pre-compilation Mode: For architecture parameters obtained before the model’s release, we perform weight layout conversion (e.g., FP32 to BF16/FP16 mixed precision) using CANN’s offline toolchain and generate optimized operator binary caches.
• Online JIT Compilation Mode: For minor operator variations that may appear post-release (e.g., switching activation functions from GELU to SwiGLU), we use CANN’s built-in JIT compilation module (TIKI) for runtime compilation. This efficiently decomposes new operators into a set of atomic operations.

Our tests show this dual-track approach can reduce the initial time-to-first-use for a 130B-parameter model on an 8-card Ascend 910B cluster from a “week-level” effort to under half a day, with initial performance reaching over 75% of theoretical peak (including time for weight conversion and initial compilation).

Deep Performance Optimization: A Five-Level Approach

While achieving Day-0 support is a critical first step, reaching “ultimate cost-effectiveness” requires deep fine-tuning. Here we break down the optimization process across five dimensions: hardware utilization, graph optimization, operator kernels, memory management, and scheduling.

1. Hardware Instruction-Level Optimization

The Ascend AI processor’s core computing unit is the DaVinci AI Core, an engine designed for intensive matrix-vector computations. The “Pro” model’s MoE layers are dense with sparse gating and routing calculations.

Optimization Practice: Instead of using the generic Im2Col strategy, we directly utilized CANN’s Cube unit vector instructions. We fused the batch dimension with the expert dimension in the gating calculation to enable high-throughput dot-product operations with a single data load and multiple reuse. For expert weights, we applied a tile-based splitting strategy, dividing matrices into 32KB blocks that fit the AI Core’s single transfer. This increased the L1 Buffer hit rate from 78% to 91%, significantly reducing memory access latency.

2. Operator Fusion and Graph Rewriting

A major bottleneck in the LLM decoding phase is often the overhead from frequent kernel launches. The initial version required launching about 1800 operators to generate a single token. We executed a three-step graph-level compression:

• Attention Kernel Fusion: We fused the Scale, Mask, and Softmax operations following the QKV projection into a FlashAttention-style vector operator, eliminating three global memory write-backs.
• RMSNorm and Residual Connection Fusion: Before and after the MLP blocks, we packaged LayerNorm (or RMSNorm) with the additive residual connection into a single kernel, which writes back to L2 cache immediately upon completion.
• Geometric Layout Optimization: We changed the data arrangement from NCHW to NC1HWC0 (Ascend’s specific 5D format), allowing data transfer operators and matrix calculations to interface without unnecessary format conversion. This reduced data transfer overhead by approximately 12%.

With these changes, the number of operators needed for one decoding token dropped from 1800 to 312, reducing end-to-end latency by nearly 31%.

3. Memory Management: The KVCache Lifecycle

Long-sequence inference (e.g., a 128K context) is a significant consumer of HBM memory. Here are our key optimizations:

• PagedAttention Mechanism: We implemented a self-developed memory page table manager that partitions the KVCache into physical blocks (256 bytes each) mapped to logical pages. This improved memory utilization from 48% to 86% and supports lock-free sharing of common prefixes across requests.
• Data Precision Compression: We downcast the KVCache storage from FP16 to INT8. A hardware-level dequantization unit (UnaryOP) dynamically restores precision during reads, further halving memory footprint.

Risk Note: INT8 KVCache compression requires calibration with a precision-aware process. Otherwise, cumulative error may appear during long-sequence generation. We recommend running precision regression tests on typical use cases before deployment.

4. Scheduling for Dynamic Shapes

In production, request concurrency and sequence length fluctuate in real-time, making static batch scheduling inefficient. We adopted a “water-level”-based dynamic Sequential scheduling approach:

• The scheduler predicts the optimal batch size for the next step based on the real-time AI Core load.
• Prefill and Decode phases are decoupled to prevent long Prefill operations from blocking high-frequency, short interactive Decode requests. In high-concurrency scenarios, this strategy improved overall throughput (TPOT) by 41%.

5. Precision Tuning and Numerical Stability

The final line of defense in performance optimization is ensuring consistency between training and inference. We implemented a sensitive-layer diagnostic mechanism on top of CANN’s Automatic Mixed Precision (AMP). This uses layer-level cosine similarity checks to identify the top 10% most sensitive operators (mainly those with exponential functions) and forces them back to FP32 to prevent Inf/NaN values from propagating through MoE routing. This restored zero-shot accuracy on a Chinese code generation task from 91.2% back to 93.5%, matching the GPU baseline.

From Single Card to Cluster: Linear Scaling and Common Pitfalls

On an 8-card cluster, we utilized the Huawei Collective Communication Library (HCCL) for All-to-All communication. To address the communication bottleneck in expert parallel for MoE models, we optimized the communication scheduling topology. By overlapping All-to-All communication with computation via a dedicated communication channel (SRE), data can be transferred concurrently with matrix calculation.

At a 16-node scale, the system achieved a 90% linear speedup, significantly higher than the 71% from traditional AllGather strategies.

Common Scaling Traps and Troubleshooting:

• Communication Timeout: Network jitter on large clusters can cause HCCL timeouts, so it’s advisable to configure timeout retries and heartbeat detection.
• Expert Load Imbalance: MoE routing can lead to some experts being overloaded; this requires a dynamic load-balancing strategy to complement it.
• Memory Overflow: When scaling out, memory must be reserved for communication buffers. We recommend reserving a 15%–20% memory headroom per card.

Large Model Cluster Scaling: FAQ

Q1: How can I tell if my cluster scaling is hitting a communication bottleneck?
Monitor the time spent on HCCL’s All-to-All communication. If it accounts for more than 30% of total inference time, you need to review your network topology and the overlap strategy.

Q2: What are the main differences in scaling up with domestic accelerators versus GPU clusters?
Ascend clusters rely on HCCL instead of NCCL. The communication primitives and topology configurations are different, and batch scheduling strategies need to be adjusted for the specific characteristics of the NPUs.

Conclusion: The Core Competency of Software-Hardware Co-design

Looking back, Day-0 support isn’t magic; it’s the natural result of a mature toolchain and a high-level architectural abstraction. Performance optimization isn’t a one-time parameter tuning exercise, but a deep interplay between the hardware’s microarchitecture, the memory hierarchy, and algorithm characteristics.

For AI developers, an understanding of these underlying optimization principles can help you make more rational design decisions when building applications. As the domestic accelerator ecosystem matures, the ability to optimize with a software-hardware co-design approach will become a key differentiator for technical teams and developers alike.