Day-0 Adaptation for Domestic AI Chips: Deep Dive into Sparse Attention and Mixed-Precision Inference

From “Usable” to “Optimized”: The Evolution of LLM Inference Adaptation

As large language model (LLM) parameter scales surge into the trillions and context windows break the one-million-token barrier, inference infrastructure is facing unprecedented engineering challenges. For AI chip vendors, merely supporting mainstream operators at the hardware level is no longer sufficient. The true watershed lies in whether they can complete full-stack adaptation—spanning framework integration, operator mapping, and compute path reconstruction—within the shortest possible time after a new model’s release. This capability is widely known in the industry as “Day-0” support.

This article uses the successful adaptation of a 3-trillion-parameter model with a 1-million-token context window on a domestic AI chip and its proprietary MUSA software stack as a blueprint. It provides an in-depth analysis of the core technical hurdles involved, particularly the underlying adaptation principles of sparse attention, the online conversion mechanism for mixed-precision expert weights, and memory/bandwidth optimization strategies under a unified training-inference architecture.

I. The Architectural Foundation of Day-0 Adaptation: Unified Training-Inference Cards and Full-Stack Compiler Synergy

1.1 Design Philosophy of Unified Training-Inference Hardware

This adaptation work is based on a computing accelerator card designed specifically for dual-scenario AI training and inference. Its core advantage lies not in the extreme stacking of a single metric (such as peak compute power), but in its adaptability across the entire model lifecycle—from gradient descent in pre-training and efficient parameter updates during fine-tuning, to low-latency inference in deployment. This design allows the same hardware resource pool to be flexibly scheduled, avoiding the massive data transfer overhead between separate training and inference cards found in traditional decoupled architectures.

In terms of memory management, unified training-inference cards are typically equipped with high-bandwidth memory (HBM) and support fine-grained memory pooling technologies. This is crucial for accommodating the weights of a 3-trillion-parameter model and the KV Cache required for million-token-level contexts. Empirical tests show that the KV Cache memory required to support a million-level context often amounts to hundreds of gigabytes, demanding hardware that not only “fits” the data but can also “read” it quickly.

1.2 MUSA Software Stack: Multi-Tier Compiler Synergy

Faced with the rapid iteration of new models, the traditional approach of hand-writing CUDA-style kernels is unsustainable. The MUSA software stack adopts a strategy of multi-tier Intermediate Representations (IR) and compiler toolchain synergy, building a “translation highway” from high-level frameworks down to underlying hardware:

• Framework Integration Layer: Through the MUSA backend of the SGLang Serving framework, it directly interfaces with upstream model definitions and scheduling logic, ensuring API compatibility for service-oriented deployment.
• Operator Library Layer: The MATE high-performance operator library provides critically optimized basic kernels such as dense matrix multiplication and FlashAttention, serving as “plug-and-play” components for critical paths.
• Programmable Layer: Triton-MUSA and TileLang-MUSA offer Python-like Domain Specific Languages (DSLs), allowing algorithm engineers to write custom operators with pseudo-code-like simplicity, which the compiler then automatically maps to the underlying instruction set.
• Compiler Layer: The MTCC compiler is responsible for lowering the aforementioned high-level IRs into efficient machine code, executing optimizations including loop unrolling, vectorization, and shared memory reordering.

The core value of this combined approach lies in the flexibility of “path selection.” For standard operators, high-performance implementations from MATE are called directly. For unconventional operators introduced by new models, DSLs are used for rapid writing and compilation verification, eliminating the lengthy low-level C++ development cycle.

II. Deep Dive: Four Key Challenges in DSA Adaptation

The sparsification of the attention mechanism is a critical technological route for current LLMs to break through ultra-long contexts. The Dynamic Sparse Attention (DSA) mechanism employed by this model does not use a simple fixed sparse pattern; instead, it dynamically selects the subset of tokens to participate in computation based on the input. This imposes stringent requirements on the inference system’s memory management and operator implementation.

2.1 Index Selection: From CPU Fallback to Real-Time GPU Computation

In traditional dense attention, there is no need for an index generation phase. DSA, however, requires the system to determine the relevant Key positions for the Query before computing each layer. The first decision point for the adaptation team was: where should the index generation module reside?

• Initial Solution: Execute index computation on the host CPU and transfer the results via PCIe. This method is simple to implement but introduces additional synchronization latency. The round-trip communication required at every layer severely bottlenecks the Time-To-First-Token (TTFT) during the Prefill phase.
• Optimized Solution: Rewrite the index generation logic as a MUSA kernel, executing it entirely on the GPU. Specifically, Tensor Cores are utilized to accelerate the preliminary similarity scoring between Query and Key, followed by a custom reduction operator to select the Top-K indices. This change compressed the index generation overhead to the microsecond level, completely eliminating the CPU-GPU synchronization bottleneck.

2.2 KV Cache Layout: Storage Reorganization for Sparse Access

The data layout of the KV Cache dictates memory access efficiency. Dense attention stores KV pairs sequentially in continuous order, resulting in highly regular access patterns. However, sparse attention only requires access to specific KV positions, leading to:

1. Non-contiguous Access: A continuous storage layout results in extremely low cache line utilization.
2. Memory Fragmentation: Dynamic selection means the active token set differs for each layer. Reusing a static layout would require reserving a large number of empty slots.

The adaptation solution employs a Block-wise Padding strategy: the KV Cache is divided into fixed-size blocks (e.g., 64×128), with each block storing KV values of continuous indices. After index computation is completed, the GPU’s DMA engine performs a batch Gather operation, aggregating the KV blocks corresponding to the sparse indices into a continuous, high-speed buffer. This “gather first, compute later” paradigm transforms irregular memory access into regular batch copies and matrix multiplications, significantly improving empirical bandwidth utilization.

