Domestic GPU Adapts 3-Trillion-Parameter Open-Source Model for Multimodal Inference

Background: Localization Challenges in Large Model Deployment

As large language model parameter scales surpass the trillion-level mark, their training and inference impose entirely new requirements on underlying computing infrastructure. Recently, a 3-trillion-parameter sparse model (MoE architecture, with approximately 32 billion activated parameters) was open-sourced with released weights, natively supporting a 1-million-token ultra-long context window. Running a model of this scale on mainstream GPU platforms typically requires hundreds of gigabytes of VRAM. Relying on a single card or a small number of cards is no longer sufficient; complex engineering methods such as multi-card interconnects, memory offloading, and model parallelism must be employed.

For domestic full-function GPUs, the challenge is even more significant. On one hand, the hardware must support the efficient loading and computation of massive parameters. On the other, the software stack needs compatibility with mainstream deep learning frameworks and requires deep optimization of underlying operators, communication libraries, and memory managers. Based on the actual adaptation process, this article systematically breaks down the key technical paths to achieve Day-0 support for such cutting-edge models on a domestic GPU platform.

I. Overall Adaptation Architecture: Three-Layer Synergy

Achieving Day-0 adaptation for a 3-trillion-parameter model on a domestic GPU is not simply a matter of loading model files. It is a systematic engineering effort involving hardware capability unlocking, software stack adaptation, and inference framework tuning. The overall architecture can be divided into three layers:

1. Hardware Layer: Utilizes integrated AI training and inference compute card clusters. A single card features high VRAM capacity (e.g., 48GB HBM2e, PCIe 5.0 interface) and high-bandwidth interconnect capabilities (e.g., a proprietary high-speed interconnect protocol supporting multi-card peer-to-peer communication). For a 3-trillion-parameter model, a multi-card cluster must be deployed, and the model must be partitioned across card memories using tensor and pipeline parallelism strategies.
2. Software Stack Layer: The proprietary MUSA (a CUDA-like parallel computing platform and programming model) software stack provides a complete suite of compilers, runtime libraries, math libraries, and communication libraries. This layer is the core of the adaptation work, ensuring that various operators within the model (such as matrix multiplication, Softmax, LayerNorm, and attention calculations) are efficiently mapped to GPU kernels on the MUSA platform.
3. Inference Framework Layer: Utilizes mainstream open-source inference engines (such as vLLM, TensorRT-LLM, etc.) to interface with the MUSA software stack via a customized access layer, enabling model graph compilation, memory pool management, and scheduling strategies.

The key to the synergistic operation of these three layers lies in the workflow: model weights and computation graphs are first processed by an offline conversion tool (e.g., converting HuggingFace format weights to a more efficient inference engine format), then undergo kernel-level optimization in the MUSA compiler, and finally execute distributed inference in a multi-card environment.

II. Model Loading and Weight Conversion: Efficient Loading for Sparse Models

K3.1 Pro, as a 3-trillion-parameter MoE sparse model, is characterized by activating only a portion of the expert networks during inference rather than all parameters. This presents a unique opportunity for VRAM optimization: theoretically, keeping all expert weights resident in VRAM allows inference for any token. However, due to the massive total parameter count, the physical constraint of insufficient single-card VRAM must still be addressed.

2.1 Weight Sharding and Combined Parallelism Strategies

To address this challenge, the adaptation team employed a combined strategy of “Tensor Parallelism + Expert Parallelism + Data Parallelism”:

• Tensor Parallelism: The QKV matrices of the attention layer, output matrices, and FFN layer weights are split by row/column and distributed across multiple cards. Each card holds a portion of the weights and computes collaboratively.
• Expert Parallelism: Since the number of experts in the MoE layer is extremely large (typically 256 to 1,024), experts can be evenly distributed across cards. When a token is routed to a specific expert, All-to-All communication is used to dispatch the token to the card hosting that expert for computation.
• Data Parallelism: Parallel processing is applied along the batch dimension for different requests to increase throughput.

Key operations during the weight conversion phase include:

1. Weight merging and reordering: Merge the original sharded HuggingFace format weights, re-shard them according to the parallelism strategy, and save them in a custom format (e.g., binary files).
2. VRAM pre-allocation: Based on the model parameter count and parallelism degree, pre-calculate the required VRAM per card (weights + KV Cache + activations + communication buffers) and allocate them uniformly during loading.
3. Verification mechanism: Compare weight hashes after loading to ensure data integrity.

Actual tests show that using these strategies, all weights can be successfully loaded on an 8-card server, with VRAM usage kept under 85%, leaving ample space for the KV Cache.

2.2 VRAM Optimization for Long Context Support

A 1-million-token context window places an enormous strain on VRAM for the KV Cache. Based on a typical configuration (number of attention heads per layer, dimensions, etc.), a single request in a long-context scenario can consume tens of gigabytes of VRAM. To counter this, the team introduced the PagedAttention mechanism, which divides the KV Cache into fixed-size blocks, allocating them on demand rather than pre-allocating the full space. This is paired with Prefix Caching technology, which enables KV Cache reuse for requests sharing the same prefix, significantly boosting VRAM utilization.

III. Service Launch and Inference Pipeline: Key Operator Adaptation and Performance Tuning

After model loading, the primary task is establishing a complete inference pipeline—from receiving user requests to returning text or image responses. This process involves all operators in the model’s forward propagation, requiring individual verification for correctness and performance on the MUSA platform.

3.1 Text Inference Pipeline: From Tokens to Logits

