Adapting Ultra-Large-Scale MoE Model Inference and Optimizing the Compute Pipeline

As large model parameter counts move toward the trillion scale, mixture-of-experts architectures and ultra-long-context techniques are becoming standard in the industry. However, the massive sparse parameters and increasingly complex attention mechanisms place unprecedented demands on memory bandwidth, compiler toolchains, and software stack optimization for underlying accelerators. This article explores how to achieve Day-0 adaptation of ultra-large-scale MoE models from community code to production-grade inference services on a fully self-developed AI computing foundation, and provides general optimization strategies for sparse attention and low-precision expert computation.

1. Core Inference Challenges for Ultra-Large-Scale MoE Models

At the trillion-parameter scale, models typically adopt MoE architectures to balance compute requirements and inference quality. Adapting such models presents three major challenges:

1. Memory wall and bandwidth bottlenecks: Even with low-precision storage, the weight volume of trillions of parameters remains enormous. At the same time, during the generation phase, KV cache memory usage rises sharply as context length grows, such as to millions of tokens, putting dual pressure on memory capacity and bandwidth.
2. Dynamic routing and computational sparsity: During forward propagation, MoE architectures dynamically select activated experts based on gating networks. This makes the compute graph difficult to fully optimize statically at compile time, increasing the complexity of operator scheduling and memory management.
3. Adaptation of new attention mechanisms: To reduce computational complexity for ultra-long contexts, models often introduce mechanisms such as sparse attention. These mechanisms depart from conventional CUDA Core/Tensor Core computation patterns and require deep customization of the underlying operator libraries.

2. Software Stack Coordination and Compiler Pipeline Integration

Converting large community model implementations into actual inference services on underlying hardware requires not only hardware support, but also close coordination between compilers and the software stack. In this process, the full-stack software system plays a critical role.

1. Intermediate representation and operator migration

When integrating upstream open-source model implementations, engineering teams typically rely on high-performance inference frameworks such as SGLang. At the lower level, self-developed operator libraries, high-level language extensions compatible with Triton, TileLang, and dedicated compilers such as MTCC are used to migrate and connect the critical compute paths of new architectures.

Through these toolchains, developers can gradually lower high-level operators in the model compute graph to underlying hardware instructions. An illustrative code invocation pattern is shown below:

# Pseudocode: loading an MoE model through a customized compilation pipeline
import sglang_musa as sgl
# Initialize the underlying MUSA runtime and operator library
engine = sgl.Engine(
    model_path="/models/trillion_moe",
    musa_kernel_lib="MATE",
    compiler_backend="MTCC",
    enable_tilelang=True
)
# Launch the inference service
server = engine.launch_service()

2. Runtime state and scheduling management

During inference execution, the scheduler must handle complex routing and dispatch logic. To ensure stability, the software stack provides dedicated device scheduling and state management mechanisms, similar to KDA-style mechanisms, responsible for maintaining expert routing tables, managing intermediate states across residual connections between layers, and ensuring seamless integration of vision-multimodal attention with the forward path, thereby constructing a complete execution graph.

3. Deep Optimization of Sparse Attention DSA Mechanisms

For sparse attention mechanisms such as DSA used in ultra-long-context models, the underlying software stack must be deeply adapted. Conventional dense attention computation patterns are no longer applicable here.

1. Index selection and KV cache layout restructuring

In sparse attention, each query does not interact with all keys and values. Instead, relevant tokens are selected through specific algorithms, such as top-K selection or hash-based clustering. During low-level adaptation, efficient index selection operators must be implemented in the operator library, such as MATE, to ensure that the selection process does not become a performance bottleneck.

At the same time, the memory layout of the KV cache must be restructured. Traditional contiguous storage cannot efficiently accommodate sparse access patterns. Block-based or distributed layouts are typically introduced to reduce memory fragmentation and improve actual memory bandwidth utilization.

2. Separation of prefill and decode computation paths

The prefill and decode stages of large model inference differ significantly in computational density. During prefill, the model processes a large number of input tokens and is compute-bound. During decode, tokens are generated one by one and the process is memory-bound.

For DSA mechanisms, the adaptation approach separates these two paths for optimization:

• Prefill path: Fully utilize the parallel computing capability of the hardware to execute sparse attention computations in batches and improve throughput.
• Decode path: Focus on precise retrieval of the KV cache and use fused operators to reduce the number of memory access instruction dispatches, substantially lowering latency during long-context generation.

4. In-Memory Conversion and Computation Strategy for MXFP4 Expert Weights

To aggressively compress the memory footprint of trillion-parameter models, leading-edge models are beginning to store expert weights in the MXFP4 format. However, the Tensor Cores of many current hardware platforms natively support only FP8 or higher-precision matrix multiplication. Efficiently executing MoE expert computation without native FP4 support is a significant engineering challenge.

1. Avoiding offline conversion: preventing additional memory overhead

A conventional approach is to use offline scripts to dequantize FP4 weights into FP8 before model loading and then deploy them to memory. However, once a trillion-parameter model is fully expanded to FP8, the memory increase is often tens or even hundreds of gigabytes, which is catastrophic for single-card or single-node memory capacity.

2. Online conversion and cross-layer workspace reuse

A more advanced adaptation strategy uses online conversion with cross-layer reuse. Specifically:

• Online conversion: When computing a layer’s experts, MXFP4 weights are read from memory and converted to FP8 in real time inside the operator using lightweight instructions. The converted FP8 data is then fed into Tensor Cores for matrix multiplication. Once computation completes, the temporary FP8 data is immediately discarded.
• Cross-layer workspace reuse: Because MoE models have many layers, the size of the temporary FP8 data blocks generated by each layer’s expert computation is fixed. Therefore, the low-level scheduler allocates only a single fixed-size FP8 workspace in memory. After layer N finishes computation, layer N+1 directly reuses the same workspace for weight conversion.

This strategy logically eliminates the need to pre-generate massive offline FP8 weights and keeps the additional memory footprint within a very small range. At the same time, it relies on the underlying compiler to deeply fuse instructions for the online conversion operators, hiding the overhead introduced by dequantization.

5. Summary

In the era of ultra-large-scale MoE models combined with ultra-long contexts, competition among computing foundations is no longer just about transistor density and memory capacity. It is also a deep-water contest involving low-level software stacks, compilers, and operator ecosystems.

The successful adaptation of trillion-parameter models demonstrates that customized compiler toolchains, refined KV cache layouts, and innovative in-memory online conversion techniques can smoothly support complex architectures of frontier large models within limited hardware resources. This kind of system-level engineering practice not only provides an efficient inference path for real-world large model applications, but also offers valuable reference points for the coordinated evolution of future hardware architectures and software algorithms. In general-purpose AI computing infrastructure such as Baidu AI Cloud’s Qianfan large model platform, similar full-stack optimization systems are helping more enterprises access cutting-edge AI capabilities at lower cost and with greater efficiency.