Rapid Adaptation and Inference Optimization for Ultra-Large Parameter Models Based on the MUSA Software Stack

In the field of artificial intelligence, the exponential growth of model scale and the fragmentation of hardware ecosystems are posing unprecedented challenges for infrastructure teams. When a cutting-edge model with trillions of parameters, a Mixture of Experts (MoE) architecture, and sparse attention is released, the ability to quickly bring up a complete inference pipeline without sacrificing performance has become a litmus test for a computing platform’s technical depth and engineering strength. This article reconstructs a real adaptation case and analyzes how a highly compatible heterogeneous computing software stack (hereafter referred to as “the software stack”) was used to complete deep adaptation and inference performance optimization for an ultra-large multimodal model.

1. Engineering Challenges Posed by Ultra-Large Parameter Models

This class of frontier models is not only massive in parameter count; its architectural design also goes far beyond traditional Transformer variants, placing multidimensional demands on the underlying compute resources:

1. Communication bottleneck of the Mixture of Experts (MoE) architecture: The model contains a large number of expert networks, but only a subset of experts is activated during each inference pass. This leads to extremely frequent All-to-All communication between devices, placing stringent requirements on multi-device interconnect bandwidth and communication scheduling strategies.
2. Complexity of the sparse attention mechanism: Unlike conventional dense attention, sparse attention requires dynamically computing correlation indices between tokens. This involves index generation, sorting, and KV Cache layout rearrangement, which breaks the conventional matrix multiplication pipeline.
3. Fragmentation between mixed precision and weight precision: To control memory usage while maintaining model quality, such models often adopt a layer-wise precision strategy. For example, some core parameters use higher precision such as BF16, while certain expert weights may use extremely low-bit compressed formats such as MXFP4. This places demands on the precision types supported by the operator library and the runtime data conversion logic.
4. Long sequences and multimodal inputs: Supporting context lengths of up to a million tokens means that the Prefill stage involves an enormous amount of computation, while also imposing strict requirements on KV Cache management efficiency. In addition, the heterogeneous computation flows of the vision encoder and text decoder require the software stack to provide unified task scheduling.

2. The Foundation for Day-0 Adaptation: A Highly Compatible Software Stack

Achieving adaptation on the same day a new model is released does not depend on rewriting all the code. Rather, the key lies in whether the software stack has strong “absorption” and “transformation” capabilities. The core advantages of this software stack are reflected in three aspects:

• Seamless framework-level integration: By maintaining deep compatibility with mainstream inference frameworks in the community (such as SGLang) and extending the backend on top of them, the new model can reuse most of the community’s native code. The team does not need to write server-side logic from scratch and can focus on modifying the compute paths that are tightly coupled to the underlying hardware.
• Collaboration between the compiler and operator library: A comprehensive operator library (such as MATE) addresses the performance of frequently used standard operators such as matrix multiplication and normalization. When facing custom operators or fused operators introduced by new models, the software stack must provide flexible programming entry points. By supporting domain-specific languages such as Triton and TileLang, combined with the underlying compiler for kernel code generation, developers can efficiently write high-performance kernels tailored to specific compute patterns, greatly shortening the operator development cycle for new models.
• Seamless runtime and communication integration: Distributed inference for large models requires a stable runtime environment. Through standardized device management interfaces and communication libraries, the software stack abstracts away underlying topology differences, allowing upstream frameworks to invoke communication primitives without awareness of hardware details. This provides a stable foundation for multi-device collaborative computation.

3. Migration and Deep Optimization of Critical Inference Paths

After completing the basic code migration, deep optimization targeting the model’s unique computational characteristics is essential to ensure that it not only “runs” but also “runs fast.”

3.1 Fully Enabling the Sparse Attention Computation Path

This was the most challenging part of the adaptation. The standard FlashAttention operator could not be used directly, so the attention computation pipeline had to be restructured:

