Engineering Practices for Day-0 Adaptation of Million-Token Context Models on Domestic AI Chips

Background: LLM Inference Enters the “Long Context + Sparse Architecture” Era

As large language models evolve toward million-token context windows and trillion-parameter scales, the quadratic complexity of traditional dense attention mechanisms has become an unsustainable computational bottleneck for inference systems. Approximate attention structures, represented by Sparse Attention, compress the computational volume from quadratic to nearly linear relative to the sequence length by allowing each query token to attend only to a select few key historical tokens. This has become a mainstream technical direction for supporting ultra-long context inference.

Concurrently, weight precision during model training and inference is shifting toward lower bits. Low-precision representations like MXFP4 (4-bit microscaling floating-point format) significantly reduce memory footprint and bandwidth pressure while maintaining model accuracy. However, the online loading of low-precision weights, format conversion, and coordinated adaptation with compute kernels impose unprecedented demands on an inference system’s software stack.

As mainstream models enter the open-source community on a “release-day open-source” cadence, the ability of AI chip vendors to complete adaptation immediately directly impacts their competitiveness in the AI infrastructure market. Day-0 adaptation refers to completing the entire pipeline—from model loading and service startup to the generation of the first token—within an extremely short timeframe after new model weights are publicly released. This is not merely a contest of engineering speed, but a comprehensive test of software stack maturity.

Overall Adaptation Strategy: Upstream Integration and a Three-Pronged Migration System

Faced with adaptation tasks for newly released models, a mature inference software stack typically adheres to the principles of “reusing upstream implementations, decoupling by layer, and heterogeneous parallelism.” This adaptation effort inherited this approach: the inference service layer continued to utilize a localized branch of the widely-adopted SGLang inference framework rather than building a proprietary inference stack from scratch. This maximized the reuse of mature community capabilities such as scheduling strategies, continuous batching, and prefix caching.

The real technical challenges were concentrated in two areas. First, the unique sparse attention implementation in the new model architecture—DSA (Deep Sparse Attention)—which dictates the entire organization of the KV Cache and the execution modes of both the Prefill and Decode stages. Second, the online processing solution for MXFP4 expert weights, which directly affects peak memory usage and the single-card throughput limit.

To address these challenges, the software stack adopted a three-pronged migration strategy:

1. Operator Library Route: High-frequency operators such as dense matrix multiplication (GEMM), LayerNorm, and RoPE were migrated to the proprietary high-performance operator library, MATE, ensuring efficient execution of the base computational paths.
2. Programmable Kernel Route: Leveraging two layers of programmable kernel frameworks—the Triton dialect and TileLang—rapid prototyping and optimization were performed for irregular computational logic in sparse attention, such as index selection and KV reordering.
3. Compiler Route: The underlying compiler, MTCC, was used for global compute graph optimization, including operator fusion, memory reuse, and loop transformations, eliminating the extra overhead of cross-kernel invocations.

These three routes complement each other, collectively transforming community model code into a highly efficient inference service on actual hardware. This layered strategy—where instruction-level optimization is handled by the operator library, irregular logic is managed by programmable kernels, and global optimization is delegated to the compiler—ensures the optimal engineering balance within a limited adaptation cycle.

Four Key Layers of Sparse Attention Kernel Adaptation

While sparse attention macroscopically reduces attention computation, its implementation details are far more complex than dense attention. In this adaptation, the implementation of DSA was broken down into four key layers:

1. GPU-Friendly Refactoring of Index Selection Logic

The core premise of DSA is that each query token performs attention calculations on only a small subset of historical tokens. Therefore, the first step is determining “which historical tokens deserve attention.” This is typically accomplished by an index prediction module learned within the model. On NVIDIA GPUs, this logic is usually implemented as a custom CUDA kernel, and when migrating to a new hardware platform, it must be rewritten to an implementation supported by programmable kernel frameworks.

The key engineering practice here is that index prediction often contains numerous data-dependent branch operations. Executing these element-by-element directly leads to severe thread divergence. By grouping index selections by row, using shared memory for intermediate buffering, and utilizing vectorized memory access instructions to batch process adjacent token index queries, the adaptation team improved the performance of this irregular computational path by approximately 2.3x. (Note: This is an internal measurement during the adaptation process and is provided for reference.)

2. Memory Rearrangement of KV Cache Layout

Sparse attention implies that the actual KV positions required by each query token are not contiguous and may even span distant sequence fragments. If the contiguous KV layout of dense attention is retained, the GPU will experience numerous cache misses and bandwidth waste during memory access.

The adaptation solution employed a block-sparse storage format: the KV cache is partitioned into fixed block sizes (e.g., 64 or 128 tokens), maintaining an independent metadata header for each block that records the number of valid tokens and the offset address within the block. Once the index selection module determines the attention scope, the inference engine only needs to establish a mapping between the query tokens and the target KV blocks. This enables efficient batch loading during the Prefill stage and on-demand, block-level incremental updates during the Decode stage.

Another advantage of this design is its support for dynamic release and reuse of the KV cache. When a sequence completes generation, its occupied KV blocks can be directly reclaimed by subsequent requests, avoiding the fragmentation issues inherent in traditional contiguous caches.

3. Asynchronous Pipeline Computation for Prefill and Decode

