Kimi K3.1 Pro Inference Acceleration: Day0 Optimization on Heterogeneous AI Hardware

Background: New Inference Optimization Challenges from Model Architecture Evolution

As large language models continue to scale in parameter count and context window length, inference bottlenecks have shifted from raw operator compute efficiency to a more complex interplay of memory bandwidth, communication overhead, parallelism strategy, and scheduling orchestration. The newly released Kimi K3.1 Pro model adopts an interleaved architecture combining KDA (state recurrence attention) and DSA (sparse attention). KDA achieves linear-complexity sequence modeling through hidden state recurrence, while DSA uses a Lightning Indexer (LI) to select the historical Top-K positions and Sparse Flash Attention (SFA) to read and compute over sparse contextual KV pairs. This hybrid design significantly reduces the complexity of long-sequence inference, but it also imposes demanding adaptation requirements on the underlying compute platform: the hardware must efficiently support both dense recurrence and sparse indexing patterns while minimizing pipeline bubbles when the two alternate during execution.

In this context, the ability of a heterogeneous AI compute platform to deliver stable and efficient inference on the day of a model’s release (Day0) has become a key indicator of software stack maturity. This article focuses on how a typical heterogeneous compute platform (Ascend CANN) adapted to Kimi K3.1 Pro, examining how it achieved rapid conversion from architectural innovation to inference performance across four dimensions: operator fusion, quantization compression, parallelism strategy, and scheduling orchestration.

Operator-Level Fusion: Reducing Memory Movement and Kernel Launch Overhead

Lightning Indexer Three-Stream Coordination: Eliminating Serial Dependencies in Input Preparation

In the DSA compute flow, the Lightning Indexer must perform three steps: relevance scoring, head-weight weighting, and Top-K index selection. In traditional implementations, these steps are typically split into multiple independent kernels, with intermediate results written back to global memory after each step and re-read at the next kernel launch. This not only increases memory bandwidth pressure but also weakens the execution efficiency of small operators due to fixed kernel launch overhead.

The optimized approach fuses scoring, weighting, and selection into a single SuperKernel, with data flowing directly through on-chip cache. Going further, three-stream coordination decomposes the input preparation stage into three parallel paths: Q projection and K normalization execute in parallel, Q quantization and head-weight projection execute in parallel, and Top-K selection is pipelined with paged K-cache reads. Specifically, head-weight projection is loaded onto on-chip memory before Q projection completes; once Q quantization output is available, relevance scoring can begin immediately without waiting for the entire input batch to be ready. This design compresses the roughly 30% serial waiting time to near zero, with particularly pronounced benefits at larger batch sizes.

Two-Stage Fusion for Attention Residuals

The model’s Attention Residuals aggregate cross-layer representations. The original implementation required three separate kernels for residual addition, layer normalization, and downstream projection. The optimized version uses two-stage fusion: the first stage adds the current layer output and historical layer residuals at the register level, and the second stage fuses the normalization operator with the first matrix multiplication (GEMM) of the downstream linear layer, launching computation directly while the normalized data is still in on-chip cache. This fusion strategy not only eliminates two global memory read/write cycles but also improves arithmetic intensity, allowing the compute-intensive GEMM to hide memory access latency.

Hierarchical Quantization: A Systematic Tradeoff of Precision for Bandwidth

Full MXFP8 Quantization with Per-Group Dynamic Scaling for LI

The Q/K vectors produced during indexing are used only for relevance ranking, and their numerical precision requirements are lower than those of the final attention computation. Therefore, LI adopts full MXFP8 quantization with per-group dynamic scaling (group size = 32). Each scaling factor covers 32 elements, meaning that even if the data distribution shifts abruptly, quantization error remains confined within a small block. Empirical results on long-sequence benchmarks such as LongBench show that at this quantization granularity, LI’s Top-K recall drops by less than 0.1% compared to the BF16 baseline, while index data memory usage is halved.

FP8 KVCache with Per-Token-Head-Tile Quantization for SFA

KV cache read by SFA is stored in FP8 format with per-token-head-tile quantization granularity (tile size = 128). This granularity balances quantization accuracy and hardware efficiency: each 128-element tile has an independent scaling factor, avoiding the sensitivity of global scaling to extreme values while ensuring Tensor Cores can load data efficiently at fixed tile sizes. During inference, KV cache memory usage is reduced by 50%. For a 128K context window, a single request can free up several gigabytes of memory, allowing larger batch sizes or longer sequence processing.

Preserving a Mixed-Precision Path

It is worth emphasizing that quantization does not sacrifice projection compute precision. DSA retains BF16 for Q/K/V projections and final attention weighting, with quantization applied only to the indexing and caching stages. This mixed strategy of high precision on critical paths and low precision on auxiliary paths maintains model output quality while systematically compressing storage and data movement overhead.

Long-Sequence Parallelism: CPxTP Combination and Load Rebalancing

Combining Sequence Parallelism (CP) and Tensor Parallelism (TP)

