Moonshot AI Open-Sources FlashKDA, Making Kimi Linear 2x Faster on GPUs

Moonshot AI just open-sourced FlashKDA, a high-performance CUDA kernel for Kimi Delta Attention (KDA), the linear attention mechanism at the core of their Kimi Linear architecture. The repo is available under an MIT license.

KDA was already a strong architecture on paper, but its Triton-based implementation in flash-linear-attention was leaving significant GPU performance unrealized. FlashKDA is the production answer: a fused CUTLASS kernel that slots in as a drop-in backend with no model changes required.

Why KDA?

Kimi Linear is a hybrid linear attention architecture that outperforms full attention under fair comparisons across short-context, long-context, and reinforcement learning scaling regimes. That claim rests on a specific architectural choice.

Standard softmax attention is expensive because every token attends to every other token, so cost scales quadratically with sequence length. Linear attention replaces this with a recurrent state updated token-by-token, making cost linear. The tradeoff is that a fixed-size state cannot perfectly remember everything. KDA extends Gated DeltaNet with a finer-grained gating mechanism that controls, channel by channel, what gets written into memory and what gets discarded, making the fixed memory budget go further.

The hybrid architecture uses a 3:1 KDA-to-global MLA ratio, which reduces KV cache requirements by up to 75% and boosts decoding throughput by up to 6x for contexts as long as 1M tokens. A bespoke chunkwise algorithm achieves high hardware efficiency through a specialized variant of Diagonal-Plus-Low-Rank (DPLR) transition matrices, substantially cutting computation compared to the general DPLR formulation.

The performance gap FlashKDA closes

The original KDA kernel in flash-linear-attention was written in Triton, which is well-suited for rapid development but surrenders performance for complex recurrent operations where fine-grained control over memory layout, register allocation, and instruction scheduling matters.

FlashKDA is the CUTLASS replacement. On an H20 GPU with sequence length 8192, 96 heads, and head dimension 128, the benchmarks show:

• Fixed-length sequences: 1.85x speedup over fla_chunk_kda
• Variable-length sequences (mixed batch): 2.06x speedup
• Uniform variable-length (1024 x 8): 2.29x speedup

Moonshot cites a 1.72x–2.22x range across the full spread of configurations.

What makes the kernel fast

The design deep-dive is unusually transparent about the engineering decisions:

• Chunk size of 16, not 64. Flash Linear Attention uses 64-token chunks. FlashKDA uses 16. This keeps accumulated gate values within bf16's representable precision, eliminating rescaling tricks inside each chunk and making the required matrix inversion far cheaper: a 16x16 inversion via Neumann series expansion rather than a 64x64 one.
• Two-kernel split. Early prototypes fused everything into one kernel. The recurrence in the second stage (K2) has much lower parallelism than the token-parallel first stage (K1), so fusing them left most GPU streaming multiprocessors idle. Splitting into two independently tunable kernels yielded at least a 15% end-to-end speedup.
• bf16 on-chip state. The recurrent state is stored in bf16 rather than fp32, cutting shared memory footprint roughly in half and removing a cast on the critical path of every matrix multiply feeding the state.
• fp16 matrix inversion. The 16x16 inverse is computed in fp16. Because elements are bounded within [-1, 1], fp16's narrower dynamic range is sufficient and avoids an extra cast that bf16 MMA would require.
• Base-2 exponent. The gate activation rebases to base 2 and uses the ex2.approx.ftz.f32 PTX instruction, which has higher throughput than the standard exp and eliminates a change-of-base multiply.
• Register-file transposes in K2. The MOVM_T instruction transposes operands directly in the register file, eliminating every intermediate shared-memory round trip between stages of K2.

Drop-in integration

If you're already using flash-linear-attention, integration is automatic once FlashKDA is installed. The library auto-dispatches to the CUTLASS kernel when it detects a compatible GPU (SM90+).

git clone https://github.com/MoonshotAI/FlashKDA.git flash-kda
cd flash-kda
git submodule update --init --recursive
pip install -v --no-build-isolation .

Call chunk_kda from flash-linear-attention under torch.inference_mode() and FlashKDA will dispatch automatically. Verify by adding logging.basicConfig(level=logging.INFO): a successful dispatch logs [FLA Backend] kda.chunk_kda -> flashkda. To fall back to the Triton path, set FLA_FLASH_KDA=0.

Hardware requirements: SM90 or above (Hopper-class GPUs and newer), CUDA 12.9+, and PyTorch 2.4+. The repo includes benchmark files for both H20 and GB200.

Limitations before you integrate

• Head dimension is fixed at K = V = 128; other sizes are unsupported in v1.
• Only the forward pass is implemented. Training still requires the Triton path.
• SM90+ only; Ampere GPUs (A100, A10) are not supported.
• The 3:1 KDA-to-MLA ratio requires empirical tuning when adapting the architecture. Too few full-attention layers causes context drift; too many erodes efficiency gains.

Where this fits in the broader kernel landscape

FlashKDA joins a growing set of CUTLASS-based attention kernels, including FlashAttention-3, that treat Triton implementations as prototypes rather than production targets. The release reflects a broader shift: for novel attention variants running on Hopper-class hardware, architecture-aware CUDA kernels with persistent, asynchronous pipelines consistently outperform what Triton can generate. For teams running Kimi Linear or building hybrid linear attention models, FlashKDA makes the performance case for the architecture substantially stronger.