2.3 Differentiated Implementation for Prefill and Decode Paths

Sparse attention exhibits vastly different computational characteristics during the Prefill phase (processing the complete prompt) and the Decode phase (generating tokens one by one):

• Prefill Phase: Requires parallel processing of all input tokens, resulting in a massive computational load. The adaptation team fused index selection, KV aggregation, and sparse attention computation into a single kernel. By parallelizing the scan of all Query index sets, it maximizes the parallel throughput of Tensor Cores. The key here is pipelining the overlap of “Gather and Compute”—prefetching the next block’s KV data while computing the current block’s attention.
• Decode Phase: Generates only one token at a time, resulting in low computational density and high memory access density. The focus here shifts from computational throughput to data transfer efficiency. The optimization strategy involves loading only the active set of KV blocks for the current token, utilizing the MUSA software stack’s async copy instructions to prefetch all required KV data into registers or shared memory, thereby avoiding repeated accesses to the HBM.

2.4 Integration with KDA State Management, Cross-Layer Residuals, and Visual Attention

The complete execution chain is not an isolated computation of sparse attention; it must be deeply integrated into the model’s overall computational graph:

• KDA State Management: The Key-Value Dynamic Allocator (KDA) is responsible for dynamically allocating and reclaiming KV Cache blocks. In sparse mode, state management must be aware of Layer ID and Head ID, maintaining an independent block bitmap for each head to prevent write conflicts across multiple heads.
• Cross-Layer Residual Connections: Residual stream computations exist as dependencies both before and after sparse attention, and their data flow is independent of the indices. By inlining and reusing the input memory of the residual addition with the parsed output buffer of the sparse attention, two large data transfer overheads are eliminated.
• Visual Attention Co-processing: Since this model supports multimodal inputs, the layout of visual tokens and text tokens in the KV Cache must be uniformly managed. The adaptation solution treats visual features as a sequence prefix and forces this prefix to remain in “full attention” mode (i.e., excluded from sparse index filtering) to ensure cross-modal semantic alignment.

III. Mixed-Precision Weight Loading: The Art of Online Conversion for MXFP4 Expert Weights

In the Mixture of Experts (MoE) architecture, the explosive growth in model parameters primarily stems from expert networks. To control memory footprint, this model stores expert weights in the MXFP4 format, where each weight value is represented by only a 4-bit mantissa plus a shared scaling factor. However, underlying hardware compute units typically only support higher-precision multiply-accumulate operations, necessitating an ingenious online conversion strategy.

3.1 Why Abandon Offline FP8 Conversion?

A naive approach would be to offline convert all expert weights from MXFP4 to FP8 after downloading the model, saving them to disk, and directly loading FP8 weights during inference. However, this solution has fatal flaws:

• Double Storage Overhead: FP8 requires 1 byte, while MXFP4 requires only half a byte. A full conversion doubles the weight volume, placing immense pressure on local disk and memory bandwidth.
• Lack of Flexibility: If the model is fine-tuned or new experts are fused in the future, the offline conversion process cannot dynamically adapt.

3.2 Online Conversion and Temporal Overlap

The adaptation team ultimately adopted a “convert-on-demand” lazy loading strategy:

1. The MTT S5000 utilizes hardware decode instructions during expert weight loading to online execute mantissa bit-width expansion (4-bit to 8-bit) and synchronously calculate FP32-level scaling factors.
2. To achieve cross-layer reuse, the inference engine maintains a global FP8 workspace. When an expert in a certain layer is activated, its converted FP8 weights are written into this workspace on demand. When the computation for that layer is complete, the workspace is not immediately released but marked as “reusable” for subsequent layers (since expert layers generally exhibit local similarity) to continue using.
3. The scheduler uses a heuristic algorithm to preset a working “cache Time-To-Live” (TTL) to balance cache hit rates against memory footprint.

The advantages of this approach are twofold: it eliminates the need to pre-generate complete offline FP8 weights, reducing disk storage by approximately 50%; and the weight loading process runs completely in parallel with the matrix multiplication computation of the previous layer, achieving an efficient pipeline of “computing Layer A while loading Layer B.”

IV. Evaluation and Outlook: From Adaptation to Excellence

Upon completing the aforementioned adaptation work, the model achieved scores consistent with hardware-software co-optimization expectations on standard agent benchmarks (such as Terminal-Bench) and code generation-specific benchmarks, with all text and visual foundational inference chains operating stably.

This Day-0 adaptation practice validates three industry trends:

1. Hardware-Software Co-design Will Become the Core Competitive Barrier for AI Infrastructure: Hardware defines the capability boundary, while the software stack determines the actual utilization rate within that boundary. Possessing full-stack capabilities from compilers to operator libraries is essential to cope with the rapid iteration of model architectures.
2. Sparsification and Mixed Precision are the Necessary Paths to Longer Contexts: Given a constant cluster scale, sparse attention and 4-bit weight storage offer better cost-effectiveness than simply increasing memory bandwidth.
3. The Essence of Day-0 Capability is Abstraction and Decoupling: Abstracting the “unconventional operations” of complex models into universal computational paradigms (such as Gather-Compute-Scatter) and mapping them to hardware pipelines via compiler technology provides stronger long-term competitiveness than hardcoding adaptations for every single model.

For developers facing a new frontier LLM, Day-0 adaptation is no longer a question of “can it run,” but rather a challenge of “how to extract maximum performance within 72 hours.” The key to solving this puzzle lies in the hands of engineers who seamlessly master underlying compilers, operator libraries, and inference frameworks.