Running a 3-Trillion-Parameter LLM on a Chinese AI Chip: A Day-0 Full-Stack Adaptation

Background: Why 3-Trillion-Parameter Models Are a Watershed

In recent years, the parameter scale of large language models has grown far faster than Moore’s law. Each jump—from hundreds of billions to trillions, and now to three trillion parameters—places extremely demanding requirements on the underlying compute platform. At three trillion parameters, the model weights alone require several terabytes of storage; even with mainstream high-bandwidth memory, no single card can hold the complete model. More critically, during inference, every Transformer layer’s matrix multiplications, attention computations, and activation caching generate intermediate data flows that far exceed the static weight footprint.

Under these conditions, model serving is no longer as simple as “load the weights and run a forward pass.” It involves combining complex strategies such as distributed tensor parallelism, pipeline parallelism, and sequence parallelism. Beyond raw compute, the chip must provide sufficient inter-chip interconnect bandwidth, a flexible memory management mechanism, and a software stack that can efficiently express parallel computation graphs and automatically handle operator splitting and scheduling. This is precisely the core value demonstrated by this adaptation case: achieving “Day-0” support for a frontier large model on a Chinese AI chip and its proprietary compute architecture—meaning usable compute support is available on the same day the model is released.

Model Characteristics: The Dual Challenge of Compute Intensity and Memory Intensity

The 3-trillion-parameter model features a million-token context window, which means the key-value cache (KV Cache) it must handle during inference is enormous. For long sequences, attention has quadratic time complexity in sequence length, while cache space grows linearly. A million-token context can push KV Cache usage beyond hundreds of gigabytes—fundamentally changing traditional memory planning logic. At inference time, you must not only accommodate the weights but also reserve large contiguous, quickly addressable cache regions.

In addition, the model’s reinforcement learning post-training and long-horizon agent task optimization give its inference patterns two distinct characteristics. The first is the “long chain of thought”: before producing a final answer, the model performs multi-step reasoning, and each step’s intermediate results are kept in memory for later backtracking and correction, placing sustained pressure on both memory bandwidth and capacity. The second is “high-concurrency branching”: in scenarios such as code generation and agent task planning, the model may need to explore multiple candidate paths simultaneously, creating a beam-search-like computation pattern in which multiple inference requests execute interleaved on the chip—raising the bar for the scheduler’s multi-stream concurrency capability.

Adaptation, therefore, is far from simply translating model weights into chip instructions. It requires a deep understanding of the model architecture—the routing mechanism of the mixture-of-experts layers, long-context attention computation, and memory usage hotspots—and systematic adjustments to the chip’s compute kernels, compilation optimization strategies, and runtime resource management.

The Critical Path of Software-Hardware Co-Adaptation

To address these challenges, the adaptation work followed a four-layer structure—“front-end graph optimization, mid-end operator mapping, back-end kernel tuning, and distributed strategy orchestration”—with each layer closely built on the chip’s own compute architecture software stack.

Graph-Level Optimization: Semantic Rewriting from Computation Graph to Execution Plan

After obtaining the model’s computation graph description, the first task is not direct translation but semantic-level graph optimization. The compiler needs to identify fusable operator sequences. For example, in multi-head attention, the linear projections for queries (Q), keys (K), and values (V) can be fused with the subsequent tensor splitting and rotary position embeddings (RoPE), reducing the number of round trips tensors make to and from memory.

Graph optimization also requires special handling for long sequences. For million-token contexts, the core idea of FlashAttention-style algorithms is to perform blocked computation in SRAM (on-chip static random access memory), avoiding writing the full attention matrix back to memory. This demands that the compiler stack recognize standard attention patterns and automatically replace them with the chip’s efficient blocked execution kernels. In this adaptation, the compiler successfully applied such replacements on critical paths, significantly reducing peak memory usage.

Instruction-Level Tuning of Operator Kernels

Graph optimization determines the computation logic; operator kernels determine how efficiently the hardware executes it. The chip’s instruction set architecture (ISA) supports specific vectorized operations and matrix multiply-accumulate instructions. The adaptation team rewrote, at the instruction level, each of the core Transformer operators: large matrix multiplication, Softmax, LayerNorm, and activation functions such as Swish.

Take matrix multiplication as an example. A 3-trillion-parameter model contains many matrix operations with irregular shapes. Standard general matrix multiplication (GEMM) libraries are typically hyper-optimized for specific shapes (such as square or wide rectangular matrices) but show low utilization on narrow, elongated shapes. During adaptation, the kernel team customized dozens of shape-specific matrix multiplication kernels based on statistics of the actual runtime tensor shapes of each layer, and through instruction-level pipeline scheduling raised compute unit utilization by roughly 15%–20% compared with the generic libraries. For the Softmax operator, the team leveraged the chip’s high-speed exponential operation instructions together with on-chip memory reduction operations, eliminating memory round trips for intermediate results.

