Day-0 Adaptation of Ultra-Large-Scale Models on Domestic AI Compute Infrastructure

Introduction: The Bottom-Line Challenges of Ultra-Large Model Inference

As general-purpose AI evolves rapidly, the industry has seen the emergence of models supporting million-token contexts and parameter counts in the trillions. These models demonstrate exceptional capabilities in agent evaluation, code generation, and long-document understanding, but they impose severe demands on the underlying AI compute infrastructure. Because such models often employ novel sparse attention mechanisms and mixed-precision computation strategies, traditional software stacks and hardware architectures cannot directly support efficient inference.

In this context, achieving “Day-0 adaptation”—that is, completing the full integration from low-level operator adaptation to end-to-end inference service on the very first day a model’s weights are released—has become a critical benchmark for evaluating the technical maturity and software ecosystem readiness of an AI compute platform. This article explores how to overcome the adaptation challenges posed by sparse attention indexing, KV cache layout, and low-bit weight computation in ultra-large models, using a domestic full-function GPU and the MUSA software stack as the reference implementation.

Overall Architecture and Full-Pipeline Integration for Day-0 Adaptation

Day-0 adaptation is not simply a matter of loading model files. It requires software and hardware to work together to close the loop from computation graph compilation to actual service invocation. Given the architectural characteristics of ultra-large models, the underlying software stack must be addressed from three dimensions: inference engine integration, core operator library reconstruction, and cross-layer compilation optimization.

1. Leveraging Open-Source Inference Engines as the Foundation

To quickly inherit algorithm implementations from upstream community models, adaptation work typically uses mainstream open-source inference frameworks such as SGLang as the bridge. By porting these frameworks onto domestic compute platforms, developers can reuse their strengths in high-level computation graph construction and request scheduling. However, to make these frameworks run on domestic GPUs, the mapping between low-level operator scheduling and hardware-specific features must be resolved.

2. Multi-Level Compilation and Operator Library Co-Design

Migrating ultra-large models relies on a complete low-level toolchain, typically including a core operator library such as MATE, intermediate representation and operator development languages such as Triton and TileLang cross-platform variants, and a low-level compiler such as MTCC. With this toolchain, the development team can smoothly transform model code written for general-purpose hardware into actual inference services running on domestic GPUs. The compiler converts high-level computation semantics into low-level hardware instructions, while the operator library provides highly optimized primitive computation blocks. Together, they enable rapid integration of critical computation paths under new architectures.

Solving the Sparse Attention (DSA) Adaptation Challenge

To handle million-token contexts while controlling computational complexity, ultra-large models widely employ innovative sparse attention mechanisms such as DSA. This approach breaks the traditional fully connected attention paradigm of Transformers and imposes entirely new requirements on memory access patterns and operator implementations at the hardware level.

1. Index Selection and KV Cache Layout Optimization

In sparse attention, the model no longer performs global attention across all tokens. Instead, it dynamically selects relevant tokens through a specific algorithm to construct a local attention graph. This means the underlying operators must handle highly irregular memory access patterns.

To address this, the MUSA software stack places emphasis on optimizing the index selection mechanism and the KV cache layout. Traditional contiguous KV cache storage leads to a large number of ineffective data loads during sparse access. By redesigning the memory layout of the KV cache—dividing it into finer-grained physical blocks and establishing an index-based fast lookup table—the system can efficiently aggregate the required Key and Value tensors in non-contiguous memory space, significantly reducing memory access latency.

2. Differentiated Adaptation for Prefill and Decode Paths

Large model inference consists of two core phases: prefill and decode. These phases exhibit distinctly different computational characteristics under sparse attention:

• Prefill Phase: This is a compute-intensive task that processes a large number of input tokens at once. The adaptation focus is on optimizing the fusion of matrix multiplication with sparse selection. Custom sparse attention operators in the MATE operator library enable high-throughput parallel computation.
• Decode Phase: This is a memory-access-intensive task that generates one token at a time. The key challenge lies in efficiently updating the attention state incrementally. The software stack uses a KDA (Key-Data Allocation) state management mechanism to maintain attention context across steps, ensuring smooth execution of the decoding process.