The Prefill stage processes the user’s entire prompt sequence and is compute-intensive with high parallelism. The Decode stage generates output tokens one by one and is memory-intensive with high single-step latency sensitivity. DSA exhibits fundamental differences in computational patterns between these two stages.

During the Prefill stage, the adapted sparse attention kernel uses a “query block × KV block” 2D parallel unit, assigning the attention computation for each query token group to a single thread block. It relies on inter-block parallelism to hide KV loading latency. Inside the thread block, a divide-and-conquer strategy is employed: partial attention scores between the current query block and all candidate KV blocks are calculated first, followed by a cross-block reduction to obtain the final result.

During the Decode stage, as only a single token is processed at a time, the computational volume is small but latency requirements are strict. The adaptation team merged the DSA kernel and the KV cache update operation into a single kernel launch, avoiding the synchronization overhead between two independent kernel launches. Simultaneously, asynchronous copy instructions are utilized to prefetch the next likely attended KV block while attention computation is ongoing, achieving complete overlap of memory access and computation.

4. Collaborative Integration with Other Complex Mechanisms

DSA does not exist in isolation; it must work in tandem with KDA state management, cross-layer residual connections, and visual Attention modules to form a complete execution pipeline. KDA state management is responsible for maintaining the KV cache state machine across different sparsity levels, dynamically adjusting block validity at runtime. Cross-layer residual connections require the output addition from the previous layer to be completed before the DSA kernel executes, which necessitates the compiler to ensure fused execution of the residual addition without adding extra memory round-trips. The visual Attention module shares the same KV storage pool with DSA but has an independent number of Heads and scaling factors. During adaptation, unified entry and dispatched execution were achieved through kernel template instantiation with different configuration parameters.

Engineering Implementation of Low-Precision Expert Weights

Beyond sparse attention, the efficient execution of MXFP4 expert weights was another major hurdle. MXFP4 is a microscaling format—different tokens use different scaling factors—so it cannot be uniformly scaled at the tensor level like traditional FP8.

Online Conversion and Cross-Layer Reuse Strategy

Instead of generating complete offline FP8 weights all at once during model loading, the adaptation team chose a technical route of “on-demand conversion, instantaneous reuse.” The specific implementation logic is as follows:

• The FP4 weights of each Expert are loaded into shared memory by a specialized rearrangement kernel upon their first invocation.
• Online unpacking and conversion from FP4 to FP8 are completed within shared memory.
• The converted FP8 weights are temporarily stored in a gigabyte-scale workspace for reuse across multiple calls within the same layer.
• When subsequent layers’ expert computations require these weights, the FP8 version is read directly, avoiding redundant conversions.

The theoretical basis for this design is that there is high similarity in weight distribution across layers. Cross-layer reuse can reduce the actual number of conversions to roughly a quarter of the original amount while eliminating the extra storage space required for offline full FP8 weight generation. Furthermore, the online conversion strategy ensures that only one FP8 working copy is retained in memory at any given time, preventing memory bloat caused by the coexistence of both FP4 and FP8 weights.

Trade-offs Between Precision and Throughput

A primary potential risk of online conversion is that conversion latency appears directly on the request hot path. The adaptation team mitigated this in two ways. First, the conversion kernel is pipelined and overlapped with the matrix multiplication of the preceding layer; while the previous layer executes its GEMM, the weights for the next layer are already being converted in shared memory, ready for consumption. Second, for hot experts, call frequency statistics were introduced. A static pre-conversion strategy is applied to the most frequently used experts, ensuring that the most common weights in the task scheduling are always ready in FP8 format.

Adaptation Results and Core Benefits

Following the aforementioned adaptations, the final results can be summarized in five points:

• End-to-End Pipeline Unlocked: Model loading, service startup, text inference, and image inference can all be executed on the same hardware platform.
• High Accuracy Maintained: Online conversion accuracy matches that of offline conversion, with no observable degradation in model evaluation scores.
• Effective Memory Control: The cross-layer reuse strategy prevents the generation of full FP8 weight replicas, significantly increasing the number of concurrent requests a single card can handle.
• High Maintainability: Kernel code written using the Triton dialect and TileLang offers excellent readability and modifiability, eliminating the need to rewrite core logic during future model version iterations.
• Methodology Accumulation: The sparse attention kernel adaptation templates and low-precision weight conversion pipelines developed during this adaptation have been consolidated into a general toolset, supporting rapid adaptation for future new models.

Summary and Outlook

This Day-0 adaptation practice demonstrates that in the face of increasingly complex model architectures and low-precision computation trends, the competitiveness of AI chips depends no longer solely on peak hardware compute capacity, but rather on the efficiency of the software stack and the rapid response capabilities of the engineering team. Adapting sparse attention is by no means a simple replacement of a few operators; it is a systems engineering challenge involving index logic, storage layout, scheduling strategies, and compiler coordination.

Looking ahead, as model parameters continue to scale and context windows keep expanding, inference systems must maintain sustained investment in the following directions: adaptive KV cache compression, more flexible sparsity patterns (such as query-dependent dynamic sparsity), and better overlap of multi-expert parallel communication and computation. Only by deepening the collaborative capabilities across all levels of the software stack can AI chip vendors secure an active position in this rapidly evolving industry landscape.