Optimizing Kimi K3.1 Pro Sparse Attention Inference: Day0 Adaptation and Performance Tuning on Ascend

Introduction: The Shift from Dense to Sparse Attention

The release of Kimi K3.1 Pro marks a significant advancement in long-context processing for large language models. Compared to its predecessor, K3.1 Pro introduces a more refined sparse attention mechanism while imposing stricter requirements on memory usage and computational latency during inference. For AI acceleration platforms like Ascend, the ability to support a new model on the day of its release (Day0) is critical, as it determines how quickly developers can deploy the model in real-world applications.

The core architectural change in K3.1 Pro is that attention computation no longer iterates over all previous tokens. Instead, an index selection mechanism restricts attention to a small set of key positions. While this design dramatically reduces computation in long-sequence scenarios, it introduces two engineering challenges: the index selection process adds extra computation and memory access, and the irregular computation graph of sparse attention makes it difficult to directly reuse optimization strategies designed for dense matrix multiplication.

Architectural Features and Computational Patterns of K3.1 Pro

Interleaved KDA and DSA Modules

K3.1 Pro employs two complementary attention units—KDA and DSA—arranged in an interleaved pattern throughout the network. KDA processes sequence information through state recurrence: it maintains a hidden state that updates over time steps. When a new token arrives, computation is based only on the current hidden state and the new input, rather than recalculating attention over the entire history. This design offers a natural advantage during inference: decoding each token requires only constant-time state updates, avoiding repeated reads of historical key-value (KV) pairs during autoregressive generation.

DSA introduces explicit sparse selection in multi-layer attention. Its workflow consists of two steps. First, the Lightning Indexer (LI) scores historical positions based on the current query and selects the top-K most relevant positions. Then, Sparse Flash Attention (SFA) reads only the KV blocks corresponding to these K positions to perform precise local attention. This “coarse screening followed by precise computation” approach reduces attention complexity from linear dependence on sequence length to dependence on the value of K.

Cross-Layer Aggregation with Attention Residuals

K3.1 Pro retains the Attention Residuals mechanism from K3, which establishes residual connections across multiple attention layers, injecting shallow-layer attention outputs into deeper layers in a weighted manner. This design helps mitigate information decay in deep networks and offers new fusion opportunities for inference optimization. When attention outputs from multiple layers need to be aggregated, operator scheduling can merge these cross-layer operations to reduce intermediate tensor traffic in memory.

Technical Foundations for Day0 Adaptation: Reuse and Migration

The core premise of Day0 support is the smooth migration of existing optimization capabilities to the new architecture. Although K3.1 Pro introduces new modules like LI and SFA, its underlying computation still relies heavily on operator patterns similar to those in K3. Therefore, the Day0 adaptation strategy on Ascend can be summarized as: reuse proven operator optimizations and make localized adjustments for new structures, rather than building everything from scratch.

Inheritance and Extension of Core Operators

The FlashKDA and MegaKDA operators, validated during the K3 phase, continue to serve as the core of KDA computation in K3.1 Pro. These operators fuse multiple small-scale matrix operations in the state recurrence process into a single kernel, reducing kernel launch overhead and memory access. For the DSA path, although SFA involves non-contiguous KV reads after top-K indexing, its internal computation can still reuse the block-based strategy from FlashAttention: organizing KV into fixed-size blocks, performing precise attention within blocks, and using careful reduction across blocks to avoid the overhead of global softmax.

The two-stage AttnRes fusion is also carried forward. This optimization merges the “attention output computation” and “residual aggregation” stages into a single fused operator, allowing attention results to participate directly in residual operations in on-chip cache, avoiding the write-back and re-read of full tensors. In K3.1 Pro, since Attention Residuals span multiple layers, the fused operator’s inputs expand from single-layer attention outputs to weighted combinations of multi-layer outputs, requiring more flexible input layouts in operator design.

Hybrid Parallelism and Communication Optimization

The model size of K3.1 Pro prevents single-card inference, requiring multi-card collaboration. Inference deployment on Ascend platforms employs a hybrid parallel strategy: tensor parallelism (TP) along the attention head dimension, expert parallelism (EP) along the expert dimension, and pipeline parallelism along the sequence dimension to hide communication latency. The sparse attention structure of K3.1 Pro imposes new requirements on communication patterns—since each query reads KV from only the top-K historical positions, KV exchange between parallel paths is no longer full-scale but dynamically varies based on indices. To address this, the communication optimization layer introduces a dynamic routing mechanism based on index tables, reducing cross-card KV transmission while maintaining correctness.

Multi-stream scheduling is another key optimization direction. During autoregressive decoding, computations from different layers can execute in parallel across multiple compute streams: some streams handle KDA state recurrence, others handle DSA indexing and sparse attention, and yet others manage communication and post-processing. Through careful stream synchronization and dependency management, hardware utilization can be significantly improved, avoiding bubbles in a single pipeline.

SuperKernel and Graph-Mode Optimization

