aiGalen Guan

FlashKDA: MoonshotAI's High-Performance Kimi Delta Attention Kernels, Now Open Source

FlashKDA: MoonshotAI's High-Performance Kimi Delta Attention Kernels, Now Open Source

Attention computation in long-context LLMs is a genuine engineering nightmare. Standard self-attention is O(n²) — memory blows up first at 128K tokens, then compute. Linear attention cuts complexity to O(n), but implementation-level efficiency lags far behind FlashAttention's maturity. Most linear attention implementations run in Triton, leaving a multiple-fold gap versus hand-optimized CUDA kernels.

On April 20, 2026, MoonshotAI open-sourced FlashKDA on GitHub — high-performance CUDA kernels purpose-built for Kimi's Delta Attention (KDA, Kimi Delta Attention), built on NVIDIA CUTLASS. As of late July 2026, the project sits at 1,061 stars and 101 forks under an MIT license.

This isn't yet another general-purpose attention library. FlashKDA is the actual attention backend running in Kimi's production system — open-sourcing it means you can directly see how MoonshotAI solved the attention bottleneck in long-context inference.

1. What Is Delta Attention?

To understand FlashKDA, first understand Kimi's choice: why Delta Attention instead of standard Softmax Attention or other linear attention variants?

Standard Attention's Limits

In a traditional Transformer, every token attends to all preceding tokens:

Attention(Q, K, V) = softmax(QK^T / √d) · V

The O(n²) computation and O(n²) KV-cache memory make 128K+ context an engineering nightmare. FlashAttention pushes compute efficiency to the limit, but doesn't fundamentally change the O(n²) asymptotic complexity.

Delta Attention: Gated Incremental Updates

Delta Attention is a linear attention variant — it maintains a fixed-size recurrent state matrix, updated incrementally per step, with O(n·d²) complexity. The core formula simplifies to:

S_t = exp(g_t) · S_{t-1} + k_t^T · v_t
o_t = q_t · S_t · β_t

Where:

  • g_t (gate): Controls historical information decay rate — lets the model learn to "forget"
  • β_t (beta): Sigmoid-gated output weight
  • S_t: Recurrent state matrix of shape [K, V] (at K=V=128, this is 128×128 = 16K elements)

Key differences from standard Softmax Attention:

Dimension Softmax Attention Delta Attention
Compute complexity O(n²·d) O(n·d²)
KV-cache O(n·d), grows with sequence O(d²), fixed size
Long-range dependency Global attention, precise but expensive Recurrent compression, approximate but efficient
Forgetting mechanism None (implicit via softmax norm) Explicit gate g, model learns decay

Delta Attention's philosophy: not every token needs to precisely recall content from 100K tokens ago. For the vast majority of inference tasks, a compressed recurrent state suffices. Standard attention still wins for tasks requiring precise long-range retrieval (codebase Q&A, long-document citation).

2. FlashKDA's Design Decisions

FlashKDA v1 made several non-obvious design choices. These come directly from MoonshotAI's practical experience with Kimi long-context inference. The following key decisions were detailed in their April 20, 2026 deep-dive post (docs/20260420-flashkda-v1-deep-dive.md):

2.1 CHUNK = 16: Smaller Than Flash Linear Attention

Flash Linear Attention (FLA) uses CHUNK = 64. FlashKDA switches to CHUNK = 16 for three reasons:

  1. Numerical range fits bf16. With gate lower_bound = -5, CHUNK = 16 keeps exp(cumsum(g)) within bf16's representable range. This eliminates the need for elaborate intra-chunk rescaling tricks — the numerical instability inherent in larger chunks simply disappears.

  2. Cheap matrix inversion. Inverting a 16×16 matrix is orders of magnitude cheaper than inverting 64×64. More importantly, the 16×16 inverse can be computed directly via Neumann series expansion without any matrix decomposition.

  3. Clean SM80 MMA mapping. All CHUNK=16 operations map cleanly onto SM80 MMA instructions without architecture-specific features, ensuring good portability across GPU generations.

2.2 Dual-Kernel Split: K1 + K2

FlashKDA partitions the full computation into two kernels organized along their natural parallelism axes:

  • K1 (token-parallel, grid = N × H × num_chunks): Gate activation → L2 normalization → decay application → L / Mqk construction → matrix inversion
  • K2 (head-parallel only, grid = N × H): Chunk-by-chunk delta-rule recurrence → output projection → running state accumulation

Early prototypes used a single fused kernel. The token-parallel work in K1 was bottlenecked by the much lower parallelism of the recurrence in K2 — a large fraction of SMs sat idle waiting for K2 to finish. Splitting delivered at least a 15% end-to-end speedup and made each stage independently tunable.

2.3 bf16 Recurrent State + fp16 Matrix Inversion

Two counterintuitive precision decisions:

Why store state in bf16? The standard approach stores recurrent state in fp32. FlashKDA uses bf16 instead: cutting state shared-memory footprint roughly in half, while removing the fp32 → bf16 cast from the critical path. MoonshotAI's internal testing shows: as long as the state update itself uses fp32 FMA instructions, storing the accumulator in bf16 between updates introduces no measurable accuracy loss.

