Adapting Trillion-Parameter Sparse Attention Models on Domestic GPUs: Key Technical Practices

1. Background: Adapting Trillion-Parameter Sparse Models on Domestic GPUs

In recent years, large language models have rapidly surpassed the trillion-parameter mark, and an increasing number of cutting-edge models now adopt sparse attention mechanisms and mixed-precision expert weights such as MXFP4 to balance inference efficiency and model capability. Deploying such models places extremely high demands on the underlying compute platform. The hardware must provide sufficient memory and computing power, while the software stack must be able to convert community model code into a working inference service within a very short time window. This is often referred to as “Day-0 adaptation.”

Day-0 adaptation means that on the same day model weights are officially released or open-sourced, or within a very short time frame, the platform should complete model loading, bring up the inference service, and verify the core inference pipeline. This places three major demands on the coordinated design of hardware and software:

1. Operator coverage: New operators such as sparse attention and mixed-precision computation must already have efficient kernel implementations on the target hardware.
2. Compilation and scheduling capability: A multi-level compiler stack must be able to map high-level computation graphs to low-level hardware instructions.
3. End-to-end engineering completeness: The full pipeline from weight loading and KV cache management to Prefill/Decode dual-phase scheduling must form a closed loop.

This article uses a domestic AI training and inference accelerator and its software stack as the technical background to systematically break down the key technical approaches for adapting trillion-parameter sparse attention models.

---

2. Core Challenges of Dynamic Sparse Attention and Adaptation Strategies

2.1 Introduction to Sparse Attention

Unlike standard full attention with O(n²) complexity, sparse attention, such as Dynamic Sparse Attention (DSA), reduces complexity to approximately O(n) by dynamically selecting a subset of tokens for attention computation. Its core mechanism is index selection: at each layer, the model dynamically decides which other tokens to attend to based on the current token’s importance.

This mechanism introduces three adaptation challenges:

Challenge	Description
Hardware implementation of index selection	Requires efficient Top-K selection and scatter/gather operations, and the indices are generated dynamically.
Changes in KV cache layout	Under sparse mode, KV entries are no longer part of a continuous full sequence, requiring support for irregular access patterns.
Differences between Prefill and Decode phases	The Prefill phase can compute indices in parallel, while the Decode phase requires incremental updates step by step.

2.2 Adapting Index Selection and KV Cache Layout

In the adaptation solution built on the domestic GPU software stack, the team accelerated the index selection logic through a custom sparse attention operator library. Specifically:

• Index generation stage: During Prefill, matrix multiplication and a Top-K sorting kernel are used to determine the sparse pattern for each layer in one pass.
• KV cache management: A tile-based storage strategy is adopted. Sparse KV entries are logically partitioned along the attention head and layer dimensions, so irregular access can be mapped to regular tile-level reads and writes, significantly improving cache hit rates.
• State management: A dedicated state management module tracks the sparse mask and valid KV positions for each layer, enabling the Decode phase to quickly locate entries that need incremental updates.

2.3 Differentiated Implementation for Prefill and Decode Paths

The Prefill phase is characterized by long input sequences and high parallelism, while the Decode phase is autoregressive, step-by-step, and latency-sensitive. The software stack handles the two phases differently:

# Conceptual pseudocode: Prefill sparse attention scheduling
def prefill_sparse_attention(Q, K, V, sparse_mask):
    # 1. Generate dynamic sparse indices
    topk_indices = topk_select(Q @ K.T, k=sparsity_ratio)
    # 2. Gather KV according to indices
    K_sparse = gather(K, topk_indices)
    V_sparse = gather(V, topk_indices)
    # 3. Perform attention computation
    return softmax(Q @ K_sparse.T) @ V_sparse
# Conceptual pseudocode: Decode incremental sparse attention
def decode_sparse_attention_step(q_t, kv_cache, state):
    # 1. Retrieve the sparse mask for the current layer from state
    mask = state.get_sparse_mask(layer_id)
    # 2. Compute attention only for KV positions specified by the mask
    attn_out = masked_attention(q_t, kv_cache, mask)
    # 3. Update KV cache and state
    state.update(kv_cache, attn_out)
    return attn_out