• Index selection and reordering: First, the Key-Value index ranges that need to be attended to are dynamically computed based on the semantics of the current Query. This process involves a large number of non-aligned memory accesses and sorting operations. By using the atomic operations and custom kernels provided by the software stack, we implemented an efficient Top-K index selector.
• Customized KV Cache layout: Because the indices are sparse, the KV Cache cannot use a contiguous memory layout. We designed a block-based storage scheme and mapped different index patterns (such as local windows, global tokens, and specific columns) to different block scanning strategies during the Prefill and Decode phases to ensure memory access efficiency.
• Operator fusion and execution scheduling: The sparse index computation was fused with the subsequent Gather operations to reduce kernel launch overhead. At the same time, by using the software stack’s event stream mechanism, the attention computation of the vision branch and the sparse attention computation of the text branch were placed into different compute streams, enabling hardware-level concurrent execution and effectively hiding the latency introduced by multimodal inputs.

3.2 Online “Translation” Strategy for Extremely Low-Bit Expert Weights

For the MXFP4-format expert weights used by the model, compute units cannot efficiently process them if they are stored in their original format. The traditional approach performs offline full weight conversion, but for ultra-large models, this leads to lengthy conversion times and additional storage overhead.

We innovatively adopted an online conversion and workspace reuse mechanism:

• On-demand conversion: During model loading, no global data movement is performed. Only when a specific expert is activated and needs to participate in computation are its weights loaded from main memory and converted in real time to the higher-precision format required for computation.
• Cross-layer workspace reuse: To avoid reallocating memory for each conversion, we reserved a fixed “FP8 workspace” in device memory. Because MoE models typically activate only a very small number of experts per request, this allows different network layers and different expert nodes to reuse the same physical memory across time slices. Although these buffers logically belong to different layers, their physical storage addresses are continuously overwritten. This strategy keeps peak memory usage at an extremely low level, compressing the memory requirements that would otherwise demand multiple high-end accelerators into a configuration that can run in more constrained hardware environments.

3.3 Memory Lifecycle Management for Multimodal and Long-Sequence Scenarios

With context lengths of up to a million tokens, KV Cache memory usage can reach tens of gigabytes. Combined with the intermediate activations of the vision encoder, this can easily lead to out-of-memory conditions.

We built a refined hierarchical memory management system:

• Separation of static and dynamic memory: Model weights and core structures use a static memory pool to prevent fragmentation, while KV Cache and activations use a dynamic allocation pool.
• Deep customization of PagedAttention: We introduced the paged attention mechanism at the framework level, but extended the page table structure to accommodate the model’s sparse characteristics. The extended structure records not only physical block addresses but also valid token masks within each block, enabling efficient physical block reclamation and allocation under sparse indexing.

4. Summary and Outlook

Through the multidimensional technical approaches described above, we successfully achieved rapid adaptation and efficient inference for a trillion-parameter sparse model. This not only demonstrates that the software stack possesses strong capability to absorb frontier models, but also highlights its engineering flexibility when handling complex compute patterns such as sparsity, low-bit precision, and multimodality.

From this practice, we can distill three core trends in large model infrastructure development:

1. The “compatibility” of a software stack has more strategic value than raw “performance”: In an era where models evolve rapidly, the ability to absorb new architectures with minimal code changes is a prerequisite for fast iteration.
2. Deep integration of algorithms and engineering is the key to breakthroughs: When facing sparse attention or special precision formats inside models, relying solely on general-purpose operator libraries is far from sufficient. It is necessary to go deep into the operator internals and design collaboratively with the model’s characteristics.
3. Memory strategy determines service cost: At a time when compute resources are expensive, squeezing the limits of memory through online conversion, workspace reuse, and smarter cache management is the most effective path to reducing the service cost of trillion-parameter models.

As models continue to evolve toward greater sparsity and multimodality, we will continue to refine this vertical optimization system spanning operator compilation, compute execution, and resource scheduling, providing a solid compute foundation for the deployment of next-generation AI applications.