Engineering Day-0 Adaptation for Trillion-Scale Models on Domestic GPUs: From Operator Migration to Inference Pipelines

Background: The Challenge of New Computing Paradigms in Large Model Deployment

As the parameter scale of large language models enters the trillion-level, the logic of model deployment has evolved from simple “load-and-infer” to a deep reconstruction of underlying computing architectures, memory hierarchies, and operator implementations. Taking the recently released K-series model as an example, it not only possesses over three trillion parameters but also introduces frontier designs such as Sparse Attention (DSA), Mixture of Experts (MoE), and low-bit weights (MXFP4). These fundamental changes in model structure imply that the traditional GPU computing paradigm—reliant on dense matrix multiplication—can no longer meet the requirements, necessitating deep customization of the computing stack.

The core concept of Day-0 adaptation lies in the hardware and software stack’s readiness to support the model the moment its weights are officially released to the community. This is not merely about getting a PyTorch model to run; it requires the inference service to reach a functional state on the specific computational paths of the model architecture, including but not limited to establishing pipelines for multimodal tasks like text generation and visual understanding. In this process, heterogeneous computing platforms face not only a competition of raw compute power but also a direct test of software ecosystem maturity.

Overall Technical Architecture for Heterogeneous Platform Adaptation

To address these challenges, our technical solution adopts a layered adaptation architecture. This architecture is designed around the collaboration of four layers: hardware, operator library, compiler, and inference framework, ensuring a seamless transition from upstream community models to actual inference services.

Model Access Layer: Bridging the Gap via Open-Source Inference Frameworks

At the model access level, we adopted a key strategy: adapting a specific version branch of the high-performance inference framework, SGLang. This decision was crucial because it allowed us to leverage the community’s mature advancements in continuous batching, prefix caching, and RadixAttention. The core of the adaptation work was not to modify the model’s forward pass code, but to accurately translate the scheduling instructions of the upper-layer framework into tasks recognizable by the underlying compute kernels.

In this layer, we primarily resolved model structure parsing and graph compilation issues. For the K-series model’s specific MoE architecture, we needed to reorganize the gating network logic and expert parallelism strategy. This ensures that when an inference request arrives, tokens are routed only to the relevant expert sub-networks according to sparse activation rules, thereby significantly reducing invalid computation.

Operator Implementation Layer: Application and Trade-offs of Three Toolchains

The core goal of the operator layer is to solve the “new architecture requires new operators” problem. Based on a fully proprietary software stack, we employed differentiated implementation strategies for various operator types:

1. MATE Operator Library: For the most computationally dense and frequently called operators in the model, such as standard matrix multiplication and LayerNorm, we directly invoke a deeply optimized encapsulated library. Its advantage lies in assembly-level optimization, fully utilizing the underlying hardware’s tensor cores and register files.
2. Triton-MUSA and TileLang-MUSA: For approximately 20% of the customized operators in the model (such as various non-linear activation functions or specific normalizations), we utilize a Triton-like Domain Specific Language (DSL) for programming. This drastically improves development efficiency, allowing algorithm engineers to implement highly efficient kernels without diving deep into the underlying CUDA. For instance, when implementing the index selection logic for DSA, we leveraged TileLang’s flexibility in block-level control to achieve efficient traversal of the sparse block table.
3. MTCC Compiler: This compiler is specifically responsible for automatically mapping high-level computational graphs to underlying instructions and statically optimizing memory access patterns. Particularly when handling inference scenarios involving dynamic shapes, MTCC can perform online auto-tuning to select the optimal loop unrolling factors and vectorization widths.

Key Technical Breakthrough I: Lossless Migration of Sparse Attention

The K-series model utilizes Sparse Attention (DSA) as the core mechanism to support ultra-long contexts (at the 1-million token level). Unlike traditional global attention, DSA allows the model to focus only on a selected subset of tokens in the sequence, reducing computational complexity from O(n²) to O(n). However, this mechanism poses severe challenges to the inference framework’s index management.

During our adaptation practice, we achieved the integration of three critical paths:

