Inference Optimization for Kimi K3.1 Pro: Decoding Architecture Evolution and Systems Engineering Practices

Introduction: From Architectural Innovation to Inference Performance Challenges

As large language models continue to scale in both parameter count and context length, optimizing inference efficiency has become a critical factor in determining user experience and deployment costs. The recently open-sourced Kimi K3.1 Pro model incorporates two sparse attention mechanisms at the architectural level: KDA (Kalman-style Dynamic Attention), which processes sequences through state recursion, and DSA (Dynamic Sparse Attention), which introduces index selection on top of MLA (Multi-head Latent Attention). While this hybrid architecture significantly enhances long-sequence modeling capabilities, it also imposes more complex scheduling and optimization requirements on the underlying compute stack.

Based on the architectural characteristics of this model, this article systematically outlines the key optimization technologies applied during the inference adaptation process. Rather than focusing on commercial promotion, we concentrate on the technology itself: how to reduce memory movement through operator fusion, utilize multi-stream parallelism to hide memory access latency, balance precision and throughput via tiered quantization, and accelerate the generation phase by adapting speculative decoding mechanisms. These optimization approaches are highly versatile and offer valuable reference for any developers or architects facing inference challenges related to long-sequence sparse attention.

1. Algorithmic Architecture Analysis: The Collaborative Compute Flow of KDA and DSA

A prerequisite for understanding inference optimization is accurately grasping the model’s computational characteristics. The core innovation of Kimi K3.1 Pro lies in the interlaced stacking of two different paradigms of attention mechanisms.

KDA (State Recurrent Attention) is responsible for handling the global dependencies of the sequence. Drawing inspiration from classic Kalman filtering or state space models, it maintains a hidden state vector and scans the sequence token by token in a recursive manner, thereby capturing long-distance information with linear complexity. During the Prefill phase, this mechanism operates as sequential state updates, while in the Decode phase, it acts as a single-step state synchronization. This compute pattern is highly regular and friendly to matrix multiplication (GEMM).

DSA (Dynamic Sparse Attention) focuses on the precise extraction of local key information. Its operational flow can be broken down into three steps:

1. Lightning Indexer (LI): Scores the relevance between the current query and historical keys. This scoring process is not a simple dot product but includes weighting of attention head weights.
2. Top-K Selection: LI filters the top K highest-scoring historical positions based on the scoring results. These K positions serve as the identified “important” context anchors.
3. Sparse Read and Compute: SFA (Sparse Flash Attention) directly reads the corresponding Key-Value (KV) cache based on the indices selected by LI and executes sparse attention computation.

Additionally, the model employs an Attention Residuals mechanism, aggregating attention features across different layers. This requires the inference framework to efficiently manage cross-layer data dependencies.

Implications for the Inference Engine: This hybrid architecture means the computation graph is not a single dense matrix flow, but a composite of “dense GEMM + sparse indexing + gather operations + sparse GEMM.” Optimization efforts must focus on fusing the sparse index computation (LI) with the sparse read (SFA), as well as interleaving the asynchronous execution of DSA with the state recursion of KDA.

2. Operator-Level Deep Fusion: From “Three-Stage Sequential Execution” to “Integrated Execution”

In traditional inference implementations, the LI computation typically involves three independent operators: calculating relevance scores, applying head weights, and executing Top-K selection. After each operator finishes, intermediate results must be written back to device memory. This not only increases memory bandwidth pressure but also introduces the overhead of multiple kernel launches.

Optimization Strategy: Single Fused Kernel
We designed a fused execution scheme that integrates the three logical steps above into a single kernel function:

