Adapting a Domestic GPU Stack to Ultra-Large-Scale MoE Models

1. Background: Engineering Challenges of a 3-Trillion-Parameter Model

Recently, an open-source community released the first model in its K3.1 series, with a parameter scale reaching the 3-trillion level and support for a 1-million-token context window. A model of this size delivers a significant capability leap, but it also places unprecedented demands on the underlying compute infrastructure.

At the same time, a domestic full-featured GPU vendor completed Day-0 adaptation for this model using its AI training-and-inference accelerator card and in-house software stack. The work covered model loading, service startup, and basic text and image inference paths.

For developers and enterprise users, the significance is not simply that “another model can run.” It validates a more important path: whether a domestic full-featured GPU stack can provide rapid support and full-stack coordination for frontier-scale models. This article examines the core technical problems and solutions from an engineering perspective.

2. Technical Characteristics of Ultra-Large-Scale MoE Models

2.1 Parameter Scale and Sparse Activation

A 3-trillion-parameter model typically uses a Mixture of Experts (MoE) architecture. Unlike dense models, MoE models are characterized by a large total parameter count but sparse activation per forward pass. This leads to several implications:

• High memory footprint: All expert weights need to reside in GPU memory or in a storage tier that can be loaded quickly.
• Relatively controlled compute: Only a subset of experts is activated per forward pass, so actual FLOPs are far lower than a dense model of equivalent size.
• Significant communication overhead: When expert parallelism and data parallelism are combined, All-to-All communication becomes a critical bottleneck.

Therefore, the first priority for adapting such models is not “whether there is enough compute,” but “whether the weights fit in memory and whether the interconnect can keep up.”

2.2 KV Cache Pressure from a 1-Million-Token Context

A 1-million-token context window means the KV cache can reach tens of gigabytes or more. During inference, KV cache memory consumption grows linearly with the number of concurrent requests and sequence length. Common mitigation strategies include:

• PagedAttention-style memory management: Partitioning the KV cache into pages to reduce fragmentation.
• KV cache quantization: Storing KV data in INT8 or FP8 to reduce memory bandwidth pressure.
• Prefix caching and reuse: Sharing cached prefixes for identical system prompts or long document prefixes.

These strategies also need to be supported by the domestic GPU software stack. Otherwise, long-context capability remains only theoretical.

3. Full-Stack Coordination for Day-0 Adaptation

Day-0 adaptation means completing a runnable validation on the same day the model is open-sourced. This depends on coordination across four layers: chip, driver, framework, and model, rather than a single-point breakthrough.

3.1 Model Loading: Weight Format Conversion and Sharding Strategy

Weight files for ultra-large-scale models are usually released in shards, with individual files potentially reaching tens of gigabytes. The loading process must address:

1. Format compatibility: Converting training-side weight formats into a format recognized by the inference engine.
2. Parallel sharding: Mapping weights to multiple accelerators according to tensor parallelism (TP), expert parallelism (EP), and pipeline parallelism (PP) strategies.
3. Memory planning: Reserving memory budgets for weights, KV cache, and intermediate activations.

A simplified loading flow looks like this:

# Pseudo-code: multi-device weight loading and sharding
def load_sharded_weights(ckpt_dir, tp_size, ep_size):
    state_dict = {}
    for shard in list_shards(ckpt_dir):
        # Split weights to the corresponding rank according to TP/EP rules
        local_tensors = split_by_parallel(shard, tp_size, ep_size)
        state_dict.update(local_tensors)
    return state_dict

In the domestic software stack, this layer is typically handled by the framework adaptation layer. Developers do not need to manually handle the sharding details of every tensor, but they should understand how the sharding strategy affects memory usage and communication.

3.2 Service Startup: Inference Engine and Scheduler

After model loading, the inference service needs to be started. Key components at this stage include:

• Inference engine: Responsible for operator scheduling, memory management, and batching.
• Request scheduler: Implements continuous batching to improve throughput.
• Health checks and elastic scaling: Ensuring stability during long-running service.

For MoE models, the scheduler also needs to consider expert load balancing. If certain experts are activated frequently, they can become hotspots and turn a single accelerator’s memory or compute into a bottleneck. Common approaches include load-balancing strategies introduced through auxiliary losses during training, or expert replica redundancy on the inference side.

3.3 Enabling Text and Image Inference Paths

The K3.1 Pro model supports both text and image inference, making it a multimodal model. A typical multimodal inference pipeline consists of:

1. Image encoder: Encodes images into visual tokens.
2. Projection layer: Maps visual tokens into the language model embedding space.
3. Language model backbone: Performs unified inference after concatenating visual tokens with text tokens.
4. Output decoding: Generates text results.

During adaptation, it is necessary to ensure that the image encoder operators run efficiently on the domestic GPU and that the concatenation of visual and text tokens does not disrupt KV cache management. For a 1-million-token context, the insertion position and number of image tokens must be carefully controlled to avoid memory overflow.

4. Key Tuning Points in Engineering Practice

4.1 Memory Optimization: Hierarchical Offloading and Quantization

When a single node’s GPU memory cannot hold all expert weights, hierarchical offloading can be used:

• Keep high-frequency experts resident in GPU memory.
• Place low-frequency experts in host memory or high-speed storage and load them on demand.
• Quantize the KV cache: Store it in FP8 or INT8 to reduce bandwidth pressure.

It is important to note that offloading introduces additional transfer latency, so dynamic scheduling based on request characteristics is required.

4.2 Communication Optimization: All-to-All and Overlap

In MoE inference, All-to-All communication is one of the main overheads. Optimization techniques include:

• Overlapping communication and computation: Preparing the next layer’s communication while waiting for expert results.
• Topology-aware routing: Completing expert communication within the same node first to reduce cross-node traffic.
• Hierarchical All-to-All: Performing intra-node communication first, then inter-node communication, to reduce global communication pressure.

4.3 Batching and Concurrency Tuning

Continuous batching can significantly improve GPU utilization. However, in MoE models, different requests may activate different experts, leading to inconsistent computation graphs within a batch. Possible solutions include:

• Group batching by expert: Combining requests that activate the same experts.
• Dynamic batch size: Adjusting based on remaining memory and request queue length.
• Priority scheduling: Handling long-context and short requests separately to prevent mutual blocking.

5. Takeaways for Developers

1. Full-stack capability matters more than single-point performance: Adapting ultra-large-scale models is not as simple as copying weights. It is the result of coordination across chip, driver, framework, and model layers.
2. Memory and communication are the core bottlenecks: At the 3-trillion-parameter scale, compute is often not the primary constraint. Memory capacity and interconnect bandwidth are.
3. Multimodality and long context require specialized optimization: The management of visual tokens and KV cache directly affects service availability.
4. Day-0 adaptation signals ecosystem maturity: The ability to complete adaptation on the day a new model is released indicates that the software stack’s abstraction layers and operator coverage are sufficiently general.

6. Conclusion

The deployment of a 3-trillion-parameter model marks a new phase in large-model engineering. Completing Day-0 adaptation of such a model on a domestic full-featured GPU demonstrates not only hardware capability but also the full-stack coordination of the software stack. For developers, understanding the technical path and tuning methods behind this adaptation can help deploy and optimize ultra-large-scale model services more efficiently in their own environments. As model scales continue to grow, memory management, communication optimization, and multimodal inference paths will remain key areas of ongoing evolution.