Engineering Practices and Optimizations for Deploying Trillion-Parameter Sparse LLMs

1. Background: Deployment Challenges from Ultra-Long Context and Sparse Models

In recent years, the parameter scale and context length of large language models have grown rapidly in tandem. When model parameters reach the trillion level while supporting million-token context windows, the pressure on inference systems is no longer just “can it compute” but a combination of problems: slow loading, slow first-token latency, difficult memory planning, and complex computation paths. In particular, the introduction of sparse attention mechanisms and mixed-precision expert weights means that inference engines must be specifically adapted to the model’s new structural characteristics rather than relying solely on traditional dense Transformer execution paths.

Against this backdrop, “Day-0 adaptation” has become an important indicator of the maturity of an inference software stack. Day-0 adaptation refers to completing the full verification flow—from model loading and computation graph construction to basic inference service startup—within an extremely short time after model weights are publicly released. For very large models, this goal is not easy to achieve; it requires the underlying software stack to be well prepared in advance across operator coverage, memory management, compiler coordination, and distributed execution.

2. Adaptation Goals and Overall Approach

The core goal of this adaptation effort was to complete the inference pipeline for a sparse model with over ten trillion parameters on a domestic AI training-and-inference accelerator card and its accompanying software stack. Specifically, this included:

1. Correct loading and initialization of model weights;
2. Correct execution of sparse attention computation paths;
3. Online conversion and efficient computation of low-precision expert weights;
4. Validation of basic inference capability for both text and image input modalities;
5. Stable operation of the full pipeline from service startup to first token return.

Given the model’s complex structure and enormous parameter scale, the adaptation work could not rely on layer-by-layer hardcoding. Instead, it started from the layered structure of the software stack, clarified the responsibilities of each layer, reused existing capabilities, and made targeted extensions only at key points of difference. Overall, the adaptation approach can be summarized as follows: the upper inference framework handles model structure and scheduling logic; the middle compilation and operator layers handle efficient implementation of compute-intensive operations; and the underlying device software handles memory allocation, data movement, and execution synchronization.

3. Model Integration on the Inference Framework Side: The SGLang Compatibility Layer Example

At the inference framework level, the adaptation work first required sorting out the model’s structure definition, weight mapping, and forward computation flow. Among commonly used inference serving frameworks in the community, SGLang has become the preferred foundation for many teams due to its flexible support for various model structures. During adaptation, the team used the framework’s compatibility layer to host the upstream model implementation, avoiding the need to write complete serving logic from scratch.

Specifically, framework-side adaptation included the following:

• Configuration parsing: Reading key parameters from the model configuration—such as hidden dimension, number of layers, number of attention heads, sparse attention window size, number of experts, and activation strategy—to ensure accurate subsequent computation graph construction.
• Weight mapping: Mapping tensor names and shapes from model weight files to the inference engine’s internal representation, handling sharded weights, inter-layer reused weights, and special parameter formats.
• Computation graph construction: Generating the forward computation graph based on model structure, explicitly distinguishing dense computation paths from sparse computation paths, and inserting custom operators at appropriate levels.
• Modality branch handling: For models supporting both text and image inputs, a visual encoding branch was added at the input processing stage, with visual features and text features correctly concatenated before being fed into the backbone network.

The key to framework-side adaptation is “hosting” rather than “rewriting.” By reusing the framework’s existing scheduling, batching, sampling, and streaming output capabilities, the team could focus on the model-specific computation paths, significantly shortening the adaptation cycle.

4. Compilation and Operator Layer Coordination: Completing the MUSA Software Stack

In the underlying software stack, going from high-level operator descriptions to device-executable code requires multiple layers of compilation and optimization. For special computation patterns in new models, general-purpose operator libraries often cannot directly cover the requirements, necessitating coordination across multiple tools.

During adaptation, the following categories of components were primarily relied upon:

1. Operator libraries: Providing hand-optimized high-performance computation primitives, including matrix multiplication, attention computation, normalization, and activation functions. For standard computations in the model, directly calling the operator library achieved good performance.

