Day-0 Adaptation of a Domestic AI Accelerator Card for an Ultra-Large-Scale Model: Sparse Attention and Mixed-Precision Inference Explained

Background: The Daunting Challenge of Day-0 Adaptation

In large model deployment, Day-0 adaptation means completing the entire inference pipeline—from model loading and service startup to basic inference—on the very day new model weights are released. This goal places extremely high demands on the underlying compute platform’s software stack compatibility, operator completeness, and performance tuning capabilities.

Take the recently released Kimi-K3.1 Pro as an example. The model has 3 trillion parameters, supports a 1-million-token ultra-long context, and adopts several cutting-edge architectural designs: a sparse attention mechanism (DSA) for handling long sequences, expert weights stored in MXFP4 format, and cross-layer residual connections along with a visual attention module. While these features significantly improve the model’s performance on agent tasks, coding, and long-range reasoning, they also impose entirely new adaptation requirements on the underlying hardware acceleration libraries—traditional dense attention operators cannot be directly reused, and standard FP16/BF16 inference pipelines cannot efficiently handle 4-bit precision expert weights.

For AI infrastructure teams, completing full-pipeline adaptation on model release day requires not only a deep understanding of the model architecture but also a highly modular software stack with rapid migration capabilities.

1. Modular Software Stack Design: The Key to Inheriting Upstream Model Implementations

Facing the time pressure of Day-0 adaptation, mature software stack architectures typically adopt a layered decoupling strategy. In the case discussed here, the team built on SGLang-MUSA (an inference serving framework compatible with the SGLang interface) to inherit upstream model implementations, and used multi-layered compilers and operator libraries to migrate and integrate the key compute paths of the new architecture.

# Illustration: simplified logic for model loading and operator dispatch
def load_model_and_dispatch(model_config):
    # 1. Parse model config, identify attention type and weight precision
    if model_config.attention_type == "sparse_dsa":
        dispatch_attention_to_mate_sparse_ops()
    if model_config.expert_weight_dtype == "mxfp4":
        enable_online_dequantization_pipeline()

In actual engineering deployment, this process relies on the coordination of three core components:

1. MATE Operator Library: A pre-implemented collection of performance-critical operators, including highly optimized kernels for sparse attention, matrix multiplication, and other common computations. During Day-0 adaptation, the new model’s DSA attention pattern needs to be quickly mapped onto MATE’s existing sparse operator templates rather than writing CUDA-like kernels from scratch.

2. Triton-MUSA and TileLang-MUSA: These two DSL (domain-specific language) toolchains allow developers to write high-performance compute kernels using near-Python syntax that automatically maps to underlying hardware instructions. When MATE lacks a ready-made operator, the team can use these tools to rapidly prototype custom operators before iteratively optimizing them.

3. MTCC Compiler: As the underlying compilation infrastructure, it optimizes DSL code or high-level IR into machine code targeting specific microarchitectures. For the new model’s special computation patterns (such as MXFP4 online conversion), the compiler must support corresponding instruction scheduling and memory access optimization.

2. DSA Sparse Attention Adaptation: Unifying Indexing, Caching, and Compute Paths

The DSA (Dynamic Sparse Attention) mechanism adopted by Kimi-K3.1 Pro introduces dynamic index selection on top of standard attention—for each query position, attention scores are computed only against a small set of selected key-value pairs whose positions are dynamically generated by the model rather than following a fixed pattern. This design dramatically reduces computational complexity in long-context scenarios, but it also makes it difficult for the underlying computation to directly leverage Tensor Core dense matrix multiply units.

2.1 Index Selection and KV Cache Layout

One of the core adaptation tasks is passing dynamic index information from the model layer down to the compute kernel layer. In practice, the team needs to complete the following steps:

• Index Selection: DSA uses an additional score predictor to output the key-value indices that each query position should attend to. These indices exist as variable-length lists and cannot be directly passed to standard operators as fixed-length tensors. In the adaptation scheme, the software stack converts these indices into a pre-filled index matrix and determines the mask padding strategy based on actual sparsity.

• KV Cache Layout Optimization: In long-context scenarios, KV caches are typically managed using PagedAttention-style management to save GPU memory. However, for sparse attention, each query block accesses only a subset of KV pages, so the index selection results must be mapped to physical page numbers while ensuring cache access locality. The MUSA software stack uses a two-level page table mechanism to quickly translate logical indices into physical locations while also accommodating cache reuse across multi-turn inference.

2.2 Unified Abstraction of Prefill and Decode Paths

Sparse attention has vastly different computation patterns in the Prefill and Decode phases:

• Prefill Phase: The input is a complete sequence of length L, requiring attention computation across all positions. Here, DSA’s sparse index matrix can be batch-processed, leveraging matrix multiply units for high throughput.

• Decode Phase: Only one token is generated at a time, reducing the query length to 1. But the KV cache keeps growing, and the sparse indices differ each time, resulting in irregular memory access patterns.

The MATE sparse attention operator switches between the two modes via an internal state machine: in Prefill mode it uses blocked sparse matrix multiplication, while in Decode mode it uses index-based gather instructions to directly collect the required KV blocks, avoiding reads of invalid data.

# Illustration: mode-switching logic of the sparse attention operator
if is_prefill:
    out = sparse_batched_matmul(Q, K, V, index_matrix)