• Index Selection and KV Cache Layout Reordering: DSA requires the inference framework to dynamically maintain a sparse index table to determine which historical tokens are visible to the current query token. We reconstructed the KV cache data structure within the MUSA software stack, changing the contiguous memory layout to a blocked sparse layout. By embedding the index selection algorithm directly within the operator, we avoided frequent copying of index data between the GPU host and device, significantly reducing scheduling latency.
• Divide-and-Conquer for Prefill and Decode Paths: In the Prefill phase, since a large number of prompt tokens need to be processed in parallel, we adopted a sparse block table-based matrix multiplication approximation algorithm. In the Decode phase, as only a single token needs to be computed per step, we switched to a vectorized Gather operator to aggregate context vectors from the KV cache by index. This divide-and-conquer strategy ensures peak hardware floating-point utilization across different computational stages.
• Synergy with KDA State Management: DSA also involves the fusion of cross-layer residual connections and visual attention modules. We fused the output of sparse attention with the KDA (Key-Value Delta Attention) state. Through a custom fused kernel, we completed interval updates and residual accumulation in a single kernel launch, avoiding redundant global memory reads and writes.

Key Technical Breakthrough II: Online Conversion Strategy for Low-Bit Expert Weights

This model employs MXFP4 (4-bit microscaling floating-point) approximate precision storage in the MoE layers. While loading the model directly in FP4 format into VRAM maximizes parameter capacity, it leads to a loss of computational precision, especially as errors accumulate progressively during long-sequence inference decoding.

To solve this problem, we abandoned the traditional approach of “offline full conversion” and instead designed an online conversion + workspace reuse mechanism:

1. Online Conversion Pipeline: During the model loading phase, only non-expert shared weights (such as the Attention and gating layers) are directly converted to FP16. For expert weights, their FP4 storage format is maintained. During inference, when a specific expert is activated by the gating network, its weight blocks are loaded in real-time into a dedicated decoder preceding the compute unit. This decoder decompresses them into BF16 (BFloat16) floating-point format for computation within a very short latency (a few microseconds).
2. FP8 Workspace Cross-Layer Reuse: The decompression operation itself is highly resource-intensive. To reduce the decompression frequency, we allocated a global FP8 workspace decoupled from the number of expert layers. When an expert layer finishes computing, this workspace is not immediately released; instead, via a cross-layer cache management strategy, it checks if the next layer has a possibility of activating the same expert. If there is a hit, the decompressed weights are directly reused, thereby nearly halving the conversion compute overhead between adjacent layers.
3. Numerical Stability Assurance: To ensure that online conversion does not introduce precision degradation, we additionally designed a zero-overhead scaling factor correction table. This table is bound to the KV cache as a sparse tensor and calibrated once during each Prefill phase, ensuring that numerical drift during long-context inference remains within a controllable range.

Engineering Results and Performance Observations

Following the deep adaptation of the aforementioned tech stack, we successfully completed the entire pipeline from model loading and service startup to basic text and image inference. In actual stress testing, the platform consistently output compliant results on conversational tasks with a 1-million token context length, as well as multimodal understanding tasks involving a single image.

It is worth noting that the significance of Day-0 adaptation is not a breakthrough in a specific benchmark score, but rather a validation of the software stack’s universal承接 capability. Previously, adapting to a new model architecture took months; currently, through the combination of toolchains and the accumulation of operator templates, this cycle has been shortened to a matter of weeks. This lays an engineering methodological foundation for tackling future models with even larger parameter scales and more complex attention variants.

Conclusion and Future Outlook: From Adaptation to Native Development

This practice validates that the domestic heterogeneous platform possesses the technical prerequisites to serve the latest generation of frontier large models. From the fine-grained adaptation of DSA operators to the dynamic decompression of MXFP4 low-bit weights, every step requires deep vertical integration at the compiler, runtime, and operator library levels.

Future evolution will focus on “reverse adaptation” capabilities—meaning that rather than just making new models fit existing hardware, hardware design phases will proactively consider synergy with sparse computing and dynamic reconfiguration of low-bit precision. At that point, large model inference will no longer be a “compromise” for hardware, but will become a true “native application.”