2. Triton-like language layer: Allowing developers to write custom operators using Python-like syntax, with the compiler automatically generating device code. This approach is suitable for implementing non-standard but relatively regular computation patterns, striking a balance between development efficiency and execution efficiency.

3. TileLang-like tensor programming layer: Oriented toward finer-grained tensor operation descriptions, suitable for expressing complex data layout transformations and on-chip data reuse strategies. When adapting sparse attention, this type of tool could be used to describe index selection and KV cache reordering logic.

4. Device compiler: Responsible for compiling intermediate representations into instruction sequences executable on the accelerator card, performing register allocation, instruction scheduling, and memory access optimization. For computations involving special data precision or custom data paths, the compiler needed to provide corresponding support.

The benefit of this multi-layer coordination is that most conventional computations in the model can directly reuse mature operators, while a small number of special computations can be quickly implemented through the programming layer without waiting for low-level library updates. This “long-tail completion” approach is a key engineering technique for rapid adaptation of large models.

5. Core Issues in Sparse Attention Adaptation

The sparse attention mechanism (commonly referred to as DSA) employed by Kimi-K3.1 Pro-class models was the focus and difficulty of this adaptation effort. Unlike standard multi-head attention, sparse attention computes only a subset of KV positions, thereby greatly reducing computation and KV cache usage in ultra-long context scenarios. However, this sparsity also brings three core issues to the inference system: how indices are selected, how the KV cache is laid out, and how the computation paths for Prefill and Decode stages are unified.

5.1 Index Selection

Sparse attention requires determining the set of KV positions that each query position attends to. Different models may adopt different selection strategies, including fixed stride, sliding window plus global tokens, and score-based dynamic selection. The inference engine needs to efficiently translate the index information provided by the model into data access patterns during computation.

In this adaptation, the index selection logic was encapsulated in an independent preprocessing stage. For each attention head, the KV position indices requiring computation were calculated based on the model configuration and current sequence length. These indices were then used to guide KV cache reads and attention score computation. The correctness of index computation directly determines the validity of model output, so strict alignment verification with the reference implementation was necessary.

5.2 KV Cache Layout

Traditional dense attention typically stores the KV cache contiguously by layer, by head, and by sequence position. However, with sparse attention, each query accesses only a subset of KV positions. If a contiguous layout is still used, memory accesses become scattered and bandwidth utilization drops.

The approach adopted in this adaptation was: after index selection is completed, reorganize the KV cache. Specifically, this included:

• Compactly arranging KV blocks that are frequently accessed;
• Establishing independent cache regions for global tokens and local window tokens;
• Maintaining independent offset tables for each attention head to record sparse access patterns.

Through this layout adjustment, the computation kernel can access KV data in larger contiguous blocks when reading, reducing performance loss from random access.

5.3 Segmented Processing of Prefill and Decode

The Prefill stage requires a one-time attention computation over the entire input sequence. Since the sequence length can be very large, even with sparse attention, the Prefill computation remains significant. In this adaptation, Prefill was split into multiple chunks for processing, with dense computation kernels used within each chunk and sparse indices used to determine dependencies between chunks.

The Decode stage processes only one newly generated token at a time. At this point, the query count is 1, but relevant positions must be selected from a massive KV cache. For sparse attention, the optimization focus in the Decode stage is reducing index lookup overhead and avoiding reading too much invalid data for a single query.

By connecting the sparse computation paths of Prefill and Decode to dedicated operator implementations and handling stage switching properly at the framework level, acceptable inference performance was achieved while ensuring correctness.

6. Online Conversion and Reuse Strategy for Low-Precision Expert Weights

Ultra-large-scale models typically adopt a mixture-of-experts architecture, with expert weights stored in low-precision formats (such as MXFP4) to save memory and storage space. However, most current computation kernels cannot yet perform efficient operations directly on MXFP4 format. The common practice is to convert to FP8 or higher precision before computation.

If an offline conversion approach were used—converting all expert weights to FP8 format in advance and saving them—two problems would arise: first, the converted weight files would be substantially larger, increasing storage and loading pressure; second, for models with parameter counts reaching the trillion level, offline conversion itself would require significant time and computational resources, conflicting with the rapid Day-0 adaptation goal.

