MoonshotAI/FlashKDAPublic

FlashKDA: high-performance Kimi Delta Attention kernels

AI summary: High-performance Kimi Delta Attention kernels built on CUTLASS for modern NVIDIA GPUs.

Stars: 1.2K; +11 today
Forks: 111
Watchers: 6
Open issues: 13
Open PRs: 5
Contributors: ~3
Commits: 9
Branches: 1

CudaMITCreated Apr 20, 2026Last push 8d ago+58 stars this week+72 this month

Star history

since Apr 19, 2026

1.2K stars as of Aug 7, 2026, tracked back to Apr 19, 2026. Historical curve reconstructed from public GitHub event archives, calibrated to the current total.

Contribution activity

commits per day, last 52 weeks

9 commits in the last yearLessMore

Signals and awards

derived from tracked data

Permissive license
MIT

What FlashKDA does

FlashKDA provides highly optimized CUDA kernels implementing the Flash Kimi Delta Attention algorithm, a variant of the attention mechanism designed for extreme performance on modern GPU architectures. Built on top of the CUTLASS template library, it offers significant speedups for large language model inference and training. It specifically targets NVIDIA SM90 architectures (Hopper) and requires newer CUDA toolkits to leverage hardware-specific tensor core instructions.

AI researchers, ML systems engineers, and HPC developers optimizing transformer models for massive GPU clusters. It requires deep knowledge of CUDA and PyTorch.

CUTLASS foundation: leverages NVIDIA's C++ template library for optimal tensor core utilization.
SM90 optimization: highly tuned specifically for modern Hopper-class GPU architectures.
PyTorch integration: ships with direct bindings for seamless use in PyTorch 2.4+ models.
Delta attention: implements a specific, highly-efficient variant of standard dot-product attention.
High-performance kernels: reduces memory bandwidth bottlenecks during LLM generation.

Where teams use it

LLM training

Accelerate the attention calculation phase when pre-training large foundation models on Hopper clusters.

Inference optimization

Deploy high-throughput text generation pipelines with reduced latency per token.

Kernel research

Study advanced CUTLASS implementations for custom neural network operations.

Hardware benchmarking

Test the theoretical memory bandwidth and compute limits of H100 GPUs.

Getting started: pip install -v --no-build-isolation .

README

master branch

FlashKDA

FlashKDA: Flash Kimi Delta Attention — high-performance KDA kernels built on CUTLASS

News

2026-04-22 — Deep-Dive Blog: the design decisions behind FlashKDA v1, read it here.

Requirements

SM90 and above
CUDA 12.9 and above
PyTorch 2.4 and above

Installation

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

By default, the build detects the current CUDA device and compiles for that architecture. For wheel or CI builds, compile all supported architectures explicitly:

FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .

Supported values are auto (default), all, or a comma-separated arch list such as 90a,100a.

Using FlashKDA as an FLA backend

Once installed, FlashKDA is auto-dispatched from flash-linear-attention's chunk_kda. See fla-org/flash-linear-attention#852 for integration details.

Requirements

Install flash-linear-attention >= 0.5.0:
```
pip install -U flash-linear-attention
```

Call chunk_kda under torch.inference_mode()

import torch
from fla.ops.kda import chunk_kda

with torch.inference_mode():
    out, final_state = chunk_kda(
        q=q, k=k, v=v, g=g, beta=beta,
        scale=scale,
        initial_state=h0,
        output_final_state=True,
        use_gate_in_kernel=True,
        use_qk_l2norm_in_kernel=True,
        use_beta_sigmoid_in_kernel=True,
        safe_gate=True,
        A_log=A_log, dt_bias=dt_bias,
        lower_bound=lower_bound,
        transpose_state_layout=True,
        cu_seqlens=cu_seqlens,
    )

Opt out: set FLA_FLASH_KDA=0 to fall back to the Triton path.

Debug dispatch: add logging.basicConfig(level=logging.INFO) to see [FLA Backend] kda.chunk_kda -> flashkda on hit, or ... rejected: <reason> on miss.

Performance

See BENCHMARK_H20.md.

Tests

bash tests/test.sh

tests/test_fwd.py — correctness tests (exact match against the torch reference; compared with flash-linear-attention)

Kernel API

`flash_kda.fwd`

flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound,
              initial_state=None, final_state=None, cu_seqlens=None)

Parameters:

Parameter	Dtype	Shape	Description
`q`	bf16	`[B, T, H, K]`	Query
`k`	bf16	`[B, T, H, K]`	Key
`v`	bf16	`[B, T, H, V]`	Value
`g`	bf16	`[B, T, H, K]`	Gate before activation
`beta`	bf16	`[B, T, H]`	Beta logits (pre-activation; sigmoid applied internally)
`scale`	float	scalar	scaling factor
`out`	bf16	`[B, T, H, V]`	Output tensor
`A_log`	fp32	`[H]`	Log-gate parameter
`dt_bias`	fp32	`[H, K]`	Gate bias
`lower_bound`	float	scalar	Gate lower bound (range from -5.0 to 0)
`initial_state`	bf16/fp32/None	`[B, H, V, K]` or `[N, H, V, K]`	(optional) Initial recurrent state
`final_state`	bf16/fp32/None	`[B, H, V, K]` or `[N, H, V, K]`	(optional, output) Final recurrent state
`cu_seqlens`	int64	`[N+1]`	(optional) Cumulative sequence lengths for variable-length batching

Currently requires K = V = 128.
initial_state / final_state accept None (stateless), bf16, or fp32 tensors. When both are provided, their dtypes must match.
When cu_seqlens is provided, B must be 1, T is the total length across all sequences, and initial_state / final_state have shape [N, H, V, K].
When cu_seqlens is None, each batch element is treated as an independent sequence, and the state shape is [B, H, V, K].

Development

To set up IntelliSense (clangd) for the CUDA/C++ sources, run:

bash setup_clangd.sh

This generates a .clangd file with the correct repository paths and installs the global clangd config.yaml to ~/.config/clangd/.

Citation

@misc{flashkda2026,
      title={FlashKDA: Flash Kimi Delta Attention},
      author={Yutian Chen, Zhiyuan Li, Yucheng Wang, Ming Wei},
      year={2026},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/MoonshotAI/FlashKDA}},
}