Tiered Memory Management and an Elastic KV Cache Strategy

During 3-trillion-parameter inference, memory is the scarce resource that determines success or failure. The adaptation practice adopted a tiered memory management system:

• Weight-resident tier: Model weights reside in memory in compressed or sharded form, with fast random access provided by the high-bandwidth memory controller.
• Activation staging tier: For intermediate activations produced by the forward pass, a cost-model-based automatic eviction policy dynamically swaps infrequently used activation data out to host memory or high-speed storage, reloading it only when needed.
• Elastic KV Cache tier: For the million-token KV Cache, on-demand allocation and pooled reuse are applied. By managing the KV Cache in blocks and leveraging the chip’s memory management unit for efficient virtual-to-physical address mapping, multiple inference requests can safely reuse the same physical memory regions, greatly improving memory utilization.

Deep Integration of Distributed Parallelism Strategies

A 3-trillion-parameter model necessarily relies on multi-card, multi-node cluster inference. Building on the chip’s high-speed inter-chip interconnect protocol, the adaptation implemented flexible combinations of tensor parallelism, pipeline parallelism, and data parallelism.

Along the tensor parallelism dimension, each Transformer layer’s parameter matrices are split across multiple cards, and partial sums are synchronized through All-Reduce operations during computation. To reduce communication overhead, the compiler stack overlapped the scheduling of communication and compute operators—during the register write-back phase of the second half of a matrix multiplication, prefetch transfers for the data needed by the next computation are already initiated, achieving pipelined overlap between computation and communication.

For pipeline parallelism, given that the model’s layer count does not divide evenly across stages, the adaptation team used a dynamic programming-based layer splitting algorithm to find the optimal layer allocation ratio, minimizing the variance of compute and communication times across compute nodes. In the final inference solution, multi-card collaboration efficiency reached mainstream industry levels, and latency remained stable under high-concurrency, multi-request scenarios.

Measured Results and Stability Assurance

After completing the Day-0 adaptation, the team evaluated the model’s performance on the chip across multiple dimensions. The evaluation sets covered code generation (involving long-range dependency reasoning over multi-file syntax trees), mathematical and logical reasoning (multi-step derivations), and long-text summarization (requiring accurate capture of details within million-word-level contexts). All achieved accuracy at industry-usable levels, and inference latency met real-time interaction expectations.

Notably, in long-horizon agent scenarios, the model must continuously execute dozens of “think–call tool–observe result–think again” loops. Through an enhanced context-switching mechanism, the chip achieves microsecond-level rapid switching between multiple inference branches with no significant memory fragmentation, ensuring stability for long-running agents. Meanwhile, the adapted inference service ran a continuous 72-hour high-load stress test without memory leaks, compute unit overheating, or thermal throttling, confirming the engineering maturity of this software-hardware co-design solution.

Takeaways: A General Methodology for Deploying Trillion-Scale Models

From this adaptation practice, several general lessons can be distilled for landing ultra-large models on Chinese compute platforms:

1. Early involvement: Compute platforms should establish a “model architecture preview” mechanism, obtaining the computation graph description at the architecture finalization stage so that compiler stack and kernel library development can proceed in parallel with model training, rather than starting adaptation only after the model is released.
2. Deep exploitation of quantization and sparsity: Although this model did not explicitly publish quantization details, inference practice can explore dynamic weight compression techniques. For mixture-of-experts architectures, function-level activation sparsity can be perceived by the scheduler, so that only the data blocks of routed experts’ weights are loaded—this requires the runtime system to offer fine-grained, data-aware scheduling capabilities.
3. The value of an ecosystem compatibility layer: The efficiency of this adaptation owed much to compatibility with the high-level APIs of mainstream deep learning frameworks (such as general-purpose distributed training frameworks). This allows model development teams to leverage the chip infrastructure without rewriting core business logic.

Conclusion

The Day-0 adaptation of a 3-trillion-parameter large model marks a key threshold for Chinese AI chips in supporting ultra-large models. It represents not merely an increase in hardware compute figures, but a concentrated demonstration of full-stack system software strength: the compiler stack, runtime scheduling, memory management, and distributed coordination. As model scale continues to balloon and context windows stretch ever longer, deep chip-model co-design will become the central theme of AI infrastructure evolution. The AI compute platforms of the future will inevitably be unified entities in which compute architecture and model architecture interpenetrate and define each other—and this practice provides a highly valuable pioneering sample for that direction.