In actual engineering implementation, this logic is compiled into low-level GPU kernels and executed efficiently through tile-level scheduling.

---

3. Online Inference for MXFP4 Expert Weights

3.1 Why Online Conversion Instead of Offline Pre-conversion

Trillion-parameter models typically use a Mixture of Experts (MoE) architecture, where expert weights are stored in MXFP4 format to save memory. The traditional approach is to convert MXFP4 weights offline to FP8 or FP16 before deployment, but this introduces two problems:

1. Storage overhead: The full offline FP8 weights are roughly twice the size of the original MXFP4 weights, which is unfriendly for memory-constrained inference scenarios.
2. Lack of flexibility: If precision needs to be changed or fine-tuning is required later, the offline weights must be regenerated.

3.2 Online Conversion and Cross-Layer Workspace Reuse

The adaptation solution adopts an online conversion strategy. During inference runtime, MXFP4 expert weights are converted on demand to an FP8 intermediate representation, and the workspace is released after computation. The core design includes:

• Cross-layer FP8 workspace reuse: Because experts in different layers are not computed simultaneously, with layer-by-layer serial execution, the same FP8 workspace memory can be reused across multiple layers, significantly reducing peak memory usage.
• Streaming conversion pipeline: Weight loading, dequantization, computation, and release form a pipeline that hides conversion latency.
• Precision assurance: A calibration mechanism is introduced during online conversion to keep the numerical error from MXFP4 to FP8 within a controlled range.

This approach makes it possible to deploy trillion-parameter MoE models under limited memory resources and represents an important engineering breakthrough for the domestic GPU software stack.

---

4. Multi-Level Compiler Collaboration: From High-Level Models to Hardware Instructions

Another key factor for achieving Day-0 adaptation is the coordinated operation of the compilation toolchain. A typical technology stack includes multiple layers:

1. High-level scheduling framework: Takes the computation graph representation from upstream models and handles operator fusion, memory planning, and task scheduling.
2. Intermediate representation and tile-level compiler: Further decomposes fused operators into tile-level instructions tailored to the target hardware’s memory hierarchy and parallelism model.
3. Low-level code generation and optimizer: Generates the final GPU kernel code and performs low-level optimizations such as register allocation and instruction scheduling.

In this adaptation effort, the team used the coordinated operation of the multi-level compiler stack to progressively lower community model code from the Python/framework layer into an executable inference service on the target accelerator, achieving end-to-end conversion from model weights to an online service.

---

5. Complete Execution Pipeline and Key Engineering Practices

Overall, a successful Day-0 adaptation for a trillion-parameter sparse model requires establishing the following complete pipeline:

1. Model loading: Support sharded loading and verification of extremely large weights.
2. Sparse attention execution: Index selection → KV gather → attention computation → state update.
3. MoE expert computation: MXFP4 online conversion → FP8 workspace reuse → expert routing and aggregation.
4. Cross-layer residuals and visual attention: For multimodal models, the attention path of the visual encoder must also be adapted.
5. Service startup and inference verification: Complete end-to-end testing of basic inference pipelines such as text generation and image understanding.

For enterprise users, choosing a domestic compute platform with a well-developed software stack means shorter adaptation cycles for bringing frontier models into real business scenarios and reduced dependence on a single technology ecosystem.

---

6. Summary and Outlook

Adapting trillion-parameter sparse attention models on domestic GPUs is fundamentally a systematic engineering effort spanning operator development, compilation optimization, memory management, and end-to-end pipeline construction. From sparse attention index and cache adaptation to online inference with MXFP4 expert weights, and to coordinated scheduling across multi-level compilers, every step requires deep hardware-software integration. As domestic AI compute platforms continue to evolve, this kind of challenging adaptation work will gradually become routine, providing a solid technical foundation for the domestic deployment of large models.