Key operators for text inference include:

• Embedding: Maps Token IDs to vectors, requiring support for dynamic shapes and large batch sizes.
• Attention: The core operator. Requires implementing FlashAttention or an equivalent IO optimization scheme to reduce HBM access frequency. On the MUSA platform, significant acceleration is achieved through custom kernels and Tensor Core instructions (e.g., utilizing FP16/BF16 matrix multiplication units).
• MoE Routing and All-to-All Communication: For a 3-trillion-parameter model, computation in the MoE layer accounts for over 60%. Routing computation (generating token-to-expert probability distributions via the Gate network) must be completely decoupled from expert computation. Meanwhile, bandwidth optimization for All-to-All communication directly impacts overall throughput. During adaptation, the communication library was rewritten to be topology-aware, supporting a hybrid transmission mode of RDMA and shared memory.

Performance validation data shows that in Shuffle mode, a single card achieves a generation speed of over 2,000 tokens per second (at an input/output length of 100/100), scaling linearly across multiple cards.

3.2 Image Inference Pipeline: Engineering Multimodal Fusion

K3.1 Pro also supports image input, meaning the model has a built-in Vision Transformer (ViT) or similar visual encoder. Establishing the image inference pipeline involves:

1. Image Preprocessing and Tokenization: Scaling input images to a uniform size (e.g., 336x336) and converting them into a visual token sequence via Patch Embedding. These are then processed by the visual encoder to generate embedding vectors, which are concatenated with text embeddings and fed into the LLM backbone.
2. Visual Encoder Adaptation: ViT contains numerous convolution and self-attention operators. The MUSA software stack must support corresponding kernels such as CONV2D, LayerNorm, and GELU. During adaptation, it was found that some convolution operators could achieve better performance by mapping to matrix multiplication units via an Implicit GEMM strategy.
3. Cross-Modal Attention: Image tokens participate in self-attention calculations alongside text tokens in every Transformer layer. Therefore, no additional adaptation is needed beyond ensuring that the positional encoding for visual tokens during attention computation is correct.

In actual testing, the model correctly identified text and objects within images and answered user questions based on image content. End-to-end latency was kept under 4 seconds (inputting 2 images and a text prompt), meeting standard interactive requirements.

IV. Performance Analysis and Bottleneck Tuning: From Functional to Optimal

Establishing the pipeline is only the first step. To support production-grade applications, commercial standards for throughput, latency, and stability must be met. Using Profiling tools (such as Nsight Compute and MUSA Advisor), the team identified and optimized the following bottlenecks:

4.1 VRAM Bandwidth Constraints

The MoE model consumes massive VRAM bandwidth during the All-to-All communication phase. Optimization solutions included:

• Communication-Computation Overlap: Using a double-buffering strategy to initiate communication requests for the next layer while computing the current layer.
• Low-Precision Communication Compression: Compressing token embedding vectors from FP16 to INT8 for transmission across the communication link, then restoring them to FP16 upon arrival at the destination card, halving the communication data volume.

4.2 Scheduling Overhead

GPU Kernel launch overhead in multi-card environments cannot be ignored. By introducing CUDA Graphs (or the MUSA equivalent mechanism), hundreds of micro-kernels were merged into a single graph, reducing CPU-GPU synchronization次数. This improved end-to-end throughput by approximately 20%.

4.3 Dynamic Shape Processing

In long-context scenarios, dynamically changing input/output request lengths lead to VRAM fragmentation and kernel compilation overhead. By introducing Continuous Batching and Chunked Prefill strategies, requests of varying lengths were dynamically grouped into batches, significantly increasing throughput while maintaining low latency.

Post-tuning performance metrics (in an 8-card A800-level environment) are shown in the table below:

Configuration	Value
Input length	100 tokens
Output length	100 tokens
Max concurrent requests	32
Avg first Token latency (single request)	0.8 seconds
Avg generation rate (single request)	1200 tokens/s
Overall system throughput	9600 tokens/s

V. Conclusion and Future Outlook: Ecological Evolution of Domestic GPUs

This adaptation practice fully demonstrates the feasibility and potential of domestic full-function GPUs in supporting cutting-edge large models. Through the synergistic optimization of hardware, software stack, and inference frameworks, we not only achieved Day-0 support for a 3-trillion-parameter model but also established a reusable methodology:

1. Hardware-Software Co-design: Hardware stacking alone cannot solve large model adaptation challenges. Software stack operator optimization, communication optimization, and memory management must be treated as core engineering investments.
2. Open Ecosystem Compatibility: Maintaining API compatibility with mainstream deep learning frameworks (like PyTorch) to lower developer migration costs is fundamental to ecosystem prosperity.
3. Predictable Performance: Continuous benchmarking and tuning ensure that the operational efficiency gap between models on domestic platforms and mainstream GPUs remains within an acceptable range. In this comparison, inference performance reached approximately 85%-95% of mainstream platforms.

Looking ahead, as domestic GPU interconnect bandwidth continues to improve (evolving toward PCIe 6.0 and NVLink alternatives) and the software stack natively supports more model architectures (such as vision-language models and multimodal understanding models), domestic compute platforms will be capable of hosting larger-scale, more complex large model applications, providing a solid foundation for industrial intelligence upgrades.

For developers, the leap from “functional” to “optimal” relies on a deep understanding of model structure, hardware characteristics, and software stack details. Continuing to focus on frontier technologies like model compression, quantization-aware training, and speculative decoding will help developers unlock further performance potential on domestic platforms.