Day-0 Adaptation for Large-Scale Sparse Models on Domestic AI Chips: Technical Pathways and Engineering Practices
As LLM parameters scale to trillions, hardware platforms face unprecedented challenges in memory, bandwidth, and adapting to sparse architectures. This article explores a domestic AI chip's successful Day-0 adaptation of a 3-trillion-parameter model. It details the software-hardware co-design strategies employed, including a decoupled software stack, optimized computation paths for sparse attention, and on-the-fly MXFP4 weight conversion to balance precision and performance.
Background: New Compute Challenges in LLM Inference
As the parameter scale of Large Language Models (LLMs) surpasses the trillion level, the computational complexity and resource consumption during the inference phase exhibit exponential growth. A 3-trillion-parameter model not only demands higher peak compute but also introduces systemic challenges in memory capacity, memory bandwidth, and communication efficiency. Notably, these ultra-large-scale models typically adopt a Mixture of Experts (MoE) architecture combined with sparse attention mechanisms. Their computational paradigm shifts from traditional dense matrix multiplication to a new model of “sparse activation + dynamic routing.”
For AI chip vendors, the ability to quickly adapt to new models directly determines their competitive position in the market. “Day-0 adaptation” refers to the capability where, upon the public release of new model weights, the hardware platform can immediately complete model loading, service initialization, and basic inference pipeline validation, achieving an “out-of-the-box” experience. This capability is a comprehensive reflection of software stack maturity, operator library richness, compiler optimization capabilities, and the engineering team’s response speed.
Hardware Platform and Software Stack Co-Design
Facing models with hundreds of billions or even trillions of parameters, the compute and memory of a single chip are insufficient, necessitating parallel computing across multi-card clusters. The adaptation case analyzed in this article is based on a domestic AI training and inference accelerator card and its accompanying software stack. The architectural design of this system follows the principles of “layered decoupling and ecosystem compatibility”:
- Low-level hardware layer: Provides high-bandwidth memory, high-speed interconnect interfaces, and compute units optimized for Transformer-series operators.
- Intermediate compiler layer: Includes a high-performance operator library, multiple programming entry points (such as Triton-like and TileLang-like frontend languages), and a low-level code generation compiler. It is responsible for efficiently mapping high-level model code to low-level hardware instructions.
- Upper framework adaptation layer: Interfaces with mainstream inference frameworks (like SGLang-like solutions) to ingest open-source community model code, thereby reducing model migration costs.
The key advantage of this layered architecture is that model developers do not need to understand low-level hardware details; they only need to follow standardized programming interfaces to complete model deployment. Meanwhile, chip vendors can continuously optimize the compiler and operator libraries to provide rapid support for next-generation model architectures.
Reconstruction of the Sparse Attention Computation Path
Ultra-large-scale models like Kimi-K3.1 Pro introduce a sparse attention mechanism (DSA). Its core idea is to break the quadratic computational complexity of traditional attention mechanisms—where each token must attend to all other tokens in the sequence—and instead selectively attend to a subset of key tokens based on specific strategies. While this mechanism significantly reduces computational overhead for long sequences (e.g., a 1-million context window), it introduces new complexities for hardware adaptation.
Index Selection and KV Cache Management
The first step in sparse attention is determining which attention scores between tokens need to be computed. This involves an index selection process, where the model dynamically filters a subset of tokens to participate in the attention calculation based on the relevance scores of query and key vectors. At the hardware implementation level, this means the computation graph is no longer static and pre-plannable; it must be dynamically generated at runtime based on input data.
The engineering team’s specific approaches were:
- Index generation and storage: During the Prefill phase, attention indices are computed based on the input sequence. The results are stored in a dedicated memory area for reuse during the subsequent Decode phase, where tokens are generated one by one.
- KV cache layout optimization: Because sparse attention only accesses key-value pairs of certain historical tokens, the KV cache cannot use a simple contiguous memory layout. The team designed a blocked storage scheme, rearranging the KV cache according to index clustering. This ensures better locality during Decode phase access, reducing cache miss rates.
- Prefill and Decode path separation: The Prefill phase processes the entire input sequence and is highly compute-intensive, while the Decode phase generates tokens sequentially and is highly memory-access intensive. Given the vast difference in computation modes between the two phases, the team performed operator-level optimizations for each path separately to ensure both achieve peak hardware performance.
Low-Level Implementation of Sparse Attention Operators
At the specific operator implementation level, the team decomposed sparse attention into several basic operations: index clustering, blocked matrix multiplication, mask application, and result scatter write-back. These operations were integrated through dedicated sparse attention operators in the high-performance operator library. Specifically, cross-layer residual connections and visual attention modules were incorporated into a unified execution pipeline management framework, ensuring the coherence of the overall inference process and the global optimization of resource scheduling.
Online Conversion Strategy for MXFP4 Expert Weights
In LLM training, to reduce memory footprint, an increasing number of models adopt low-precision weight storage schemes. MXFP4 (4-bit floating-point format) is one such emerging quantization standard. It can compress weight volume to a quarter of FP16 while maintaining model accuracy. However, the dynamic range of a 4-bit floating-point number is limited. Directly using it for matrix multiplication would result in severe precision loss, so in actual computation, weights typically need to be converted to a higher precision (such as FP8 or FP16).
Engineering Wisdom to Avoid Offline Conversion
The traditional approach is to pre-convert all low-precision weights into high-precision versions, storing them as separate weight files. However, for a 3-trillion-parameter model, this implies a massive additional storage space and conversion time. The ingenuity of this case lies in:
- On-the-fly Conversion: During the process of loading expert weights from main memory to the compute unit, the format conversion from MXFP4 to FP8 is executed in real-time. This approach only requires adding a lightweight conversion module on the data transfer path, occupying no extra memory space.
- Cross-layer FP8 workspace reuse: Because the decoding logic for MXFP4 is relatively unified, and multiple expert layers may share the same weight distribution characteristics, the team designed a cross-layer reusable FP8 workspace. Different layers can reuse the same FP8 memory region when executing expert computations, avoiding the need to allocate separate converted weight buffers for each layer. This drastically reduces peak memory requirements, allowing more parameters to be accommodated within limited memory capacity.
Balancing Precision and Performance
Online conversion is not without cost. The conversion process itself consumes additional compute cycles, and quantization errors may be introduced during the restoration of low-precision weights. The team mitigated these issues through the following means:
- Custom hardware instructions were designed for the MXFP4 decoding logic, hiding the latency of the conversion operations within the data transfer process, achieving “zero-cost” conversion.
- A mixed-precision strategy was adopted on critical computation paths—using FP16 for precision-sensitive operators and FP8 for operators with higher fault tolerance, maximizing throughput while ensuring output quality.
Ecosystem Adaptation: Bridging Open-Source Models to Production Services
The true difficulty of Day-0 adaptation lies not in optimizing individual operators, but in seamlessly migrating the entire open-source model ecosystem’s code to the target hardware platform. This requires the software stack to have a high degree of “ecosystem compatibility”:
- Framework layer compatibility: By interfacing with mainstream LLM inference service frameworks, model developers can deploy and invoke using familiar APIs without rewriting model definitions.
- Automatic operator mapping: The compiler can automatically recognize operator patterns in the model computation graph and map them to deeply optimized local operator libraries. For newly emerging operators (such as sparse attention variants), high-performance implementations are written using domain-specific languages similar to Triton, and then optimized by the compiler into low-level instructions.
- Dynamic shape support: In LLM inference, input sequence lengths and batch sizes change dynamically. The software stack needs to efficiently handle memory allocation and computation scheduling issues caused by dynamic shapes at runtime, avoiding performance jitter from frequent memory reallocation.
In this adaptation practice, the team ingested the upstream model implementation through the inference framework’s extension interfaces and utilized multiple programming tools to collaboratively complete the migration of key computation paths for the new architecture. The pattern formed during this process—framework adaptation layer acquiring model code, compiler executing operator optimization, and operator library providing high-performance primitives—can be reused for the rapid adaptation of future new models, forming accumulative technical assets.
Performance Validation and Benchmark Analysis
Upon completing the adaptation, the team conducted systematic performance validation. The testing scenarios covered multimodal tasks such as text generation and image understanding, with key metrics including:
- Time To First Token (TTFT): Reflects the processing efficiency of the Prefill phase. At a 3-trillion-parameter scale, through the effective combination of sparse attention and parallel strategies, TTFT was controlled within an acceptable interactive latency range.
- Token generation rate: Reflects the throughput capability of the Decode phase. Thanks to FP8 workspace reuse and KV cache layout optimization, the generation rate increased several fold compared to the naive implementation (without sparse optimization).
- Memory footprint: Through the MXFP4 online conversion strategy, a single card can load a larger proportion of model parameters without using complex parallel strategies like Tensor Parallelism (TP) or Pipeline Parallelism (PP).
Furthermore, in Agent workload tests and code generation benchmarks, the adapted inference service maintained accuracy consistent with the original model, with no degradation in output quality caused by hardware adaptation.
Summary and Future Outlook
A new co-evolutionary relationship is forming between large models and AI chips—innovations in model architecture (sparse attention, MoE, low-precision quantization) create new demands on hardware, while the expansion of hardware capabilities feeds back into the enhancement of model capabilities. The adaptation case analyzed in this article demonstrates that:
- The software stack is the core hub for hardware deployment: The computational potential of a chip must be translated into actual usability through a mature compiler, operator libraries, and framework adaptation.
- Pre-processing mindset for new architectures: Rather than passively optimizing each operator, it is better to establish an adaptation methodology of “model architecture feature recognition -> computation pattern analysis -> targeted optimization,” thereby building the capability to handle unknown new models.
- Hardware and algorithm co-design: Approaches like MXFP4 online conversion and sparse attention KV cache management essentially integrate algorithmic characteristics into hardware design considerations. In the future, the depth of this software-hardware co-design will determine the upper limit of AI infrastructure.
As model scales continue to expand, inference optimization will no longer be limited to single-card internal computation optimization, but will extend towards multi-card collaboration, heterogeneous computing, and compute-in-memory architectures. It can be foreseen that compute platforms with rapid iteration capabilities and deep hardware co-design capabilities will become the critical infrastructure of the large model era.