SuperKernel is an operator merging strategy designed for Ascend hardware. Its core idea is to fuse multiple adjacent small operators in the computation graph—such as LayerNorm, residual addition, and activation functions—into a unified kernel. K3.1 Pro’s network contains numerous element-wise operations of this type, each with modest computation but frequent kernel launches and data movement that can severely degrade execution speed. SuperKernel reduces kernel launch overhead by completing these operations within a single kernel.

Graph-mode optimization operates at a higher level to optimize the inference computation graph. During offline preprocessing, dynamic structures such as variable-shape indexing operations and conditional branches are transformed into static graph representations, paired with a compilation cache mechanism to avoid repeated graph optimization and operator selection during inference. For K3.1 Pro, since the K value in LI’s top-K selection is fixed during inference, this logic can be converted into static sorting and truncation operations embedded in the static graph, eliminating dynamic graph dispatch overhead. The NZ weight layout reorders weights according to Ascend hardware’s memory access characteristics, aligning weight read patterns with the hardware cache structure to improve bandwidth utilization in matrix multiplication.

Key Mechanisms for Inference Efficiency

DSpark Speculative Decoding

Speculative decoding is an effective technique to mitigate the serial dependency of autoregressive decoding. The basic idea is to use a lightweight draft model to quickly generate multiple candidate tokens, then have the main model verify the candidate sequence in parallel. If verification passes, multiple tokens are accepted in a single forward pass, achieving “one forward pass, multiple token outputs.” DSpark is the engineering implementation of this approach on the Ascend platform. It does not modify the main model’s structure but introduces a separate candidate generation module in the inference pipeline, along with efficient acceptance/rejection logic.

In the K3.1 Pro scenario, the benefits of DSpark are particularly pronounced. Since the main model already performs sparse attention in single-step decoding, computation time per step is relatively low, making the serial dependency a larger proportion of latency. Speculative decoding, which generates multiple candidate tokens in parallel, can more fully utilize idle compute capacity. The draft model for DSpark can be a reduced-layer or reduced-head variant of K3.1 Pro, or an auxiliary module that performs lightweight sampling of the KDA state. A key engineering consideration is that the draft model’s output distribution must overlap sufficiently with the main model to maintain a high acceptance rate; otherwise, the extra computation from speculative decoding can yield negative returns.

ReplaySSM State Replay Mechanism

The KDA path in K3.1 Pro relies on state recurrence. During long-sequence inference, state updates occur incrementally, but cross-layer aggregation through Attention Residuals requires access to intermediate states from different time steps. Recomputing historical states for each aggregation would result in substantial redundant computation. The ReplaySSM mechanism avoids repeated state recurrence by caching and replaying states at key time steps.

Specifically, ReplaySSM maintains a state buffer that saves complete KDA hidden-state snapshots at preset intervals. When cross-layer aggregation in a later layer requires a state from a historical time step, the system starts from the nearest snapshot and replays only the recurrence process from that snapshot to the target time step, recovering the required state at minimal computational cost. When combined with speculative decoding, state replay computation can be scheduled during idle communication periods, further reducing its impact on overall inference latency.

Developer Recommendations

For developers deploying K3.1 Pro on Ascend platforms, the following practical recommendations are worth considering:

Leverage existing K3 optimization assets. If your team has already completed operator tuning and parallel strategy configuration for K3, adapting K3.1 Pro is largely an incremental effort. Prioritize verifying compatibility of FlashKDA/MegaKDA and AttnRes fused operators, then make localized adjustments for the LI/SFA sparse paths. This approach can significantly shorten deployment time.

Balance communication and computation. The sparse attention in K3.1 Pro substantially reduces computation, but index selection and cross-card KV routing can introduce new communication overhead. When configuring parallel paths, evaluate communication costs based on actual sequence lengths and top-K values to avoid communication becoming a new bottleneck.

Use speculative decoding to improve throughput. In online service scenarios, the acceptance rate of DSpark directly determines throughput gains. Calibrate the draft model based on actual request distributions and dynamically adjust the speculation depth (the number of candidate tokens generated per step) under different loads to achieve the optimal balance between latency and throughput.

Enable graph mode and compilation caching. For fixed-shape inference in production environments, staticizing the top-K logic of LI and enabling compilation caching can eliminate per-request dynamic dispatch overhead. This is especially important for controlling tail latency in high-concurrency scenarios.

Conclusion

Kimi K3.1 Pro, through the interleaved design of KDA state recurrence and DSA sparse attention, significantly reduces the computational burden of long-sequence inference while maintaining model expressiveness. The Ascend platform, building on operator accumulation and engineering optimization experience from the K3 phase, achieved Day0 inference support for the new architecture. This rapid adaptation capability stems from the synergistic combination of FlashKDA/MegaKDA operator reuse, AttnRes fusion, hybrid parallelism, multi-stream scheduling, SuperKernel graph optimization, speculative decoding, and state replay mechanisms. For developers, understanding the underlying logic of these optimization techniques will help more effectively harness the inference performance of the new model in specific business scenarios.