Kimi New Model Open Source: Day0 Inference Adaptation and Performance Optimization on Ascend

1. New Model Architecture Evolution and Adaptation Challenges

As context lengths of large models continue to grow, the computation and memory access overhead of attention mechanisms has become the primary bottleneck for inference performance. A recently open-sourced large model extends the architectural accumulation of its predecessor and introduces an interleaved design using two attention paths: one path processes sequence information through state recurrence, avoiding the explicit maintenance of a full historical key-value cache; the other path builds on Multi-head Latent Attention (MLA) and introduces an index selection mechanism, where a lightweight indexer first scores the relevance of historical positions and selects Top-K positions, after which a sparse attention kernel reads only the KV data of the selected positions to complete the computation. In addition, the model continues to strengthen deep-layer representation through a cross-layer attention residual aggregation mechanism.

This hybrid architecture places new demands on the inference engine. On one hand, the state-recurrence path requires efficient recursive operators and a state replay mechanism. On the other hand, the sparse attention path involves a multi-stage pipeline of index scoring, Top-K selection, and sparse reads; if each stage launches operators independently, it brings substantial intermediate result write-back and kernel launch overhead. For an inference platform, whether Day0 support can be completed at the moment a new model is released, and whether computation, communication, scheduling, and storage can be co-optimized, directly determines whether developers can quickly validate and deploy on their own compute.

2. Reusing Prior-Generation Optimizations for Rapid Architecture Adaptation

Adapting the new model did not start from scratch. Leveraging the engineering accumulation of the predecessor model in inference optimization, this support reused several mature techniques: the state-recurrence path reuses existing fused operator implementations, merging recursive computation and gating logic to reduce kernel launches; the two-stage fusion of cross-layer attention residuals continues to be used, lowering the memory access overhead of inter-layer aggregation; and grouped operators and communication-computation overlap schemes related to the mixture-of-experts structure are also directly inherited, avoiding repeated tuning.

At the system level, the adaptation combines hybrid parallelism strategies, communication optimization, and multi-stream orchestration: independent branches in the computation graph are split onto different CUDA streams (referring here to device-side parallel execution streams) for parallel scheduling, with event synchronization controlling dependencies so that memory-bound and compute-bound operations mask each other. Meanwhile, SuperKernel technology aggregates multiple small operators into a unified scheduling unit, and graph-mode execution reduces host-side dispatch overhead; a compilation cache mechanism avoids startup latency from repeated compilation; and weights adopt the NZ layout to match the hardware matrix unit’s data layout preference. For the state-recurrence path, speculative inference and a state replay mechanism are also retained, reducing serial waiting during the decoding phase through draft token prediction and fast state recovery. Together, this accumulation supports the new architecture in achieving usability adaptation within a short cycle.

3. Kernel Fusion and Multi-Stream Coordination for the Sparse Attention Path

Sparse attention is the core change in this architecture, and its execution flow can be decomposed into two stages: indexing and reading. The indexing stage needs to complete three things: relevance scoring between queries and keys, head-weight weighting, and Top-K position selection. In traditional implementations, these three steps often correspond to multiple independent operators, and intermediate results must be written back to global memory and then read by the next operator, which both increases bandwidth pressure and introduces multiple kernel launch latencies. The optimization fuses the three into a single kernel: it directly reads the key cache stored in pages, completes scoring, weighting, and selection within registers and shared memory, and finally outputs only the list of selected indices, significantly reducing intermediate result write-back and operator launches. The reading stage then reads data by position from the main attention KV cache according to the index list to complete the sparsified attention computation.

For quantized path support, both the preprocessing of the indexing stage and the decoding stage support BF16 and quantized paths simultaneously, so developers can flexibly choose according to accuracy requirements and hardware conditions. Noteworthy is the parallelization of the index pre-preparation stage: this stage includes mutually independent subtasks such as query projection, key normalization, query quantization, and head-weight projection. Through three-stream coordinated orchestration, the originally serial input preparation work is converted into parallel execution, effectively shortening the preparation latency before decoding.

4. Layered Quantization Strategy: Balancing Accuracy and Overhead

