How a Domestic GPU Achieved Day-0 Adaptation for a 3-Trillion-Parameter MoE Model

Background and Challenges: The Deployment Gap for Very Large MoE Models

Sparse mixture-of-experts (MoE) architectures have become the mainstream choice for very large language models. Their core idea is to dynamically activate a subset of expert sub-networks through a gating mechanism, preserving model capacity while significantly reducing the computation required for each forward pass. However, when model scale reaches the trillion-parameter level, engineering deployment is no longer a simple performance-tuning problem. It requires reworking the entire path from the software stack to hardware mapping.

A three-dimensional trillion-parameter model introduces several constraints:

• Memory capacity bottlenecks: Even with low-bit quantization, the total weight footprint far exceeds the memory of a single accelerator. Multi-device coordination and fine-grained memory management become mandatory.
• Sparse computation efficiency: The expert-routing mechanism in MoE models creates highly divergent computation patterns. General-purpose matrix multiplication libraries struggle to reach hardware peak performance under these conditions.
• Attention mechanism variations: In long-context scenarios, standard dense attention is replaced by sparse attention variants such as DSA. Their index computation and KV cache layout differ fundamentally from traditional implementations.
• Precision and throughput trade-offs: Low-precision training or inference weights such as MXFP4 reduce memory usage, but they require runtime format conversion that can introduce additional latency.

A domestic GPU vendor recently announced that it completed Day-0 adaptation for Kimi-K3.1 Pro, a 3-trillion-parameter model, using its training-inference accelerator and full software stack. Day-0 here means the inference path was ready on the same day the model weights were publicly released. This reflects not only engineering efficiency but also early and deep involvement in the model architecture design phase, along with reusable adaptation methodology.

This article reconstructs the core stages of that adaptation process and discusses the general engineering patterns behind it.

1. Overall Adaptation Approach: A Compiler Toolchain-Centric Strategy

For a newly released model with significant architectural differences from mainstream open-source models, there are typically two adaptation strategies:

1. Post-hoc adaptation: Wait for the model weights to be released, then incrementally develop compatibility on top of an existing inference engine. This approach has long cycles, high risk, and no real-time feedback from the model designers.
2. Pre-release collaborative adaptation: Before the weights are public but after the architecture is finalized, begin toolchain migration based on model documentation, source configuration, and pre-release versions. This requires an inference framework and compiler with high flexibility and extensibility to quickly absorb new operators and data flows.

This adaptation clearly followed the second approach. The technical route centered on smoothly mapping the upstream model implementation to the underlying hardware execution graph. Specifically, the team used SGLang-MUSA, a hardware-adapted branch of SGLang, to host the upstream model code. The frontend model definition, such as the Python layer, was preserved, but all tensor computation, memory allocation, and operator scheduling had to be rebuilt through the MUSA toolchain.

Four key components played decisive roles:

• MATE operator library: Provides high-performance fused operators, especially general matrix multiplication, grouped matrix multiplication, and activation functions for MoE architectures. It enables seamless replacement through a unified operator abstraction interface on the model side.
• Triton-MUSA and TileLang-MUSA: Two domain-specific compiler backends that allow developers to write custom operators in Python-like syntax and compile them efficiently into native GPU instructions. This capability was particularly important for highly customized operators such as DSA attention.
• MTCC compiler: Responsible for progressively lowering high-level intermediate representations into hardware-executable code while applying loop unrolling, shared memory reuse, and other optimizations.

Through this combination, the critical computation paths in the model, including embedding, attention, expert feed-forward networks, and gated routing, were successfully migrated to the target hardware. Notably, the model logic was not modified during the process, which kept the coupling between the model-side implementation and the underlying hardware to a minimum.

2. Deep Adaptation of Sparse Attention: Redesigning Memory Layout and Computation Paths

Kimi-K3.1 Pro natively uses DSA (Dense-Sparse Attention), a core innovation that supports its million-token context window. The central idea is that not all historical tokens contribute meaningfully to the attention computation of the current token. By combining retrieval and prediction, a small number of highly relevant tokens are selected in advance, reducing attention complexity from (O(n^2)) to (O(n\sqrt{n})) or lower.

Turning this algorithm into an efficient GPU implementation requires reworking three layers:

1. Implementing the index selection mechanism in advance
DSA first generates a candidate index set, where each query computes attention only against a pre-selected set of keys and values. The index generation process is not a simple sort. It involves approximate retrieval based on vector similarity. During adaptation, the team implemented a lowered version of the index selection operator in the MUSA software stack. The operator uses tensor cores to accelerate the dot products between query vectors and all key vectors, then collects indices through Top-K selection. This stage has relatively low compute density but highly irregular memory access patterns, so it requires an efficient parallel reduction strategy to avoid thread divergence.