Why invert in fp16? The elements of the 16×16 inverse matrix are proven bounded within [-1, 1] (see Su Jianlin's blog analysis), so fp16's dynamic range suffices. Using fp16 avoids the fp32 → bf16 cast that bf16 MMA would require and gives the Neumann series expansion extra headroom, improving inverse accuracy.

2.4 More Low-Level Optimizations

  • Base-2 exponent. In the g_act stage, the exponent base is switched to 2, using ex2.approx.ftz.f32 PTX instruction — eliminating the change-of-base FMA entirely and benefiting from ex2's higher throughput versus exp.
  • K1 occupancy. Through aggressive shared-memory reuse (unions over non-overlapping lifetimes) and __launch_bounds__(256, 8), a small amount of register spilling is traded for a substantial increase in thread blocks per SM.
  • Register-file transposes. K2 uses MOVM_T instruction to transpose operands directly inside the register file, eliminating every intermediate shared-memory round trip between stages.

3. Performance Benchmarks

The following are from the April 22, 2026 H20 benchmark and May 26, 2026 GB200 benchmark. All tests at T=8192 sequence length, D=128 dimension, warmup=30, iters=200, repeats=5.

H20 (Hopper)

Config flash_kda (ms) fla_chunk_kda (ms) Speedup
H=96, fixed-length 2.62 4.84 1.85×
H=96, varlen batch=[1300,547,…] 2.34 4.83 2.06×
H=96, varlen batch=1024×8 2.04 4.67 2.29×
H=64, fixed-length 1.62 3.17 1.95×
H=64, varlen batch=1024×8 1.40 3.22 2.31×

GB200 (Blackwell)

Config flash_kda (ms) fla_chunk_kda (ms) Speedup
H=96, fixed-length 1.01 2.33 2.31×
H=96, varlen batch=[1300,547,…] 0.86 2.33 2.71×
H=96, varlen batch=1024×8 0.71 2.31 3.27×
H=64, fixed-length 0.92 1.58 1.70×
H=64, varlen batch=1024×8 0.48 1.54 3.21×

Two trends are clear:

  1. Variable-length batching sees the largest speedups. K1's token-level parallelism has better load-balancing advantages under variable-length workloads.
  2. Blackwell's relative speedup exceeds Hopper's. Blackwell's larger L1/shared memory and higher warp scheduler throughput directly amplify FlashKDA's fusion and state-reuse advantages.

4. Positioning vs. Alternatives

Solution Attention Type Specific Variant Implementation Hardware Required
FlashKDA Delta Attention (linear) Kimi KDA CUTLASS CUDA kernels SM90+, CUDA 12.9+
FlashAttention-3 Softmax Attention Standard/MQA/GQA CUTLASS CUDA kernels SM80+ (H100 optimized)
Flash Linear Attention (FLA) Multiple linear variants chunk_kda, etc. Triton kernels Broad compatibility
FlexAttention (PyTorch) Custom attention User-defined Torch compiler Broad compatibility

FlashKDA doesn't need to "compete" with FlashAttention — they solve different problems. FlashAttention optimizes memory and compute for standard Softmax Attention, but the O(n²) ceiling remains. FlashKDA targets linear attention — specifically KDA — with CUDA-level optimization for ultra-long-context scenarios. It fills the gap where FlashAttention runs out of compute at 128K+ and FLA's Triton kernels leave efficiency on the table.

Scoring (Multi-Dimensional Assessment)

Dimension Score Notes
Single-kernel efficiency 9/10 Deep CUTLASS optimization, near hand-tuned PTX level
Generality 4/10 Only supports KDA, hard requirement K=V=128, SM90+
Integration 8/10 Auto-registers as FLA backend, one-line switch
Documentation quality 7/10 Excellent deep-dive post, but lacks multi-scenario examples
License openness 10/10 MIT license, commercially friendly

5. Our Take

FlashKDA is an exemplary engineering case study in long-context LLM inference optimization. It doesn't try to be a "general attention library" — it focuses on one specific scenario (KDA + long-context Kimi) and achieves excellence in that narrow domain.

But its limitations are equally clear: hard requirement K=V=128, SM90+ GPUs only (Hopper and newer), no training support (forward only currently). This means its audience is narrow — essentially teams using Kimi architectures or similar Delta Attention variants.

Our recommendations:

  • If you're deploying Kimi or FLA's chunk_kda for inference, install FlashKDA immediately. Integration is automatic (FLA backend discovery), and the performance gains are real.
  • If you're optimizing inference for other linear attention variants (Mamba, RWKV, RetNet), study FlashKDA's design patterns (chunk selection strategy, dual-kernel split, precision trade-offs) — but the kernel itself can't be used directly.
  • If you only need standard Softmax Attention, keep using FlashAttention-3. FlashKDA isn't solving your problem.

MoonshotAI's choice to open-source — MIT license + a detailed deep-dive post — is uncommon among Chinese LLM companies. Laying a production-grade attention kernel bare for everyone to inspect is, in itself, a statement of technical confidence.

References

  1. MoonshotAI — FlashKDA GitHub Repository (1.1K+ stars, MIT license)
  2. MoonshotAI — FlashKDA v1 Deep-Dive Post (2026-04-20)
  3. MoonshotAI — H20 Benchmark (BENCHMARK_H20.md) (2026-04-22)
  4. MoonshotAI — GB200 Benchmark (BENCHMARK_GB200.md) (2026-05-26)
  5. FLA-org — Flash Linear Attention PR #852: FlashKDA Integration
  6. Su Jianlin — Johnson-Lindenstrauss Lemma: Applications (fp16 matrix inversion bound analysis)
  7. Tri Dao et al. — FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision