Adapting Trillion-Parameter MoE Models on Domestic AI Compute Cards: A Technical Guide

1. Background: Engineering Challenges of Ultra-Large-Scale Model Deployment

Recently, a Mixture of Experts (MoE) architecture model with a parameter scale reaching 3 trillion was open-sourced. It supports an ultra-long context window of 1 million tokens and has achieved significant capability improvements in code generation, 3D content creation, financial research, knowledge QA, and long-range Agent tasks. According to public benchmark data, the model scored 65.7% on terminal interactive benchmarks and 61.2% on long-range code task benchmarks, representing an increase of about 5 percentage points compared to the previous generation of the same series. Its overall performance is now approaching the level of top-tier closed-source models.

However, for infrastructure developers, improved model capabilities also mean an exponential increase in deployment complexity. Inferring a 3-trillion-parameter MoE model involves complex computational patterns such as expert routing, dynamic activation, and ultra-long KV Cache management. This places extremely high demands on the memory capacity, interconnect bandwidth, operator efficiency, and software stack maturity of the underlying hardware.

This article systematically explores the key technical paths and engineering practices for completing inference adaptation on domestic AI training and inference compute cards, focusing on the architectural characteristics and computational requirements of such ultra-large-scale MoE models.

2. Model Architecture Characteristics and Computational Requirements

2.1 Core Computational Characteristics of the MoE Architecture

Unlike traditional dense Transformer models, the MoE architecture introduces multiple parallel expert networks in each layer. A gating network dynamically selects the Top-K experts for each token to perform the computation. The core advantage of this design is that the total parameter count of the model can be vastly expanded, while the actual computation per token only relates to the number of activated experts. This balances parameter scale with inference cost.

For a 3-trillion-parameter MoE model, its computational characteristics can be summarized as follows:

• Sparse Activation: A single forward pass only activates a small fraction of the total parameters, but the load balancing strategy for expert distribution directly impacts hardware utilization.
• Dynamic Routing: The output of the gating network determines which experts a token is assigned to, posing irregular challenges for operator fusion and memory access patterns.
• Ultra-Long Context: A 1-million-token context window means the storage requirement for the KV Cache is extremely massive, requiring refined memory management strategies.
• Communication-Intensive: When the model spans multiple compute devices for Expert Parallelism, All-to-All communication becomes a critical bottleneck.

2.2 Memory and Bandwidth Estimation

Taking the 3-trillion-parameter MoE model as an example, if inference is performed in FP16/BF16 precision, the model weights alone require about 6TB of storage space. Even with INT8 quantization, the weight storage requirement remains around 3TB. This means single-card deployment is almost impossible, and multi-card or even multi-node distributed inference solutions must be utilized.

In actual deployment, the following resource constraints need to be comprehensively considered:

Total Memory Requirement = Model Weights + KV Cache + Activations + Communication Buffers

The size of the KV Cache scales linearly with the context length. For a 1-million-token context, even if optimization techniques like Grouped-Query Attention (GQA) are used, the KV Cache can still reach hundreds of gigabytes. Therefore, memory management strategies—including Paged Attention, KV Cache quantization, and prefix caching reuse—become key factors determining whether the system can run stably.

3. Key Technical Paths for Inference Adaptation

3.1 Operator Mapping and Optimization at the Software Stack Level

To adapt such models on domestic AI compute cards, the primary task is to leverage the self-developed software stack to map and optimize key operators. This includes, but is not limited to:

• Attention Operators: Flash Attention implementations supporting variable-length sequences, adapting to attention variants like GQA/MQA.
• MoE Gating and Routing Operators: Efficient Top-K selection, Softmax normalization, and expert dispatch logic.
• Matrix Multiplication Operators: GEMM optimization for different precisions (FP16/BF16/INT8), fully utilizing the hardware’s Tensor Core capabilities.
• Normalization and Activation Operators: Fused implementations of common operations like RMSNorm and SiLU.

During operator migration, special attention must be paid to numerical precision consistency. Since different hardware platforms have differences in the implementation details of floating-point operations, strict precision alignment tests must be conducted to ensure inference results remain consistent with the reference implementation.

3.2 Distributed Parallelism Strategy Design

For trillion-parameter models, a single parallelism strategy often cannot meet the requirements. It is usually necessary to combine multiple parallel methods:

Parallelism Method	Applicable Scenario	Key Considerations
Tensor Parallelism	Splitting matrix operations within a single layer	High communication frequency, suitable for intra-node high-bandwidth interconnects
Expert Parallelism	Distributing MoE experts to different devices	All-to-All communication, requires load balancing
Pipeline Parallelism	Splitting the model by layers	Reducing bubbles, requires fine-grained scheduling
Data Parallelism	Concurrent processing of multiple requests	Suitable for high-throughput scenarios

In actual deployment, a hybrid strategy of Tensor Parallelism + Expert Parallelism is recommended: splitting the weights of the attention layer across multiple cards within the same node via tensor parallelism, while distributing MoE experts across different devices to achieve expert parallelism. This combination strikes a good balance between communication overhead and memory footprint.

3.3 KV Cache Management and Long Context Optimization

A 1-million-token context window is one of the core selling points of such models, but it also places extremely high demands on memory management. Here are several key optimization methods:

Paged Attention: Manages the KV Cache in pages (blocks) to avoid fragmentation caused by contiguous memory allocation. Each page stores the Key and Value vectors of a fixed number of tokens, indexed via a page table.

KV Cache Quantization: Quantizing the KV Cache from FP16 to INT8 or even INT4 can significantly reduce memory usage. However, it is important to note the impact of quantization errors on attention scores. Typically, per-channel quantization is used for Keys, and per-token quantization is used for Values.

Prefix Caching Reuse: For multi-turn dialogues or Agent scenarios, the prefix portion of the system prompt and historical context can be shared across multiple requests, avoiding redundant computation.

# Schematic code: Basic structure of Paged KV Cache management
class PagedKVCache:
    def __init__(self, num_layers, num_heads, head_dim, page_size, max_pages):
        self.page_size = page_size
        self.page_table = {}  # Mapping from logical pages to physical pages
        self.free_pages = list(range(max_pages))
    def allocate(self, seq_id, num_tokens):
        """Allocate KV Cache pages for a sequence"""
        num_pages_needed = (num_tokens + self.page_size - 1) // self.page_size
        allocated = []
        for _ in range(num_pages_needed):
            if self.free_pages:
                page = self.free_pages.pop()
                allocated.append(page)
        self.page_table[seq_id] = allocated
        return allocated
    def free(self, seq_id):
        """Release pages occupied by a sequence"""
        if seq_id in self.page_table:
            self.free_pages.extend(self.page_table.pop(seq_id))

3.4 Inference Engine Scheduling and Batching

In actual serving scenarios, the inference engine needs to handle multiple concurrent requests simultaneously, and the input and output lengths of each request vary. Continuous Batching technology allows new requests to be dynamically added or completed requests to be removed during the decoding process, thereby maximizing hardware utilization.

For MoE models, the scheduling strategy also needs to consider expert load balancing: if tokens in a certain batch are intensively routed to a few experts, it will cause some devices to be idle while others are overloaded. Therefore, the inference engine needs to perform fine-grained, load-aware scheduling at the batching level.

4. Adaptation Verification and Performance Considerations

4.1 Functional Correctness Verification

After completing the adaptation, it is first necessary to verify the correctness of the model inference results. Common methods include:

• Layer-by-Layer Comparison: Numerically comparing the intermediate layer outputs of the adapted model with the reference implementation layer by layer to locate the source of precision deviations.
• End-to-End Evaluation: Running the complete inference process on standard benchmark datasets and comparing the scores with officially published data.
• Boundary Testing: Conducting specialized tests for boundary conditions such as ultra-long contexts and extreme input distributions.

4.2 Performance Tuning Directions

On the basis of ensuring correctness, performance tuning is typically carried out from the following dimensions:

1. Operator Fusion: Merging multiple consecutive small operators into a single large operator to reduce Kernel Launch overhead and memory access times.
2. Communication Optimization: Utilizing the high-speed interconnect capabilities provided by the hardware to optimize collective communication operations like All-to-All and All-Reduce.
3. Pipeline Scheduling: Improving overall throughput through separated scheduling of the Prefill and Decode phases.
4. Quantization Acceleration: Applying INT8/INT4 quantization within allowable precision limits to increase computational throughput and reduce memory pressure.

5. Summary and Outlook

Adapting trillion-parameter MoE models for inference is a systematic engineering effort involving hardware, software stacks, parallelism strategies, and memory management. Starting from the model architecture characteristics, this article systematically outlined the key technical paths for adapting such models on domestic AI compute cards, including operator mapping and optimization, distributed parallelism strategy design, KV Cache management, and inference engine scheduling.

As model parameter scales continue to grow and context windows expand, the challenges facing inference infrastructure will continue to escalate. In the future, the following directions are worth noting:

• More Efficient Sparse Computation: Hardware acceleration and algorithm co-design tailored to MoE routing characteristics.
• Smarter Memory Management: Dynamic KV Cache allocation and reclamation strategies combined with request characteristics.
• More Flexible Parallel Combinations: Automatically searching for optimal parallel schemes based on model structure and hardware topology.
• End-Cloud Collaborative Inference: Offloading part of the computational load to edge devices to reduce cloud pressure.

For developers, understanding these underlying technical principles not only helps in completing specific adaptation tasks but also enables quick identification of key issues and the formulation of effective engineering solutions when facing new generations of model architectures.