Quantization of the sparse attention path adopts a layered design rather than a one-size-fits-all approach for all steps. The projection computation retains BF16 precision to avoid quantization errors accumulating in the projection transformation and affecting the attention distribution. Queries and keys in the indexing stage are dynamically quantized in MXFP8 format, with scaling factors shared per group—32 elements per group share one scaling coefficient—compressing the storage and memory access overhead of index data to roughly half while preserving the relative accuracy of index scoring.

The KV cache uses FP8 format, with quantization granularity divided into three levels: per token, per head, and per block, with a block size of 128—that is, within each attention head of each token, independent quantization is performed in blocks of 128 elements. This granularity choice balances two conflicting goals: too coarse a granularity causes a single scaling factor to cover elements with large numerical differences, resulting in noticeable accuracy loss; too fine a granularity increases the storage and computation overhead of the scaling factors themselves. Partitioning by token-head-tile confines quantization error to a local range while maintaining hardware-friendly memory access alignment. Overall, layered quantization significantly reduces the storage footprint and bandwidth demand of the two types of frequently accessed data—index data and KV cache—providing direct benefits for cache capacity and throughput in long-context scenarios.

5. Context Parallelism and Long-Sequence Load Balancing

In long-input scenarios, single-card activation and intermediate tensor occupancy are key factors constraining sequence length. This support adopts a combined deployment of context parallelism and tensor parallelism, splitting long inputs along the sequence dimension so that each card processes only one segment of the sequence, and further shares intermediate tensors through sequence parallelism within the tensor parallel group. This combination reduces per-card memory usage while preserving the benefits of tensor parallelism for weight partitioning.

Causal index computation in sparse attention has an inherent load imbalance: positions later in the sequence need to retrieve a longer historical range, resulting in greater index computation. If segments are simply split contiguously, the later cards will bear significantly heavier computation, creating a pipeline bottleneck. To address this, the optimization adopts a Zigzag layout, assigning two paired segments from the head and tail of the sequence to the same card—the head segment has a short history and the tail segment a long history, so their computation loads compensate for each other, balancing the index computation load across cards; after computation, the sequence layout is restored in the original order to ensure the input order for subsequent layers is correct. This scheduling technique requires no change to model semantics and achieves cross-card load balance purely through task allocation strategy.

6. Ultra-Long Sequence Slicing for the State-Recurrence Path

The prefill computation time of the state-recurrence path grows approximately linearly with sequence length, directly affecting first-token latency in ultra-long sequence scenarios. When only tensor parallelism split by attention head is used, each card still needs to process the full sequence length, yielding limited splitting gains. To address this, the support scheme, based on an operator development framework implemented with a domain-specific language, adds context parallelism capability to this path: the input is split along the sequence dimension across multiple cards, each card handles only the recurrence computation for one segment, and cards are connected through state passing. This scheme aims to reduce first-token latency under ultra-long sequences while reducing per-card activation and intermediate tensor occupancy, complementing the context parallelism of the aforementioned sparse attention path, so that both attention paths can benefit from splitting in long-sequence scenarios.

7. Practical Recommendations and Summary

For developers planning to deploy this model on their own compute, it is recommended to proceed in the following order: first verify the availability of basic operators and the graph-mode execution path, and confirm that the weight layout and compilation cache configuration are correct; second, based on the actual business sequence length distribution, evaluate whether to enable the quantized path and context parallelism—short-sequence, high-concurrency scenarios can prioritize stream orchestration optimization in the decoding phase, while long-sequence scenarios should focus on configuring Zigzag load balancing and sequence slicing for the state-recurrence path; finally, combine speculative inference and the state replay mechanism to squeeze decoding-phase latency.

Overall, the core approach of this Day0 support can be summarized in three points: reuse mature optimization assets to shorten the adaptation cycle, perform kernel fusion and pipeline restructuring for the new architecture’s computational characteristics, and balance memory usage and load balancing in parallelism strategies. This methodology offers broadly applicable reference value for the inference deployment of similar sparse attention models—architectural innovation and engineering optimization must advance in tandem to truly translate a new model’s capabilities into inference services that developers can use.