3. Fusing Cross-Layer Residuals and Visual Attention

Modern multimodal ultra-large models include not only text processing but also integrated visual encoders. During adaptation, sparse text attention must form a complete execution chain together with cross-layer residual connections and visual attention modules. By performing operator fusion at the computation graph level within the compiler, the system eliminates frequent write-back and re-reading of intermediate results, ensuring efficient multimodal inference.

Online Conversion and Computation Optimization of Low-Bit Expert Weights

Beyond sparse attention, ultra-large models commonly adopt Mixture-of-Experts (MoE) architectures in their FFN layers and store expert weights in ultra-low-bit formats such as MXFP4. This design significantly reduces memory footprint but presents considerable challenges for domestic GPU adaptation.

1. Abandoning Offline Conversion in Favor of Online Decoding

Typically, if hardware does not natively support the MXFP4 format, offline scripts are used to convert weights to the FPGA-supported FP8 or FP16 format. However, for models with trillions of parameters, generating a full set of offline FP8 weights is extremely time-consuming and consumes enormous additional storage space, defeating the purpose of low-bit compression.

To address this, the MUSA software stack takes an innovative approach: online conversion with cross-layer reuse of a fixed FP8 workspace. Instead of pre-computing all FP8 weights, the system converts only the currently activated expert weights from MXFP4 to FP8 at the moment of inference, during MoE expert computation.

2. Cross-Layer Workspace Reuse and Zero-Copy Memory Management

To compensate for the computational overhead of online conversion, the adaptation team introduced a cross-layer reuse strategy. Since MoE architectures activate only a subset of experts per inference step, the system pre-allocates a fixed-size FP8 workspace in memory. When a layer completes its computation, this workspace is immediately reused by the experts activated in the next layer, eliminating the need for repeated memory allocation.

The following pseudocode illustrates the concept:

# Pseudocode: MXFP4 online conversion and cross-layer workspace reuse
workspace_fp8 = allocate_fixed_workspace() # Pre-allocate shared cross-layer workspace
for layer in model.layers:
    if layer.has_moe:
        # Get indices of experts activated in the current layer
        active_experts = get_active_experts(layer.input)
        for expert_id in active_experts:
            # Load MXFP4 weights from memory, convert to FP8 in the shared workspace
            mxfp4_weights = load_expert_weights(expert_id)
            convert_mxfp4_to_fp8_inplace(mxfp4_weights, workspace_fp8)
            # Directly use the FP8 workspace for expert computation
            output = compute_fp8(layer.input, workspace_fp8)
            # Mark the workspace as reusable for the next expert or layer
            mark_workspace_reusable(workspace_fp8)

With this approach, the system avoids the massive storage cost of offline conversion while achieving high-performance low-bit computation through fine-grained memory management.

Conclusion and Outlook

Completing Day-0 adaptation for ultra-large-scale models at the level of Kimi-K3.1 Pro marks a significant step forward for domestic AI compute platforms in terms of software-hardware co-optimization and ecosystem compatibility. By using the MUSA software stack to restructure sparse attention indexing and KV cache behavior, and by implementing online cross-layer conversion and reuse of MXFP4 expert weights, developers have successfully opened the full pipeline from low-level operators to high-level serving—enabling efficient inference with million-token contexts and multimodal inputs.

Looking ahead, as model parameters continue to grow and network architectures keep evolving, adaptation of the underlying infrastructure will need to happen even earlier in the release cycle. For enterprise users, combining general-purpose cloud-native platforms with strong compute scheduling and containerized deployment capabilities—alongside deep hardware-level optimization for sparse computation and low-bit quantization—will become the standard approach for supporting the efficient iteration of next-generation AI workloads. This end-to-end, full-stack optimization capability is the key to bridging the compute gap and unlocking the full potential of ultra-large models.