Therefore, an online conversion scheme was adopted in this adaptation: during the inference service startup phase, only the original low-precision weights are loaded; before actual expert computation is executed, the required expert weight segments are converted on the fly to FP8 format and placed into a pre-allocated workspace for computation. Additionally, considering that the same experts may be activated across different layers and different tokens, the system implemented a cross-layer reuse mechanism—expert weights that have already been converted are retained in the workspace, and when the same expert is encountered subsequently, it is directly reused, avoiding repeated conversion overhead.

The key to this scheme lies in workspace capacity management. Due to limited memory, it is not possible to keep all converted expert weights resident in memory, so a reasonable cache eviction strategy must be designed. In this adaptation, the system dynamically adjusted workspace contents based on expert activation frequency, prioritizing high-frequency experts. This strategy effectively controlled memory usage while ensuring computation correctness.

7. Integration of the Vision Modality and Cross-Layer Residual Pipeline

For models supporting multimodal input, the inference system must handle not only text tokens but also image inputs. In this adaptation, the visual encoding portion was encapsulated as a separate module that receives image data and outputs a visual feature sequence. This sequence is concatenated with text token embeddings and fed uniformly into the backbone network.

It is worth noting that visual attention and text sparse attention differ structurally. Visual attention typically uses dense computation because the number of image patches is relatively manageable. Therefore, in the computation graph, the visual encoding branch and the text sparse attention branch were clearly separated, calling different computation kernels, and finally merged at the feature dimension.

Furthermore, the cross-layer residual structure in the model requires the inference engine to correctly pass and accumulate hidden states between layers. For special layers containing residual connections, the system needed to maintain residual states in memory and update them correctly after forward propagation completes. Although this logic appears simple, synchronization errors can easily occur in distributed execution and pipeline parallelism scenarios. In this adaptation, a unified state management module was used to handle this centrally.

8. Debugging, Validation, and Performance Tuning Experience

After completing functional adaptation, systematic debugging and performance tuning were still necessary. Here are several lessons from practice:

Layer-by-layer alignment verification: Split the model by layer and compare the output of the adapted implementation against the reference implementation layer by layer. First verify that the numerical error of a single layer’s output is within tolerance, then gradually chain multiple layers together, and finally verify end-to-end output. This hierarchical localization method can quickly pinpoint the source of problems.

Memory peak monitoring: Continuously monitor memory usage during service startup and inference, identifying memory peaks in the loading stage, Prefill stage, and Decode stage. For cases exceeding limits, resolve by adjusting workspace size, enabling recomputation, or optimizing cache layout.

First-token latency optimization: For very large models, first-token latency is a key metric for user experience. By optimizing weight loading order, prefetching high-frequency experts, and reducing synchronization waits during startup, the time from request to first token return can be significantly shortened.

Edge-case testing: In addition to standard inputs, test boundary scenarios such as ultra-long input, empty input, and ultra-long generation to ensure the system does not crash or produce erroneous output under abnormal conditions.

9. Summary and Future Directions

This adaptation practice demonstrates that for sparse large models at the trillion-parameter level, coordination across the inference framework, compiler, and operator libraries can achieve full pipeline completion from model loading to basic inference service startup in a relatively short time. Core lessons include: hosting rather than rewriting at the framework layer, long-tail completion at the operator layer, dedicated handling of sparse attention, online conversion and reuse of low-precision weights, and independent processing and integration of multimodal branches.

Future work can proceed in the following directions: further optimizing the computational efficiency of sparse attention and exploring more refined KV cache compression strategies; improving cache scheduling algorithms for low-precision online conversion to reduce conversion overhead; enhancing batching capability in multimodal input scenarios; and strengthening memory and communication coordination optimization in distributed multi-card scenarios.

For builders of large model inference systems, establishing an adaptation methodology that can quickly respond to new model structures has more long-term value than deep optimization for a single model. As models continue to evolve, a software stack with a well-designed layered architecture and flexible operator extension capabilities will become the foundational guarantee for efficiently delivering inference services.