Large Language Model Inference Acceleration in Practice: From Sparse Attention to Cluster Deployment Optimization

Introduction: Challenges and Opportunities in LLM Inference Optimization

As large language models (LLMs) grow in parameter scale and context window, inference bottlenecks have shifted from simple compute shortages to a multi-faceted constraint involving compute, memory bandwidth, capacity, and communication efficiency. The industry has widely recognized that hardware iterations alone cannot solve the root problem. Algorithm-system co-design has emerged as the key lever for improving inference efficiency.

The recently open-sourced K3.1 Pro model (representative of the Kimi K3.1 Pro family) offers an excellent research sample for inference optimization. Architecturally, the model fuses two sparse attention mechanisms—KDA (Kernel-based Dynamic Attention) and DSA (Dynamic Sparse Attention)—while retaining cross-layer attention residuals. How to enable such a complex new architecture to run efficiently on Ascend AI hardware while reusing existing optimization assets is a compelling technical topic. This article reconstructs a reusable inference acceleration methodology based on the adaptation practice of this model.

1. Architecture Deep Dive: Understanding K3.1 Pro’s Sparse Attention Design

Before discussing optimization, it is essential to clarify the core innovations of the model architecture. K3.1 Pro is not simply a stack of layers; it introduces two complementary sparse attention paths to reduce the complexity of long-sequence inference.

1.1 KDA: Linear Attention via State Recurrence

The core of KDA is the use of a fixed-size hidden state to replace the traditional KV Cache. Through a state recurrence formula, it compresses historical information into each time step’s computation. Mathematically, this resembles linear attention, turning per-step decoding complexity from sequence-length-dependent to constant-time. However, this compression is lossy—it cannot precisely recover fine-grained information from distant positions.

1.2 DSA: Precise Sparse Attention with Index Selection

To compensate for KDA’s precision loss, the model introduces DSA. Built on top of traditional MLA (Multi-head Latent Attention), DSA adds two critical sub-modules:

• LI (Lightning Indexer): A lightweight index generator that quickly scores the relevance between the current Query and historical Keys, then selects the Top-K most relevant positions.
• SFA (Sparse Flash Attention): Using the indices produced by LI, it gathers only the corresponding positions from the main KV cache and performs sparse FlashAttention computation.

This design lets the model maintain near-global attention accuracy while reducing attention complexity from O(N²) to O(N·K), where K ≪ N.

1.3 Attention Residuals: Cross-Layer Information Aggregation

The model continues the cross-layer attention residual approach: attention matrices or hidden states computed at lower layers are passed directly to higher layers via residual connections. This structure aids gradient flow through deep networks and, at inference time, enables reuse of intermediate computation results.

2. Day-0 Adaptation Strategy: Maximizing Reuse of Existing Optimization Assets

When facing a new architecture, starting from scratch is the last thing you want. The efficient path is to abstract common model patterns, identify new operators, and reuse proven kernels. The Ascend CANN platform provides comprehensive support for this.

2.1 Operator Level: Reuse and Fine-Tuning

The optimization team performed a detailed breakdown of K3.1 Pro’s operator requirements:

Model Sub-structure	Core Operator	Optimization Strategy
KDA state recurrence	FlashKDA / MegaKDA	Full reuse of fused kernels from prior versions—no changes needed
Attention Residuals	AttnRes	Two-stage fusion to reduce global memory reads/writes
MoE expert modules	MegaMoE / MC2	Reuse mature MoE parallelism and communication optimizations
DSA index selection	LI-related operators (scoring, TopK, Gather)	New custom fused operators, following scheduling patterns of existing FlashAttention

2.2 System Level: Multi-Dimensional Co-Optimization

Beyond single-operator performance, system-level optimization is equally critical:

• Hybrid parallelism: Combining data parallelism (DP), tensor parallelism (TP), and expert parallelism (EP), with All-to-All communication tuned for MoE structures.
• Multi-stream execution: Assigning independent compute tasks (e.g., new Q projection and old K normalization) to different compute streams for concurrent execution, hiding memory latency.
• Compilation and storage optimization: Applying SuperKernel (large-scale fused kernels) to reduce kernel launch overhead; using Graph Mode to reduce Python dispatch overhead; introducing compilation caching to avoid redundant builds; using NZ weight layout to optimize matrix multiplication memory access.

Key Insight: Day-0 support is fundamentally about rapid mapping at the intermediate representation (IR) level. By partitioning model structure into known and unknown patterns, atomizing the unknown parts, and leveraging CANN’s auto-scheduling engine to map them onto efficient kernels, adaptation cycles are dramatically shortened.

3. Deep Optimization Case Study: DSA Fusion and Multi-Stream Decoding

DSA is both the performance priority and the hardest nut to crack. Its compute flow involves multiple memory accesses and operator switches, making it a bottleneck risk.

3.1 Fused Execution of LI and SFA

In a naive implementation, after LI computes relevance scores and TopK selection, index information is written back to global memory and then re-read by SFA. Under long sequences, this creates enormous memory bandwidth pressure. The optimized approach fuses LI and SFA at the kernel level:

