Day-0 Adaptation of a Domestic GPU for a 3-Trillion-Parameter MoE Model: A Full Technical Breakdown

1. Background: What System-Level Challenges Arise When Model Scale Reaches Trillions of Parameters

As large language model parameter counts approach the trillion scale, traditional dense Transformer architectures hit compute and bandwidth bottlenecks, making efficient inference within limited hardware resources increasingly difficult. The industry has broadly shifted toward sparse Mixture-of-Experts (MoE) architectures, which activate only a subset of expert sub-networks through a gating routing mechanism, controlling actual computation while maintaining model capacity. However, this architectural evolution introduces new system-level challenges:

• Memory bandwidth pressure: Expert weights are distributed across multiple devices or a unified memory pool. During inference, parameters must be frequently swapped, and bandwidth consumption grows exponentially.
• Sparse attention computation: To support ultra-long contexts (e.g., 1 million tokens), the model introduces Dynamic Sparse Attention (DSA). Index selection and irregular KV cache access patterns make memory access optimization difficult.
• Low-bit weight quantization: Expert weights often use the MXFP4 format (4-bit floating point), but existing hardware compute units lack native support for low-precision data types, requiring additional conversion overhead.
• Broken compilation pipeline: New operators introduced by the model architecture cannot be directly mapped to existing computation graph optimizers, requiring new operator integration at the compiler level.

These issues mean that simply swapping hardware or reusing an existing inference framework is insufficient. A full-stack “Day-0 adaptation” is required — completing the entire path from model definition to service deployment at the time of model release.

Scope note: This article focuses on inference adaptation for a 3-trillion-parameter MoE model on a domestic training-and-inference integrated accelerator card. Conclusions are based on specific hardware and model configurations. Other parameter scales or architecture variants require independent validation.

2. Overall Adaptation Architecture: Co-Design of the Training-and-Inference Accelerator Card and Unified Software Stack

This adaptation work is based on a domestic training-and-inference integrated intelligent compute card and a domestic parallel computing platform compatible with the CUDA-like ecosystem. The core design principles are as follows:

Application Layer  →  Kimi-K3.1 Pro (3T-parameter MoE, 1M context)
  ↓ Model definition dispatch
Framework Layer  →  Adaptation branch based on an open-source inference engine (supports async scheduling, continuous batching)
  ↓ Computation graph transfer
Operator Layer  →  General operator library + sparse attention专用 operators + low-bit compute runtime
  ↓ IR conversion and optimization
Compiler Layer  →  Unified IR → Performance optimization → Code generation (backend supports different architectures)
  ↓ Instruction dispatch
Hardware Layer  →  Training-and-inference integrated accelerator card (multi-precision compute, large memory, high-speed interconnect)

Key design decision: The team did not implement an inference engine from scratch. Instead, they adapted the backend of a mature inference framework with tensor parallelism, pipeline parallelism, and speculative decoding capabilities. This reduced ecosystem integration costs and allowed the team to focus on operator-level adaptation and performance tuning. At the computation graph level, a unified intermediate representation (IR) converts upstream model definitions into hardware-executable computation graphs: standard operators map directly to the general operator library, while new operators (such as the DSA index module and KDA state management) are implemented through a dedicated compiler with hand-crafted customization.

3. Restructuring DSA Sparse Attention: From Index Selection to Operator Integration

3.1 Sparse Attention Computation Path Analysis

Kimi-K3.1 Pro uses Dynamic Sparse Attention (DSA). Before computing attention scores at a given position, a lightweight index selection module dynamically determines which context windows to attend to. This results in: index generation dynamism that makes KV cache access patterns unpredictable; conventional block-sparse optimization strategies that are difficult to reuse; and vastly different memory access characteristics between Prefill and Decode phases.

3.2 Core Strategies of the Adaptation Solution

