Day-0 Adaptation of Ultra-Large MoE Models on Domestic GPUs: Engineering Practice and Implementation Path

1. Background and Challenges: Why Day-0 Adaptation Is Difficult

As large language models exceed trillion-parameter scale, inference deployment must address not only growing compute demand but also coordinated software stack maturity. For AI accelerator vendors, rapid adaptation to newly released models — commonly called Day-0 adaptation — means that as soon as model weights are published, the hardware and software toolchain can immediately support inference workloads. This capability directly affects product competitiveness and ecosystem readiness.

Adapting ultra-large Mixture-of-Experts (MoE) models is particularly difficult. Kimi-K3.1 Pro, a recently released example, has roughly 3 trillion parameters, uses a typical MoE architecture, and introduces sparse attention mechanisms along with MXFP4 expert weight compression. These characteristics create three major adaptation challenges:

• Architectural diversity: MoE models combine expert networks, routing mechanisms, and shared layers within a highly complex computational graph. This places heavy demands on static graph capture and dynamic scheduling in the serving framework.
• Irregular computation patterns: Sparse attention and sparse expert activation produce many dynamic branches at runtime. Data dependencies are difficult to predict, so conventional dense operator optimization templates cannot be directly applied.
• Precision-performance trade-offs: MXFP4 (4-bit floating-point) expert weights significantly reduce memory consumption, but online dequantization and numerical precision control require careful engineering. Without proper tuning, inference outputs may show unacceptable deviation.

2. Overall Adaptation Design: Framework Integration and Compiler Bridging

A clear adaptation roadmap is essential. This Day-0 effort followed a four-layer structure: upstream model integration, critical computation path decomposition, low-level operator library mapping, and runtime service stabilization.

Layer 1: Framework integration. SGLang, a widely used inference framework, served as the primary execution environment. The team used its MUSA backend — a GPU acceleration backend based on the vendor’s unified system architecture — to load and execute the upstream model. This is the first step in Day-0 adaptation: ensuring that the model definition and weight loading can be correctly parsed and executed on a domestic accelerator. Because upstream models are typically written for the CUDA ecosystem, this layer requires automated API mapping and memory management semantic conversion.

Layer 2: Compiler and operator library support. Rather than manually porting each operator for the new architecture, the team built a multi-level operator generation system. The MATE operator library provides high-performance precompiled kernels. Triton-MUSA and TileLang-MUSA address dynamic-shape and static-shape scenarios respectively, allowing developers to write custom operators in Python-like DSLs. The MTCC compiler then optimizes and generates executable code for the target accelerator. This combination of precompiled and just-in-time compilation significantly reduces manual coding effort when adapting new operators.

Layer 3: Runtime and state management. Day-0 timelines are extremely tight, and directly reusing community model code often leads to performance bottlenecks. The team therefore used the MUSA software stack to deeply customize the core execution path of Kimi-K3.1 Pro, going beyond a shallow “it runs” level of adaptation.

3. Deep Adaptation of Sparse Attention

Sparse attention is a key technique for ultra-long-context models. It dynamically selects and computes attention over a small subset of token pairs, reducing complexity from O(n²) to near-linear. However, this irregular computation pattern is unfriendly to both hardware and software stacks.

The adaptation work addressed three aspects of the sparse attention mechanism:

1. Index selection and KV cache layout restructuring: Sparse attention requires independent index selection at each layer, and the KV cache must support dynamic, non-contiguous memory addressing. The team customized the cache management module in the MUSA software stack to enable flexible block partitioning and reordering of KV tensors, effectively hiding the memory access latency of index selection.

2. Separate prefill and decode paths: The prefill phase processes long sequences and benefits from compute-intensive batch execution. The decode phase generates tokens one at a time and is memory-bound. These two paths have very different operator requirements. The team designed separate sparse attention kernels for each path and scheduled them through dedicated entry points in the MATE operator library.

3. Joint optimization with KDA state management and cross-layer residuals: Kimi-K3.1 Pro introduces KDA (Key-Value Delta Attention) for cross-layer attention state accumulation. Combined with the vision attention branch for multimodal input, the complete execution path contains multiple attention variants. The team orchestrated these components into a unified dataflow graph to avoid synchronization overhead between heterogeneous attention operations.

4. Online Processing Strategy for MXFP4 Expert Weights

The model stores expert weights in MXFP4 format. This substantially reduces storage and bandwidth pressure but creates two problems: current accelerator compute pipelines generally do not natively support 4-bit floating-point arithmetic, and offline conversion to full FP8 (8-bit floating-point) weights would significantly increase storage overhead and model loading time.

To balance precision, loading speed, and inference performance, the team designed an online conversion strategy with cross-layer reuse:

• Online conversion instead of offline preprocessing: During model loading, MXFP4 weights are not uniformly converted to FP8. Instead, each expert’s weight block is dynamically dequantized from MXFP4 to FP8 immediately before that expert is computed. This avoids the additional storage and time cost of full conversion.
• Cross-layer FP8 workspace reuse: MoE models often share some experts across adjacent layers. The team maintained a global FP8 workspace. If an expert’s weights already exist in the workspace and its input computation has not yet completed, subsequent layers that call the same expert can reuse the existing FP8 values and skip repeated dequantization. This strategy reduces the additional overhead of online conversion to near zero.

5. Performance Validation and Evaluation Benchmarks

After completing the computation path adaptation, the team validated the core inference pipeline across three dimensions:

• Correctness: Using text generation and image description as baseline tasks, the team compared semantic consistency of output token sequences and sampling distribution similarity to ensure numerical error remained within an acceptable range.
• Stability: Extended continuous inference stress tests monitored whether memory usage kept increasing or whether the service experienced hangs or out-of-memory errors. These tests verified the robustness of KDA state management and KV cache reclamation in long-session scenarios.
• Performance: The team focused on time-to-first-token (TTFT) and inter-token latency (ITL). Under sparse attention and sparse expert activation, these metrics are highly sensitive to operator scheduling efficiency and memory bandwidth utilization.

6. Engineering Lessons and Industry Implications

This Day-0 adaptation validated the ability of domestic accelerators to support ultra-large MoE models and produced several reusable engineering lessons:

• A “sandwich” strategy for hardware-software collaboration: Between the low-level vendor operator library (MATE) and the top-level model code, there must be a flexible intermediate representation layer such as Triton-MUSA or TileLang-MUSA. This layer enables rapid experimentation with custom kernels tailored to new model characteristics.
• Prior analysis before computation path decomposition: When facing a completely new model architecture, blindly attempting to run it is counterproductive. Static code analysis and profiling tools should first identify computation hotspots and memory access bottlenecks before designing a targeted adaptation plan.
• Careful runtime state management: The hidden killer of Day-0 adaptation is often state management, including KV cache fragmentation and lifecycle maintenance of cross-layer shared state. This area accounted for nearly 40% of the total adaptation effort and requires early planning.

7. Outlook

As MoE architectures evolve toward greater depth and width, and as models integrate multimodal and strong reasoning capabilities, Day-0 adaptation will become a core measure of software stack maturity for AI accelerators. Future adaptation work will move further toward automated compilation, end-to-end graph optimization, and dynamic shape scheduling — pushing large model inference from “runnable” toward high-performance, low-cost, production-scale deployment.