1. Compute raw relevance scores between Query and Key.
2. Perform head-weight weighting and TopK selection on-chip, generating index pointers directly.
3. Use those pointers to read from the Paged KV cache and feed directly into SFA.

This eliminates intermediate global memory round-trips and reduces kernel launch counts.

3.2 Indexer Prolog: Three-Stream Concurrent Preprocessing

During decoding, generating the Q projection for the first token traditionally requires waiting for the current token’s full Q computation. The Indexer Prolog technique introduces pipelined parallelism:

• Stream 1: Q projection main computation (partial heads).
• Stream 2: K normalization and old KV update.
• Stream 3: Q quantization, running in parallel with head weight projection.

This three-stream collaboration dramatically compresses serial wait times in the input preparation phase, boosting decoding throughput.

4. Reducing Storage and Bandwidth Overhead: Layered Quantization in Practice

Supporting ultra-long contexts within limited HBM bandwidth makes quantization essential. K3.1 Pro’s quantization strategy reflects a philosophy of layered, fine-grained management.

Data Object	Quantization Format	Granularity	Rationale
Main attention Q projection	BF16	Maintained precision	Used in the main compute path for numerical stability
LI index Q/K	MXFP8	Per-group (Group Size=32)	For index scoring tasks; dynamic quantization balances precision and speed
SFA KV Cache	FP8	Per-token-head-tile (Tile=128)	Targets the main memory consumer; halves capacity and enables hardware matrix-multiply acceleration at tile=128

The elegance of this strategy: index scoring tolerates lower precision, so aggressive quantization is acceptable; KV Cache dominates memory capacity, so FP8 directly halves it, and the per-128-element tile quantization enables corresponding matrix-multiply acceleration units on specific hardware. The result is reduced memory footprint and significantly lower data-movement bandwidth overhead.

5. Ultra-Long Sequence Prefill: Context Parallelism (CP) and Load Balancing

When context length exceeds a single accelerator’s memory, multi-accelerator parallelism is mandatory. K3.1 Pro supports a combination of TP and CP.

5.1 Zigzag Layout for Load Balancing

DSA has a notable trait: the further a position is in the sequence, the larger its candidate history index set, creating an increasing compute load trend. If the sequence is simply split sequentially (e.g., accelerator 1 handles the first quarter), accelerator 4 would carry far more work than accelerator 1.

The solution is a Zigzag layout:

• Allocate both the start and end segments of the sequence to the same accelerator.
• Accelerator 1 handles segments 1 and N; accelerator 2 handles segments 2 and N−1; and so on.

This equalizes the load across accelerators. After computation, an AllGather collective communication restores the original sequence order for downstream layers.

5.2 CANNBot-DSL-Based FlashKDA CP Implementation

In this case, KDA’s prefill time scales linearly with sequence length, directly affecting Time To First Token (TTFT). With TP alone (splitting along attention heads), each accelerator still processes the full sequence length.

The practice implements a context-parallel version of FlashKDA using CANNBot-DSL (a domain-specific language for describing compute graphs). The core flow:

1. Shard input: Split the long sequence along the sequence dimension across accelerators.
2. Parallel precomputation: Each accelerator computes the state summary for its own shard.
3. State exchange: Sync prior shard state summaries across accelerators via AllGather.
4. Sequential synthesis: Each accelerator performs recurrence updates and computes final outputs based on the received full state summaries.

This approach effectively reduces per-accelerator activation memory for long sequences and significantly shortens TTFT.

6. Speculative Decoding and State Replay: The Last Mile of Decoding Acceleration

In autoregressive decoding, tokens are generated one by one, bounded by memory bandwidth. Two advanced techniques address this:

1. DSpark speculative inference: A small draft model predicts multiple future tokens in one shot, and the target large model verifies them in a single forward pass. If verification succeeds, what would have required K iterations now needs only one forward pass plus one verification, dramatically increasing throughput.
2. ReplaySSM state replay: In speculative rejection or rollback scenarios, KDA’s recurrent state must be rolled back. The ReplaySSM mechanism records state snapshots at time points, enabling rapid restoration to a prior state and avoiding recomputation caused by state corruption.

7. Summary and Outlook

Through multi-layered co-optimization spanning operators, algorithms, systems, and scheduling, we not only achieved rapid adaptation of K3.1 Pro on the Ascend platform but also validated a general-purpose LLM inference acceleration methodology:

1. Understand the architecture first: Dissect model structure to separate “unchanged parts” from “new parts.”
2. Combine reuse with fine-tuning: Maximize reuse of validated operator libraries to avoid reinventing the wheel.
3. Memory bandwidth is the primary bottleneck: Reduce data movement through fusion, quantization, and parallelism.
4. Scheduling is the invisible engine: Multi-stream concurrency and load balancing can deliver an additional 20%–30% performance gain.

For developers and architects, there is no need to obsess over peak compute of a single hardware platform. Real performance differences often lie in the fine-grained orchestration of memory access patterns and compute pipelines. This methodology is equally applicable to rapid adaptation and deep tuning of other new models, forming a core competency for keeping pace with future LLM generational shifts.