1. Joint optimization of index selection and KV cache layout: Index computation is split into two stages — “token importance scoring + top-k filtering.” KV cache is stored in blocks, with contiguous addresses within blocks and index pointer connections between blocks.
2. Prefill/Decode dual-path separation: During Prefill, sparse indices are converted into masked matrix multiplication, leveraging dense compute capability. During Decode, truly sparse operators are activated, skipping invalid KV blocks via bitmap indices.
3. Integration of a sparse attention operator library: Kernels optimized for different sparsity levels are pre-implemented, with explicit shared memory management. Index tables are preloaded into on-chip memory and copied asynchronously.
4. KDA state management and cross-layer residual fusion: KDA state information is stored in a dedicated buffer. Cross-layer residual connections directly propagate state updates, avoiding multiple global memory reads and writes.
5. Collaborative adaptation of the vision attention branch: The vision encoder uses a local window + global token pattern, embedded into the same pipeline as text, distinguished only by different indexing strategies.

Comparison with other sparse approaches:

Approach	Indexing Method	Applicable Phase	Memory Efficiency
DSA (this article)	Dynamic lightweight scoring + top-k	Prefill/Decode dual-mode	High (~30% of dense)
Block Sparse	Fixed block mask	Mainly Prefill	Medium
Top-k Routing	Global fixed-k selection	Mainly Decode	Lower

3.3 Performance Optimization and Validation

On a domestic training-and-inference integrated accelerator card (80 GB memory per card), tested across sequence lengths from 4K to 128K tokens, KV cache memory usage was reduced to approximately 30% of the dense attention solution (baseline: standard global attention at equivalent precision). In a 128K-context test sequence, the sparse attention operator achieved 2–3× the effective throughput of the dense attention operator on the same hardware, with average per-token generation latency variation during Decode under 5%, ensuring stable streaming output.

4. Loading and Computing MXFP4 Expert Weights: Online Conversion and Workspace Reuse

4.1 Problem Definition

MoE expert weights are stored in MXFP4 format (4 bits per parameter + shared scaling factor). Sending them directly to FP16/BF16 matrix multiplication units requires conversion to higher precision first, incurring significant overhead. Moreover, the set of activated experts differs per token, making pre-batch conversion impossible.

4.2 Online Conversion Engine Design

A hybrid strategy of online conversion + cross-layer reuse is adopted:

1. FP32 scaling factor extraction: MXFP4 weights are decomposed into 4-bit mantissas and a shared FP32 scaling factor, which is broadcast to cache during loading.
2. Vectorized conversion instructions: Vector load instructions pack 4-bit data into registers. Lookup-table mapping converts to FP16/FP32, with SIMD parallelization where each thread processes 8 weight values.
3. Cross-layer FP8 workspace reuse: A global Expert Least Recently Converted (LRC) cache tracks the most recent FP8 version and timestamp of each expert.

class ExpertWeightCache:
    def __init__(self, capacity):
        self.capacity = capacity  # Number of FP8 experts that can be cached
        self.cache = {}  # expert_id -> (fp8_tensor, timestamp, ref_count)
    def load(self, expert_id, mxpf4_loc):
        if expert_id in self.cache:
            self.cache[expert_id].ref_count += 1
            return self.cache[expert_id].fp8_tensor
        else:
            fp8_tensor = self.online_convert(mxpf4_loc)
            self.evict()  
            self.cache[expert_id] = (fp8_tensor, now, 1)
            return fp8_tensor

⚠️ Note: LRC cache capacity must be set reasonably based on total memory. Exceeding the limit will cause out-of-memory (OOM). It is recommended to reserve at least 20% safety margin based on available memory, and enable a joint eviction policy based on reference count and timestamp when the cache is full.

1. Overlapping computation and conversion pipeline: When memory bandwidth allows, expert weight conversion for the current batch runs concurrently with computation for the previous batch, using a multi-stage pipeline to hide conversion latency.

4.3 Measured Results

On a single inference node (batch size 32, continuous inference for 30 minutes), compared with a no-cache baseline (v1.0.0-base), expert weight loading memory bandwidth utilization increased from approximately 55% to approximately 85%, and end-to-end throughput improved by approximately 40%. These figures are based on internal benchmarks. Actual performance varies with model configuration and hardware environment. Readers are encouraged to reproduce and validate in their own environments.

