Optimizing Ultra-Large-Scale MoE Model Inference and Adapting to Domestic AI Compute

In recent years, as large models scale toward trillions of parameters, Mixture-of-Experts (MoE) architectures have become a mainstream technical path for breaking through compute bottlenecks, thanks to their sparse activation characteristics. However, while MoE models deliver performance gains, they also introduce unprecedented challenges for hardware adaptation and inference engineering deployment. This article explores how full-stack software optimization can address the inference adaptation challenges of ultra-large-scale MoE models on domestic AI compute platforms.

1. Core Challenges in Ultra-Large-Scale MoE Model Inference

1.1 Memory Wall and Memory Bandwidth Bottlenecks

Trillion-parameter models impose extremely high requirements on memory capacity. In a Mixture-of-Experts architecture, although only a subset of experts is activated per inference, all expert weights must remain resident in memory. At traditional FP16/BF16 precision, model weights alone require several terabytes of storage, far exceeding the memory capacity of a single device. More critically, the frequent expert routing operations during MoE inference generate large amounts of irregular memory access, placing immense pressure on hardware memory bandwidth.

1.2 Engineering Challenges of Sparse Attention Mechanisms

To support context lengths in the millions, emerging sparse attention mechanisms reduce complexity by dynamically selecting key tokens for computation. However, this approach breaks the regular matrix computation pattern of traditional dense attention:

• KV cache layout must support dynamic index selection
• The Prefill and Decode phases have differentiated requirements for attention computation paths
• The index selection algorithm itself may introduce additional computational overhead

1.3 Balancing Low-Bit Precision and Performance

To reduce memory footprint, the industry commonly adopts ultra-low-bit quantization formats such as MXFP4 to store expert weights. However, directly executing 4-bit computation on hardware presents challenges:

• Incomplete support in hardware compute units
• Accumulated quantization errors affecting model quality
• Data format conversion overhead when reusing workspaces across layers

2. Full-Stack Adaptation Solution for Domestic AI Compute

2.1 Computation Graph Compilation and Operator Optimization

Given the dynamic characteristics of MoE models, a full-stack software solution must provide comprehensive computation graph optimization capabilities:

• Operator library layer: Dedicated implementations such as sparse attention operators and dynamic routing operators cover the key computation paths of MoE models
• Compiler layer: Intermediate representations such as TileLang enable cross-layer operator fusion and memory reuse
• Runtime layer: Efficient memory pool management and streaming execution scheduling reduce memory fragmentation

2.2 Complete Adaptation of Sparse Attention Mechanisms

Taking Dynamic Sparse Attention (DSA) as an example, adaptation work must span multiple layers from algorithm to hardware:

# Illustrative pseudocode: sparse attention computation flow
def sparse_attention(query, key_cache, value_cache, topk_indices):
    # 1. Dynamically select relevant KV blocks based on query vectors
    selected_keys = gather(key_cache, topk_indices)
    selected_values = gather(value_cache, topk_indices)
    # 2. Compute local attention scores
    attn_scores = torch.matmul(query, selected_keys.transpose(-1, -2))
    attn_weights = softmax(attn_scores / math.sqrt(head_dim))
    # 3. Apply cache layout optimization
    output = torch.matmul(attn_weights, selected_values)
    return output

At the hardware adaptation level, the following aspects require focused optimization:

• KV cache layout: Adopt a block-based storage structure that supports non-contiguous memory access
• Index selection: Implement efficient Top-K selection algorithms in the MATE operator library
• Computation path branching: Use large-block matrix computation during Prefill and optimize for small-batch queries during Decode

2.3 Online Conversion Solution for MXFP4 Expert Weights

To address the lack of native MXFP4 compute support on hardware, an innovative online conversion solution provides a path that balances precision and performance:

• Workspace reuse: Share FP8 workspaces across layers to avoid memory waste from per-layer allocation
• Streaming conversion: Perform weight decompression synchronously during expert routing, overlapping computation and communication
• Mixed-precision execution: Maintain high-precision computation for activations while using low-bit formats only for expert weights

2.4 End-to-End Multimodal Inference Pipeline

Modern large models process not only text but also multimodal inputs such as images. This requires the underlying compute stack to provide:

• Unified computation graph representation for visual feature extraction and text embedding
• Operator fusion for cross-modal attention mechanisms
• Full-pipeline adaptation for residual connections and state management

3. Key Optimization Points in Engineering Practice

3.1 Parallelism Strategies and Memory Optimization

• Tensor parallelism: Partition expert weights along specific dimensions and synchronize intermediate results through collective communication
• Pipeline parallelism: Divide the model by layers and overlap computation with communication using micro-batching
• Expert parallelism: Distribute different experts across compute nodes and schedule execution through dynamic routing

3.2 Communication Overhead Optimization

All-to-All communication in MoE models is a primary performance bottleneck. Optimization approaches include:

• Communication-computation overlap: Prefetch the next batch of data during expert computation
• Topology-aware scheduling: Optimize data transfer paths based on network topology
• Quantization and compression: Use low-precision transmission for communication data

3.3 Dynamic Batching and Request Scheduling

For variable-length inputs and dynamic activation characteristics, the inference engine must support:

• Continuous batching: Dynamically insert new requests to improve GPU utilization
• Expert load balancing: Prevent popular experts from becoming computational hotspots
• Priority scheduling: Guarantee SLAs for high-priority requests

4. Technical Ecosystem and Future Outlook

Domestic AI compute has made significant progress in software-hardware co-optimization, but ecosystem development still requires sustained investment. Future directions include:

1. Operator library standardization: Promote interface standardization for key operators such as sparse attention and dynamic routing
2. Compiler optimization: Enhance automatic optimization capabilities for dynamic control flow
3. Multi-device collaboration: Optimize communication efficiency and fault tolerance in large-scale clusters
4. Toolchain completeness: Provide end-to-end tools covering model conversion, performance analysis, and deployment monitoring

Through continued technical innovation and ecosystem development, domestic AI compute is expected to achieve broader adoption in ultra-large-scale model inference, laying a solid foundation for building self-reliant intelligent compute infrastructure.