2. Dynamic planning of the KV cache layout
When each token uses only a small number of KV entries, a traditional block-based cache layout wastes substantial bandwidth. The team introduced an on-demand two-dimensional sparse cache combined with a KDA (Key-Data Address) state management mechanism that dynamically tracks the mapping between logical and physical KV addresses across heads. The cache therefore no longer behaves like a simple ring buffer. It functions more like an object storage system with random access, where hit rate and defragmentation strategy directly affect P99 latency. The vector states corresponding to cross-layer residual connections are also brought under unified management to ensure gradient and activation correctness across long computation chains.

3. Separate heterogeneous paths for prefill and decode
Prefill and decode have very different computation characteristics. Supporting a 1.1-million-token context means the prefill phase must process extremely long sequences, where the bottleneck is tensor-core matrix throughput. Decode, by contrast, is limited by memory bandwidth and operator launch latency. The team created separate compute kernels in the MUSA stack for these two phases. For sparse attention, prefill uses block-wise matrix multiplication combined with sparse masks, while decode uses one-dimensional parallel scanning and partial result reuse to avoid CPU-side synchronization overhead.

3. Low-Precision Expert Weight Conversion: An Online Computation Strategy Without Full Materialization

The model stores expert weights in MXFP4 format to save memory, but performing matrix multiplication directly in 4-bit precision would degrade accuracy too much. Moreover, current GPU tensor cores do not universally support FP4. The engineering team therefore needed to convert weights to FP8 or FP16 before computation.

The simplest approach would be offline conversion: convert all expert weights from MXFP4 to FP8, store them in memory, and then run inference. However, this creates two complete copies of the weights. For a 3-trillion-parameter model, even FP8 would require nearly 30 TB of memory, far beyond the physical capacity of any single server. More importantly, offline conversion means the weight data must support full dequantization and requantization workflows, consuming significant PCIe or network bandwidth and delaying service startup.

This adaptation instead used online conversion with cross-layer workspace reuse:

• Online conversion unit: A lightweight conversion module is embedded inside the expert computation core. While expert weights are loaded into shared memory, the module converts them from MXFP4 to FP8. The exponent offsets and scaling factor tables required for conversion are kept in register files to avoid external lookups.
• Cross-layer workspace reuse: Within a given MoE layer, 4-bit expert weights are converted to FP8 in real time. Because different experts across layers share almost no weights, additional FP8 workspace is needed. The team designed a global FP8 workspace pool that is dynamically allocated across different computation streams but prevents temporal conflicts, significantly reducing peak memory demand.
• Precision protection: Rounding errors introduced by online conversion are compensated through stochastic rounding and partial residual accumulation, keeping final output quality essentially consistent with offline full-FP8 computation.

The key benefit of this design is that model loading time is compressed to nearly direct memory mapping of the original weights, without requiring a full format conversion pass. At the same time, peak memory usage is constrained to an acceptable range, making single-node deployment feasible.

4. Ecosystem Coordination and Performance Validation

Beyond the core computation paths, a complete Day-0 delivery also requires a working model serving entry point, a staged rollout path, and a performance monitoring system. At the serving layer, the SGLang scheduler was extended to be aware of MoE routing sparsity, allowing it to prioritize batch requests that improve data locality across experts and reduce inter-device communication.

For performance validation, the team completed functional testing of the basic forward path and also ran benchmarks focused on agentic tasks. Kimi-K3.1 Pro ranks first among open-weight models on major industry leaderboards, with particularly strong results on long-horizon tasks such as terminal command execution and complex code generation. These capabilities are the result of coordinated optimization between the DSA attention mechanism and the MoE structure. However, it is worth noting that even if the model itself performs well, an inference engine without a matching KV cache strategy will undermine its long-context advantages in production serving. The KDA state management mechanism in this adaptation is a key element in ensuring effective utilization of the context window.

5. Summary and Implications: A Replicable Path for Cloud-Native AI Inference

The success of this Day-0 adaptation is superficially a technical integration between a single vendor and a single model. In practice, it reflects a broader engineering pattern:

1. Operators as interfaces: Logically, supporting any new model architecture should be treated as a process of registering a new combination of operators. A compiler-based workflow and a high-quality operator library can significantly reduce dependence on hand-written operator-level optimization.
2. Memory orchestration matters more than raw compute: In the era of tens or hundreds of billions of parameters, memory access patterns, data layout, and lifecycle management often determine system throughput and latency more than simply increasing FLOPS.
3. Front-loading software-hardware co-design: Day-0 adaptation implies that AI infrastructure teams need a pairing-style collaboration with model teams, getting involved in architecture review early to identify performance traps hidden in Python-level code.

For developers and architects building AI infrastructure, the main takeaway is that an inference stack with a model-operator abstraction layer, a unified memory manager, and an extensible compiler backend delivers far greater long-term return on investment than repeatedly hard-coding optimizations for individual models. This is also the technical foundation required for cloud providers to offer general-purpose LLM inference services at scale.