else:
    gathered_K = torch.gather(K_cache, dim=1, index=current_indices)
    gathered_V = torch.gather(V_cache, dim=1, index=current_indices)
    out = scaled_dot_product_attention(Q_block, gathered_K, gathered_V)

2.3 Coordination with KDA State Management and Cross-Layer Residuals

Beyond the core attention computation, the complete execution pipeline also includes KDA (KV-Driven Attention) state management—maintaining KV cache allocation and release across multi-turn conversations—and cross-layer residual connections that inject shallow-layer features directly into deeper layers to shorten gradient propagation paths. Although these modules are not specific to sparse attention, they must work closely with sparse operators during Day-0 adaptation to ensure consistent storage layouts for intermediate results and avoid unnecessary format conversions.

Additionally, the model introduces a visual attention module for handling mixed image-text inputs. During adaptation, the visual encoder’s attention computation is unified into MATE’s dense operator path, sharing the same KV cache pool with the DSA sparse path, and the two types of features are ultimately merged through concatenation.

3. MXFP4 Expert Weights: The Engineering Wisdom of Online Conversion over Offline Quantization

The MoE (Mixture of Experts) architecture is the foundation for Kimi-K3.1 Pro achieving 3 trillion parameters, with only a small number of experts activated per token to control computational cost. However, the expert weights provided at model release are stored in MXFP4 format—a 4-bit floating-point format designed to reduce model size and accelerate weight loading.

3.1 Computational Challenges of MXFP4

The MXFP4 format uses a shared exponent mechanism: several consecutive 4-bit mantissas share a single 8-bit exponent, achieving greater dynamic range within a limited bit width. But this format cannot be directly processed by Tensor Cores. The traditional solution is to pre-convert all weights offline to FP16 or BF16 before loading them into GPU memory. However, for a 3-trillion-parameter model, offline conversion is not only time-consuming but also requires several terabytes of additional storage and conversion compute.

3.2 Online Conversion and FP8 Workspace Reuse

The software-hardware co-design breakthrough lies in an online conversion strategy: during inference, only the activated expert weights are temporarily converted from MXFP4 to FP8 (8-bit floating point) and stored in a workspace before participating in matrix computation. The key optimization points of this approach are:

• Cross-Layer Reuse: Different layers may activate the same experts (especially shared experts). By maintaining a global FP8 workspace, already-converted expert weights can be directly reused, avoiding redundant conversion computation.

• Workspace Capacity Planning: FP8 format requires twice the storage of MXFP4, so the number of concurrently activated experts must be precisely estimated to reserve sufficient GPU memory. In the implementation, the workspace uses a pre-allocation mechanism and dynamically adjusts block sizes based on runtime statistics to prevent fragmentation.

# Illustration: online conversion and workspace reuse
for layer in active_layers:
    experts_to_load = get_active_experts(layer)
    for eid in experts_to_load:
        if not fp8_workspace.contains(eid):
            fp8_weights = dequant_mx4_to_fp8(expert_weights_original[eid])
            fp8_workspace.store(eid, fp8_weights)
    output = moe_linear(hidden_states, fp8_workspace.get_batch(experts_to_load))

3.3 Trade-offs in Conversion Precision

MXFP4-to-FP8 conversion is not lossless—FP8 has only 8-bit precision (typically e4m3 format). Compared to the original MXFP4’s 4-bit mantissa plus shared exponent, it actually retains finer-grained mantissa precision but discards some exponent range. In experimental validation, the impact of this conversion strategy on the model’s final output is within an acceptable range (typically less than 0.5% precision variation), while delivering significant memory bandwidth savings. Notably, online conversion does not increase end-to-end latency because the conversion process can be pipelined with the previous layer’s computation—that is, while layer N is computing, layer N+1’s weight conversion executes in parallel in the background.

4. Day-0 Adaptation Methodology: General Lessons from Practice

This adaptation practice provides a reusable methodology for AI infrastructure teams:

1. Anchor on the Inference Framework: Prioritize selecting the open-source inference framework with the best compatibility with upstream model implementations, and extend hardware-specific optimizations on top of it rather than building a serving engine from scratch.

2. Operator-Level Abstraction First: After receiving model code, first identify all reusable operator patterns in the computation graph, map new architecture components to existing operator templates, and reduce the number of custom kernels.

3. Workspace and Memory Lifecycle Management: The core bottleneck in long-context inference is often memory access. Designing allocation and reuse mechanisms for large blocks like KV caches and weight workspaces in advance is more effective than merely optimizing matrix multiplication computation.

4. Adaptive Precision Strategy: For low-bit weights, there is no need to pursue maximum average precision. Instead, differentiate precision allocation based on the sensitivity of different layers and experts, reducing storage and conversion overhead while maintaining end-to-end effectiveness.

5. Summary and Outlook

The Day-0 adaptation case of Kimi-K3.1 Pro demonstrates the immense potential of software-hardware co-design: through a modular software stack architecture, mapping cutting-edge model features such as sparse attention and mixed-precision weights onto general-purpose accelerators becomes feasible. As model parameter scales continue to grow and context windows further expand, Day-0 adaptation capability will become one of the key indicators for measuring the maturity of AI infrastructure.

For hardware vendors and software stack developers, the future challenge lies in building more automated and intelligent operator mapping and code generation pipelines that reduce manual intervention—for example, using compilers to automatically identify sparse index patterns and generate efficient kernels, or automatically selecting optimal precision conversion strategies based on performance models. This direction will become a core technical driver for democratizing large model applications.