• Data Flow Cohesion: After the relevance scores are calculated, they reside directly in on-chip caches (e.g., SRAM/L2 Cache). The head weight weighting operation follows immediately, and the weighted scores do not need to fall back to global memory.
• Inlined Top-K Selection: After weighting is complete, the Top-K selection algorithm is executed directly at the register or shared memory level, outputting only the final selected indices. Key Technique: Once the indices are output, we use them to directly read the Paged KV Cache, rather than reading full contiguous cache blocks. This “on-demand in-page addressing” resolves the index-to-data mapping at a lower level, reducing the extra overhead caused by address translation.

The direct benefit of this fusion is the complete elimination of global memory reads and writes for intermediate tensors, reducing kernel launch counts from three to one. In long-sequence scenarios, since the K value in Top-K is much smaller than the sequence length (e.g., K=128), the theoretical speedup ratio brought by sparse computation can be more authentically translated into an actual end-to-end latency reduction, without being diluted by memory movement.

3. Pipeline Orchestration: Indexer Prolog and Multi-Stream Collaborative Scheduling

The decoding process (Decode) of large models generates tokens sequentially. Each step involves a serial chain: “read input -> compute attention -> compute feed-forward network (MLP) -> output.” Any环节 without parallel optimization will directly expose the system to latency.

Bottleneck Identification: In the DSA architecture, input preparation (Q projection, normalization, quantization) is the prerequisite for all subsequent computations. If these operations are strictly executed in the order of “project, then normalize, then quantize,” it forms a long-tail latency.

Innovative Optimization: Three-Stream Parallel Architecture
To achieve zero-wait, we introduced the “Indexer Prolog” pipeline stage, splitting the input preparation phase into three micro-streams executing in parallel:

• Stream A (Main Path Stream): Executes the Q matrix projection (GEMM operation).
• Stream B (Normalization Stream): Directly executes K normalization on the input activations; this operation does not depend on the Q projection results.
• Stream C (Quantization and Index Stream): Executes Q quantization and projects it with pre-calculated head weights.

Using multi-stream technology, these three streams are assigned to different compute queues on the device. The launch times of Stream B and Stream C are deliberately advanced to overlap execution with Stream A. When Stream A completes the Q projection, the normalization results from Stream B and the quantization results from Stream C are already prepared. The three merge and immediately enter the LI computation phase.

The essence of this collaborative scheduling is using the spatial parallelism of computing resources in exchange for a reduction in temporal serialization. It effectively fills the pipeline bubbles most likely to occur during the decoding process, thereby improving overall throughput (Tokens/s).

4. Tiered Precision Control: The Art of Balancing Quantization Granularity and Numerical Stability

To reduce the memory footprint and搬运 overhead of KV cache in long-sequence scenarios, compressing (quantizing) the cache data is necessary. However, excessive quantization can lead to model precision loss. The solution provided by the Kimi model is a combination of “tiered quantization” and “dynamic group granularity.”

Fully Quantized Path for LI Indexer: MXFP8 Format
The purpose of the indexer’s computation is to select Top-K positions, not to accurately reconstruct the attention matrix. Therefore, we use the MXFP8 (Microscaling FP8) format to dynamically quantize Q and K during the indexing process. The key parameter here is Group Size=32, meaning that every 32 consecutive elements share a single scale factor. Choosing 32 over a smaller value strikes a compromise between quantization precision and the storage overhead of scale factors. Since this is dynamic quantization, the calculation of scale factors happens during inference, requiring the inference kernel to possess highly efficient Reduce operation capabilities.

SFA KV Cache Quantization: per-token-head-tile
For the KV cache that actually participates in the attention weighted sum, we retain BF16 for the core matrix multiplication to maintain precision, but the cache itself is stored as FP8. The quantization granularity is per-token-head-tile, with the Tile size set to 128. The brilliance of this design lies in:

• per-token and per-head: Ensures that each token has an independent quantization range within a single attention head, avoiding the amplification of outlier values caused by cross-head or cross-token statistics.
• Tile of 128: This is a size friendly to Matrix Multiply Units (MMU). In matrix computations on AI processors, loading and computation are typically performed in units of specific Tile sizes (e.g., 16x16 or similar). A quantization granularity of 128 allows the de-quantization operation to seamlessly integrate with the block-loaded matrix. The de-quantized BF16 data can directly fill the input buffer required by the compute unit without extra data rearrangement.

