A Technical Breakdown of Day-0 Adaptation for Domestic AI Chips Running 3-Trillion-Parameter MoE Models
This article analyzes the technical path for adapting domestic AI chips to support ultra-large-scale, 3-trillion-parameter Mixture of Experts (MoE) models for Day-0 deployment. It covers memory planning strategies for massive parameters, execution chains for hybrid attention mechanisms, and heterogeneous computing demands for multimodal inputs. The piece further details the decoupled software stack architecture, pipeline optimizations for model loading, and key inference path optimizations including sparse attention integration, mixed-precision workspace reuse, and topology-aware communication scheduling. It concludes with engineering practices for regression testing and performance benchmarking, offering insights into building scalable software ecosystems for next-generation AI infrastructure.
Background: Deployment Challenges of 3-Trillion-Parameter MoE Models
As Large Language Models (LLMs) exceed the trillion-parameter threshold, their architectural complexity continues to rise. Consider the industry’s current frontier 3-trillion-parameter Mixture of Experts (MoE) models. While they utilize a sparse activation mechanism—activating only a portion of the expert networks per inference to reduce computational overhead—this places extreme demands on the inference system’s memory management, communication bandwidth, and operator scheduling.
The private deployment of such ultra-large-scale models typically faces three core challenges:
- Massive Weight Storage: Storing the complete model weights requires terabytes of memory, posing a severe test to both per-card memory capacity and multi-card memory aggregation efficiency.
- Complex Inference Paths: The models employ sparse attention mechanisms alongside multimodal vision encoders. The inference path intertwines dense matrix computations with sparse indexing operations, which traditional inference engines cannot directly and efficiently support.
- Mixed Precision Requirements: To balance inference speed and memory footprint, model weights often use low-bit quantization formats (such as MXFP4). This requires the inference framework to possess online format conversion and mixed-precision computation capabilities.
This article systematically outlines the technical path for domestic AI chips to achieve Day-0 adaptation for ultra-large-scale models, focusing on how the software stack ecosystem supports the rapid migration and stable operation of new model architectures.
I. Model Architecture Characteristics and Computational Demands
1.1 Memory Planning Strategy for Massive Parameters
Taking a 3-trillion-parameter model as an example: even with 4-bit quantization (approximately 0.5 bytes per parameter), the complete model weights still require about 1.5TB of storage space. Given that mainstream AI accelerator cards currently offer 64GB to 96GB of memory per card, a minimum of 16 to 24 cards are required just to load the weights fully.
Crucially, MoE architectures also require additional memory for KV Cache and intermediate activations during inference. With a 1-million-token context window, the memory footprint of the KV Cache grows linearly with sequence length, creating secondary pressure on memory capacity. Therefore, during the Day-0 adaptation phase, the primary challenge is distributing model weights and dynamic caches reasonably across a multi-card cluster using a combination of Tensor Parallelism and Expert Parallelism.
1.2 Execution Chain Analysis of Hybrid Attention Mechanisms
New-generation trillion-parameter models generally adopt a hybrid architecture design: using traditional dense attention in global layers to capture long-range dependencies, and introducing sparse attention (such as Dynamic Sparse Attention, or DSA) in local layers to reduce computational complexity. Mathematically, this design is equivalent to imposing structured sparsity constraints on the attention matrix. However, in practical implementation, the inference engine must simultaneously maintain two computation graphs:
- Dense Attention Path: Executes standard QKV projection, scaled dot-product attention, and output projection.
- Sparse Attention Path: Requires completing Index Selection and KV cache layout rearrangement before executing block matrix multiplication based on sparse masks.
This hybrid execution chain requires the inference framework to have flexible graph compilation capabilities, enabling the seamless integration of different forms of attention operators into the same execution flow.
1.3 Heterogeneous Computing Demands of Multimodal Inputs
Supporting image inference means the inference engine must simultaneously handle two types of computation subgraphs: the Vision Encoder and the Language Decoder. The Vision Encoder relies primarily on convolutions or windowed attention, featuring relatively regular computation patterns. The Language Decoder generates tokens sequentially in an autoregressive manner, exhibiting stronger computational dynamism. How efficiently data exchange and pipeline parallelism are achieved between heterogeneous computation subgraphs directly determines the end-to-end latency of multimodal inference.
II. Overall Architecture of Day-0 Adaptation for Domestic AI Chips
2.1 Layered Decoupling Design of the Software Stack
The core prerequisite for achieving Day-0 adaptation is building a layered, decoupled software ecosystem. The overall architecture can be divided into four layers from bottom to top:
- Driver and Runtime Layer: Responsible for device management, memory allocation, and kernel launch, providing programming interfaces compatible with the CUDA runtime.
- Operator Library Layer: Provides deeply optimized basic computation units, including dense/sparse matrix operators, normalization operators, and activation functions.
- Compiler Layer: Supports automatic conversion from high-level languages (such as Python, Triton) to device instructions, incorporating auto-tuning and graph optimization capabilities.
- Framework Adaptation Layer: Receives calls from mainstream inference engines (such as SGLang), mapping the framework’s dataflow graphs to underlying operator execution sequences.
The key value of this layered architecture is that when a new model is released, developers only need to focus on incremental development at the framework adaptation and operator layers, without needing to re-adapt driver and runtime interfaces. This “base layer reuse + business layer customization” model is the fundamental guarantee for achieving Day-0 adaptation.
2.2 Pipeline Optimization for Model Loading and Service Startup
The main bottleneck during the model loading phase is weight file reading and format conversion. For terabyte-scale weights, directly adopting a full loading strategy can lead to service startup times spanning tens of minutes. A common industry practice is to use a sharded lazy loading mechanism:
- First, read the model configuration file and weight index to construct a Memory Mapping.
- Automatically plan weight sharding strategies based on the memory topology, mapping different expert networks to different compute cards.
- During service startup, load only the shared layer weights (such as Embedding and Attention projection layers); expert weights are loaded on demand.
Complementing this strategy, Day-0 adaptation introduces online weight format conversion technology. This converts MXFP4 format expert weights into the FP8/FP16 formats required for basic computation in real-time during the loading phase, rather than pre-generating a complete offline copy. This move not only saves disk I/O bandwidth but also prevents the impact of large temporary files on the storage system.
III. Adaptation and Optimization of Key Inference Paths
3.1 Integration Process of Sparse Attention Operators
The adaptation of Sparse Attention (DSA) is the most core technical component of this effort. The computation flow of this operator can be decomposed into three steps:
Step 1: Index Selection. Based on the attention score distribution of the current query, dynamically determine the key-value position indices to attend to. This process is similar to a Top-K selection operation in database queries and must ensure low memory access latency in high-concurrency scenarios.
Step 2: KV Cache Layout Rearrangement. Because different queries select different KV positions, a static KV cache layout generates severe memory fragmentation. During adaptation, a Block-based Cache strategy is employed. The KV cache is partitioned into fixed-size blocks (e.g., Block Size=16), and a block-level index table is maintained, enabling sparse access to complete block positioning in O(1) time.
Step 3: Execute Sparse Matrix Multiplication. After completing the index mapping, block matrix multiplication is executed via the sparse attention kernel in the MATE operator library. This kernel is deeply customized for the compute units of domestic accelerator cards, supporting a hybrid mode of intra-block dense computation and inter-block sparse skipping, effectively mitigating the bank conflicts common in sparse computation.
3.2 Workspace Reuse for Mixed-Precision Computation
Facing MXFP4 expert weights, a key engineering pain point is that repeated format conversion operations consume significant GPU compute power. The adaptation team adopted a workspace reuse strategy:
- Pre-allocate a fixed-size FP8 Work Area in memory.
- Convert each expert weight from FP4 to FP8 before its first scheduling, writing the result into the Work Area.
- If the same expert is called multiple subsequently (high-frequency experts in MoE scenarios are easily activated consecutively), directly reuse the FP8 copy in the Work Area to avoid repeated conversions.
The core benefit of this strategy is trading space for time: using limited memory overhead (one workspace is approximately 1/8 of the main model weights) to achieve tens of times speedup in format conversion.
3.3 Topology-Aware Scheduling for Multi-Card Communication
In Expert Parallelism scenarios, compute cards must frequently exchange token representations (All-to-All communication). Communication bottlenecks often become the upper limit for inference throughput. During adaptation, a topology-aware communication scheduling algorithm was utilized:
- During the model loading phase, obtain the inter-node communication bandwidth matrix via NCCL topology detection.
- Based on the popularity distribution of expert networks (derived from historical request statistics), prioritize placing high-frequency experts on compute cards within the same rack or switch domain.
- During inference runtime, employ a Double Buffering mechanism to hide communication latency. When the expert computation for the current layer is completed, immediately trigger communication transmission for the next layer, overlapping computation and communication in a pipeline.
This optimization increased the cluster’s effective compute utilization by approximately 15% to 20%.
IV. Engineering Practices for Day-0 Adaptation
4.1 Rapid Construction of Regression Testing Systems
Day-0 adaptation requires verifying the model’s functional correctness within a limited timeframe. In practice, a three-level testing node is typically constructed:
| Test Level | Verification Focus | Acceptance Criteria |
|---|---|---|
| Unit Test | Output consistency of single operator kernel | Error < 1e-3 compared to CPU reference implementation |
| Pipeline Test | Output quality of text generation pipeline | Perplexity (PPL) deviation < 5% compared to official reference |
| Integration Test | Stability of multi-card distributed inference | Continuous operation for 100 hours without abnormal interruption |
It is worth noting that for a 3-trillion-parameter model, running even dozens of test samples consumes considerable compute resources. Therefore, a “phased enablement” mechanism was introduced into the automated testing framework: first verify basic pipeline correctness through 32 standard samples, then gradually expand to the full test set.
4.2 Methodology for Performance Benchmark Evaluation
When evaluating the results of Day-0 adaptation, benchmarks commonly used in the industry must be established. Evaluation metrics should include:
- Time To First Token (TTFT): The time from sending a request to generating the first response token, reflecting system initialization and Prefill efficiency.
- Time Per Output Token (TPOT): The average generation time per token, reflecting the inference throughput capability during the Decode phase.
- Effective Throughput Rate: The ratio of the system’s actual completed token generation volume to the theoretical peak under concurrent request scenarios.
On top of these metrics, the adaptation team recommends adding a “long-context stability” indicator. Under an extreme input of 1 million tokens, the KV Cache overflow rate and inference accuracy degradation are monitored to verify the engineering robustness of the system.
V. Summary and Outlook
This article systematically outlined the technical landscape of Day-0 adaptation for domestic AI chips running 3-trillion-parameter MoE models. The core conclusions are as follows:
- Layered decoupling of the software stack is the cornerstone of rapid adaptation. The combination of universal interfaces at the driver layer and incremental development at the operator layer compresses the adaptation cycle for new models from months to days.
- Sparse attention and mixed-precision computation are the two major difficulties in adapting trillion-parameter models, requiring synergistic optimization of operator-level deep customization and memory planning.
- Topology-aware scheduling for multi-card communication plays an increasingly critical role as model scales expand, directly impacting the effective utilization rate of cluster compute power.
Looking ahead, as model parameter scales evolve toward the tens of trillions, we will see larger proportions of expert sparsity and more aggressive quantization schemes (such as 2-bit weights). This raises the bar for the software stacks of AI chips—they must not only support rapid functional adaptation but also possess automated performance tuning capabilities. Compiler-driven Auto Kernel Generation and machine learning-based memory planning strategies will be the core competitiveness of next-generation AI infrastructure. For enterprises looking to deploy ultra-large-scale models in private environments, building a sustainably evolving software stack ecosystem is of greater strategic value than merely pursuing peak per-card compute power.