5. Compiler Stack Migration and Integration: MATE Operator Library and Domestic Compilation Toolchain

5.1 Multi-Level Compiler Collaboration

• High-level pattern matching: Based on the graph optimization module of a general deep learning compiler, fusion patterns such as “Attention followed by linear layer” are automatically identified, reducing kernel launch counts.
• Operator generators: For custom operators (sparse index lookup, low-bit weight scaling), developers can write in a Triton-like sub-language corresponding to the domestic platform. The compiler automatically generates optimized kernels for different input shapes.
• Low-level instruction scheduling: Instruction-level compilers perform fine-grained tuning of register allocation, memory access instruction reordering, and loop unrolling to match hardware microarchitecture characteristics.

5.2 Engineering the Adaptation Process

1. Run the model’s native code on a GPU to collect operator call parameters and tensor shape distributions.
2. Identify unsupported operators. First verify correctness with a Python reference implementation, then rewrite in the Triton-like language. Numerical comparison tolerance is less than 1e-2.
3. Register new operators into the inference engine’s custom operator library. Establish a unified registry to avoid naming conflicts.
4. Validate model output consistency through end-to-end test sets (including logits differences under greedy decoding and sampling decoding).

5.3 Computation Path Pruning for New Architectures

• Operator fusion: Fuse LayerNorm with Quantize to eliminate two rounds of global memory round-trips between kernels.
• Memory planning: Merge tensors with non-overlapping lifetimes into the same physical address, reducing peak memory usage.
• Automatic vectorization: Automatically generate vectorized Load/Store instructions for low-dimensional hot data (sizes 2/4/8), using L2 cache-friendly layouts.

6. Key Risks and Alternative Solutions

During compiler stack adaptation, the following risks may be encountered:

Risk Scenario	Alternative Solution
New operator cannot be automatically mapped through IR	Fall back to Python reference implementation, invoked via FFI, trading performance for compatibility
Operator registration conflict causing crashes	Establish versioned namespaces, isolating operator libraries by model version
Insufficient memory causing OOM	Enable gradient checkpointing strategy, load expert weights in batches

7. Technical Takeaways and Methodology Insights from Day-0 Adaptation

This adaptation effort produced a “replicable adaptation process” for ultra-large-scale models. Core elements include:

1. Modular decoupling: Divided into five independent stages — “model structure analysis → compiler stack extension → operator optimization → pipeline scheduling → end-to-end tuning.” Each stage has defined inputs/outputs and regression test sets, supporting parallel progress.
2. Canary release strategy: Start with smoke tests on a subset of layers or small batches, then gradually expand validation scope to keep risk controllable at each step.
3. Performance observability: Embed multi-layer performance monitoring in the inference engine (per-operator latency, memory bandwidth, cache hit rate) to locate bottlenecks rather than relying solely on empirical inference.

For developers facing similar requirements, it is recommended to first inventory new operators and assess integration costs. For an entirely new model architecture, do not write an inference engine from scratch. Instead, choose a mature framework for secondary development and focus energy on operator-level optimization.

Special note: Figures such as “2–3×,” “approximately 40%,” and “approximately 30%” in this article come from internal benchmarks under specific test environments and represent validated conclusions. Statements such as “recommended to first inventory” are empirical suggestions. Readers should independently verify against their own scenarios.

8. Future Outlook

8.1 Hardware Instruction Set Nativization

Build support for common low-bit formats (MXFP4/INT4) into hardware instruction sets, eliminating software-layer conversion overhead and fundamentally improving low-precision compute efficiency.

8.2 Sparse Compute Unit Hardening

Solidify sparse attention mechanisms into hardware-configurable compute units, enabling dynamic sparse patterns like DSA to receive native acceleration at the hardware layer and reducing operator development and tuning costs.

8.3 Automated Compilation Templates

Through automated compilers, translate more complex model structures (such as state space models and hybrid expert attention variants) into general optimization templates, achieving the leap from “manual adaptation” to “platform anticipation.” On this path, the co-design of software stacks and accelerated hardware will be a key variable determining the productivity of AI infrastructure.