Engineering Benefits: Through this quantization strategy, the theoretical KV cache memory footprint is reduced by approximately 50%, and the required system memory bandwidth (for large batch processing scenarios) is simultaneously halved. This makes it possible to support longer contexts or larger batches within the same memory budget.

5. End-to-End Inference System Optimization: A Comprehensive Perspective from Operators to Hybrid Parallelism

The optimizations discussed above focus on the operator and single-card levels. However, to achieve high-concurrency, low-latency services, challenges regarding cross-card communication and overall computation graph scheduling must also be addressed.

1. Hybrid Parallelism and Communication Overlap
Given the massive model parameter count, adopting hybrid parallelism (data parallelism + tensor parallelism + pipeline parallelism) is an inevitable choice. The core of the optimization lies in overlapping communication with computation.

• At the Pipeline parallel boundary, traditional All-Reduce communication is synchronously blocking. We utilize a multi-stream mechanism to place communication operations on a separate compute stream, allowing them to overlap with the backward or forward computation of micro-batches.
• For the sparse indices generated by DSA, the index results across different devices require global synchronization. We employ a “communication pruning for index aggregation” strategy: performing a local Top-K compression before communication, transmitting only the top candidate indices from each device, thereby reducing the data volume of cross-node communication.

2. Graph Mode and SuperKernel
The execution of the Kimi model follows “Graph Mode” rather than single-operator mode. The advantage of Graph Mode is that the compiler can observe global data dependencies, allowing it to:

• Apply SuperKernel technology, merging multiple operators with similar semantics but different operands into a single large “super kernel graph” to reduce the scheduling overhead of CPU dispatching instructions.
• Implement a Compile Cache strategy. For Decoder phase operators with fixed shapes, the binary instructions generated by one-time compilation are cached. Subsequent requests with the same shape directly call the cache, avoiding the time consumption of repeated compilation.

3. Speculative Decoding and State Replay
During the generation phase, to accelerate decoding, the industry-standard speculative decoding framework is utilized. We introduced a lightweight Draft Model to parallelly speculate multiple subsequent tokens, which are then verified by the target model in a single pass.

• Under the KDA architecture, speculative verification needs to ensure the continuity of state recursion. We apply the ReplaySSM (State Space Model Replay) mechanism: when a draft token verification fails (i.e., is rejected), the system does not simply roll back to the previous token. Instead, it restores KDA’s internal state variables to a specific snapshot at the verification point and restarts the recursive computation from that snapshot. This avoids redundant forward state computations caused by rollbacks, improving both the success rate and execution efficiency of speculative decoding.

Conclusion and Future Outlook

The inference optimization of Kimi K3.1 Pro represents a deep software-hardware collaborative design centered around sparse computation characteristics.

From the computation side, we resolved the irregular memory access introduced by sparsity, transforming irregular accesses into hardware-friendly block loading through kernel fusion and Tile-based quantization. From the scheduling side, pipeline bubbles in decoding latency are effectively filled using multi-stream parallelism and insertion strategies, ensuring every microsecond is used for useful computation. From the system side, Graph Mode and Compile Cache minimize dynamic graph overhead, while ReplaySSM ensures the stability of the speculative mechanism in complex state-based models.

These practices clearly demonstrate that the optimization of new-generation large model inference has long surpassed simple tuning of GEMM or FlashAttention. It has entered a new era of “structural sparsity + fine-grained quantization + pipeline orchestration + state management.” For developers, deeply understanding the model’s computational topology (what is dense and what is sparse) is far more important than blindly applying general optimization libraries. Future inference frameworks will focus more on reserving hardware-friendly interfaces at the algorithmic level to achieve ultimate energy efficiency.