In ultra-long sequence scenarios (e.g., contexts up to 256K), a single card cannot hold the full activations in memory, and Prefill compute time grows linearly with sequence length, directly increasing time-to-first-token (TTFT). Using TP alone to split attention heads still requires each card to process the full sequence, which does not address the sequence-length bottleneck.

The solution introduces a combined CP and TP deployment: the input is split into multiple shards along the sequence dimension and distributed across compute cards, while within a TP group, the attention heads (or KV heads) of the same shard are further split across cards. This combined strategy delivers two-dimensional benefits: per-card activation memory and intermediate tensor usage decrease proportionally to the number of sequence shards, and TTFT drops significantly as compute scales down linearly with sequence length. The challenge is that after sharding, attention computation for each sequence shard requires KV states from other shards, with communication volume proportional to the sharding granularity.

Zigzag Layout: Balancing Cross-Card Indexing Compute Load

DSA’s causal indexing computation has inherent load imbalance: later sequence shards have more historical KV entries, resulting in heavier Top-K filtering workloads in the Lightning Indexer. With uniform sharding, the last shard could be several times slower than the first in indexing compute, creating a bottleneck effect.

The solution is a Zigzag layout that pairs the first and last sequence segments on the same card. Specifically, shard 1 and shard N are assigned to card 0, shard 2 and shard N-1 to card 1, and so on. Each card simultaneously processes one lightweight shard (near the sequence start) and one heavyweight shard (near the sequence end), statistically balancing indexing compute across cards. After computation, an AllGather operation restores the original sequence order, and this step overlaps with the next layer’s computation through a multi-stream mechanism, making the additional communication latency negligible.

DSL-Based FlashKDA Cross-Card Implementation

State Summary Exchange and Composition

For KDA’s Prefill stage, the cross-card implementation of FlashKDA (CP version) is key to reducing TTFT. The core challenge is that KDA’s state recurrence has sequential dependencies: the initial hidden state of shard i depends on the final hidden state of shard i-1. A naive horizontal split of the sequence would force subsequent shards to wait for all preceding shards to finish, degenerating into serial execution.

The solution proceeds in three phases:

1. Parallel pre-computation: Each card independently computes its shard’s state summary (the contribution of the input sequence to the hidden state), requiring no dependency on other shards’ results.
2. Summary exchange and composition: Through an AllGather operation, each card obtains the state summaries of all shards and composes them locally in sequence order to derive the exact starting state for its own shard.
3. Parallel post-computation: Each card completes the full attention output for its shard in parallel using the derived starting state.

DSL Abstraction and Code Generation Optimization

The above logic is implemented through CANNBot-DSL, a domain-specific language for AI operator development. The advantage of the DSL is that developers only need to describe the mathematical logic of state summary computation and state composition, while the compiler automatically handles tensor memory layout planning, communication operator insertion, and pipeline partitioning. The compiler also generates multi-stage pipelines: while waiting for the AllGather result, local computation for the shard is launched in advance, ensuring overlap between communication and computation.

System-Level Coordination: From Single Operators to Full-Path Scheduling Optimization

Beyond the targeted optimizations described above, the overall solution includes the following system-level techniques:

• Multi-stream orchestration: Prefill and Decode are placed in separate compute streams, allowing KV cache writes during Decode to execute in parallel with the next token’s Q projection, reducing decoding idle gaps.
• Graph-mode optimization: The entire inference process is constructed as a static compute graph, enabling the operator scheduler to predict memory lifetimes, reuse memory pools, and reduce dynamic allocation overhead.
• Compilation caching: Optimized kernels are pre-compiled for common K3.1 Pro input shapes (e.g., 2K/32K/128K contexts), so runtime execution hits the cache directly and skips repeated compilation and optimization.
• Speculative inference and state replay: The DSpark speculative inference mechanism uses a small model to draft multiple candidate tokens in parallel, with the large model validating them in a single pass. ReplaySSM caches intermediate state matrices during long-sequence decoding and replays key states in a skip-wise manner, avoiding full recomputation.

Summary and Outlook

Kimi K3.1 Pro’s rapid Day0 adaptation on a heterogeneous compute platform validates an efficient path: operator-level fusion eliminates memory bottlenecks, hierarchical quantization reduces storage pressure, the CPxTP combination with Zigzag layout addresses long-sequence parallel scalability, and DSL abstraction significantly shortens the development cycle for new architecture operators.

For developers, the key insight from this case is that when facing a new-generation model architecture, one should not focus solely on single-operator optimization. Instead, a systematic performance optimization framework should be built across four dimensions: compute patterns (dense/sparse), data precision (BF16/FP8/MXFP8), parallelism dimensions (sequence/tensor/head), and scheduling strategies (multi-stream/graph mode). As sparse attention, state space models, and other new architectures continue to emerge, inference optimization has evolved from brute-force computation into a comprehensive discipline of intelligent scheduling and precision engineering. Going forward, algorithm-hardware co-design will play an increasingly decisive role in determining the energy-efficiency boundaries of large model deployment.