# FlashKDA, CUTLASS kernels for Kimi Delta Attention

> FlashKDA is a CUDA kernel library for the Kimi Delta Attention operator, built on CUTLASS and designed to be picked up automatically by the flash-linear-attention library rather than called directly. The engineering choices worth examining are the architecture-detection build, the opt-out that restores the Triton path, and the state layout rules that the kernel API imposes on callers.

**MoonshotAI/FlashKDA** — FlashKDA: high-performance Kimi Delta Attention kernels

- Repository: https://github.com/MoonshotAI/FlashKDA
- Stars: 1,296 · Forks: 132
- Language: Cuda
- License: MIT
- Published: 2026-10-19 · Updated: 2026-10-19 · Language: en
- Canonical page: https://hysenlabs.com/projects/moonshotai-flashkda

## One operator, specialised kernels, no new API for the user

FlashKDA is Flash Kimi Delta Attention, a set of high-performance KDA kernels built on CUTLASS, published by MoonshotAI and licensed under the MIT licence. The project description is short: high-performance Kimi Delta Attention kernels.

The interesting thing is not the kernel but how it is meant to reach the user. The README describes it as an FLA backend. Once installed, FlashKDA is auto-dispatched from the chunk_kda operator in the flash-linear-attention library, and the integration details are tracked in a pull request in that project. A user of flash-linear-attention does not write a FlashKDA call; they keep calling chunk_kda and the library picks the faster path.

That design has consequences worth spelling out, because it is different from shipping a kernel library that people call directly. The integration is version-coupled, since the requirements call for flash-linear-attention 0.5.0 or later, so a user on an older version of that library cannot use these kernels at all. The fallback is a Triton implementation that already exists, so the difference between the two paths is performance rather than availability. And the decision to switch can be made per process, which is what the opt-out is for.

The repository itself is small and conventional: a csrc/ directory with the C++ and CUDA sources, including csrc/smxx/fwd_launch.cu for the launch layer, a flash_kda/ directory for the Python package, a cutlass/ directory that is a submodule, tests/, benchmarks/, docs/, a setup.py, and a setup_clangd.sh for editor configuration.

The project also links a deep-dive blog post on the design decisions behind FlashKDA v1, dated 2026-04-22 and kept in the repository as a Markdown document, which is the right place for a design write-up of a kernel library.

## A hardware floor that decides whether you can use it

The requirements are three lines long and they narrow the audience immediately. SM90 and above, CUDA 12.9 and above, and PyTorch 2.4 and above.

SM90 is the Hopper generation, so the supported hardware starts at the data centre parts and excludes every consumer card built on Ampere or earlier, and every Ada card. That is a real constraint for a company serving a model on a fleet of workstation GPUs, and it is the first thing to check before anything else in the README.

The CUDA 12.9 floor is a similar statement about the toolchain rather than the hardware, and it is the kind of requirement that appears when a kernel uses a feature that older compilers do not emit.

The set of architectures the build can target is defined in the setup file as 90a, 100a, 103a and 120a. Those suffixes are the architecture-specific variants, the a meaning the features that are specific to that generation, and they are what the compile flags are generated from, one -gencode pair per architecture.

The third requirement, PyTorch 2.4, is not about the maths. It is about the extension API, since the package is built as a PyTorch CUDA extension and a load call has to be compatible with the runtime that loads it.

Taken together, the requirements say that this library is for recent data centre GPUs and a recent stack, and a team that has to support a mixed fleet will find that the kernels only apply to part of it.

## The build detects your GPU, and that breaks in CI without a flag

The installation is four commands, and two of them are there because of how the build works:

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

The submodule step is for CUTLASS, the template library the kernels are written against, and it is repeated inside the setup script, which runs git submodule update --init for the cutlass directory as a side effect of being imported. The no-build-isolation flag is the other significant choice: the build needs to import torch in order to detect the GPU architecture and to compile the extension, and an isolated build environment would not have torch in it.

The architecture selection is the interesting part, and it has three modes. By default the build detects the current CUDA device and compiles for that architecture. For wheel or CI builds, you compile all supported architectures explicitly, and the README gives the command:

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

The supported values are auto, which is the default, all, or a comma-separated architecture list such as 90a,100a.

The auto mode has a behaviour that will bite somebody. The detection code reads the device capability of the current device, and if there is no visible CUDA device it raises a RuntimeError, with a message that says auto requires a visible CUDA device and points at FLASH_KDA_CUDA_ARCHS=all as the alternative. In other words, a build agent with no GPU attached cannot use the default, and the error tells you what to do. That is a defensible design, since silently building for the wrong architecture is worse than failing, but it means the environment variable is not optional in a CI setup, and a wheel that is meant to run on several architectures has to be built with all.

One more build knob appears in the setup file, an environment variable for the number of parallel nvcc threads, defaulting to 32, which is a sensible thing to have in a repository that compiles a large number of template instantiations.

## Auto-dispatch, a documented opt-out and a debug switch

The integration with flash-linear-attention is designed to be invisible, and the README documents the three things you need in order to use it.

First, the library version: flash-linear-attention 0.5.0 or later, installed with pip install -U flash-linear-attention.

Second, the call. The operator is invoked under torch.inference_mode(), and the call passes the usual Kimi Delta Attention arguments, the queries, keys, values, the gate g, the beta logits and the scale, together with the gate parameters A_log, dt_bias and lower_bound, and a set of switches that push work into the kernel, including use_gate_in_kernel, use_qk_l2norm_in_kernel and use_beta_sigmoid_in_kernel, a safe_gate flag, a transpose_state_layout flag and cu_seqlens for variable-length batching.

```python
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,
    )
```

Third, the escape hatch. Setting FLA_FLASH_KDA=0 falls back to the Triton path, which is what you want when you are debugging a numerical difference between the two implementations, or when the kernel is not the bottleneck.

There is also a debugging switch, and it is a nice one. Adding logging.basicConfig(level=logging.INFO) makes the dispatcher print a line such as [FLA Backend] kda.chunk_kda -> flashkda when it takes the kernel, or a rejection with a reason when it does not. Since the dispatch happens implicitly, without that line a user whose kernel is not being used has no way to tell, and the rejected messages are the diagnosis.

## The kernel API and the state layout rules it imposes

Beneath the auto-dispatch there is a small, explicit kernel interface, and it is documented as a table of parameters with their dtypes and shapes.

The signature is flash_kda.fwd taking the queries, keys, values, the gate, the beta, the scale, an output tensor, the gate parameters A_log, dt_bias and lower_bound, and optionally the initial state, the final state and cu_seqlens.

The tensor conventions are strict. The queries, keys, gate and beta are bfloat16, the values and the output are bfloat16, and the scale is a float scalar. The A_log parameter is a float32 vector over the heads, and the gate bias is a float32 tensor over the heads and the key dimension, which is where the per-head, per-dimension gate parameters live. The lower bound is a float scalar, described as the gate lower bound with a range from -5.0 to 0.

The shapes are all in the batch, time, head, key or value layout, for example the queries at [B, T, H, K] and the values at [B, T, H, V]. And there is one hard constraint stated plainly: it currently requires K = V = 128.

The state rules are the part that will take a caller longest to understand. The initial and final states accept None, which is the stateless case, or bfloat16, or float32, and when both are provided their dtypes must match. When cu_seqlens is provided, the batch dimension must be 1, the time dimension is the total length across all the sequences, and the states have the shape [N, H, V, K] with one state per sequence. When cu_seqlens is None, each batch element is treated as an independent sequence and the state shape is [B, H, V, K].

In other words, the state layout changes shape depending on whether you are using variable-length batching, which is the kind of rule that produces a silent bug if you get it wrong, and it is why transpose_state_layout exists as a switch on the higher level call.

## Correctness against a torch reference, and two benchmark files

The test section is one command and one sentence of description.

```bash
bash tests/test.sh
```

The test file is tests/test_fwd.py, and the README describes it as correctness tests, an exact match against the torch reference, compared with flash-linear-attention. An exact match rather than a tolerance is the right standard for a forward kernel that is meant to be a drop-in replacement, and comparing against the other implementation as well as the reference means a failure tells you which of the two things broke.

Performance is documented in separate files, BENCHMARK_H20.md and BENCHMARK_GB200.md, which correspond to the two hardware generations the project targets in practice, Hopper and Blackwell. The README links the H20 one under the performance heading. The fact that the numbers are in the repository rather than in a blog post is useful, since a kernel's speed is a function of the exact shape, the sequence length and the batch size, and a headline number without those is not reproducible.

The development tooling is minimal and the project is explicit about it. For IntelliSense on the CUDA and C++ sources there is a setup_clangd.sh script that generates a .clangd file with the correct repository paths and installs the global clangd config.yaml into the config directory in the home folder. The repository also carries a config.yaml and a .clangd.template, so the generated file is derived from a template rather than being hand written.

The citation block is a BibTeX entry naming the four authors, Yutian Chen, Zhiyuan Li, Yucheng Wang and Ming Wei, with the year 2026 and the GitHub repository as the publisher, which is the standard way to make a kernel library citable when there is no paper.

## Where FlashKDA is the wrong tool

The first limit is the hardware, and it is not a preference. SM90 and above means no Ampere and no Ada parts, so a team serving the model on anything but recent data centre hardware gets no benefit at all and still has to install the package. The auto-detection build makes that worse in an indirect way, since a developer on a workstation without a supported card cannot even build the default configuration.

The second is the head dimension. The kernel currently requires K equal to V equal to 128, and there is no fallback for other shapes. A model whose key and value dimensions differ, or that uses a different head size, simply cannot use it.

The third is the scope of the integration. The documented path is inference, under torch.inference_mode, and the project is described as a set of forward kernels with a launch file in the sources. A training path would need the backward kernels, and the README does not mention any, so the memory saving that flash-linear-attention is normally used for in training is not available here.

The fourth is the coupling to another project. The kernels are reached through flash-linear-attention's dispatch, which means the behaviour of your program depends on the version of that library, on an environment variable to opt out, and on a logging line to debug. That is a reasonable design for a fast path with a maintained fallback, and it is still a dependency on somebody else's dispatch logic.

The fifth is the release hygiene. The repository publishes no GitHub releases, the last push was on 2026-09-01, and the only versioned artefact is the source tree. A team that wants to pin a kernel version has to pin a commit, and the deep-dive document is dated April 2026, so the design has probably moved since it was written.

## Conclusion

FlashKDA fits a team serving Kimi-style models on Hopper, Blackwell or newer data centre GPUs through the flash-linear-attention library, where the change to the caller is a version bump and the rest happens behind the existing chunk_kda call. It does not fit Ampere or Ada cards, since the floor is SM90, and it does not fit a project that needs a head dimension other than 128 or a training path, since the documented integration is under inference mode. Before adopting it, check that your GPUs are in the supported list of 90a, 100a, 103a and 120a, and read the state layout rules, because initial_state and final_state have different shapes depending on whether cu_seqlens is used. The licence is MIT, the repository publishes no releases, and the last push was on 2026-09-01.

## FAQ

### What are the requirements for FlashKDA?

SM90 and above, so Hopper or newer data centre GPUs, plus CUDA 12.9 and above and PyTorch 2.4 and above. The setup file targets the architecture-specific variants 90a, 100a, 103a and 120a.

### How do I build FlashKDA for every supported GPU architecture?

Set FLASH_KDA_CUDA_ARCHS=all when installing, as in `FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .`. The default auto mode detects the visible CUDA device and fails with a RuntimeError when there is none, which is the usual situation on a build agent.

### How do I make flash-linear-attention use the FlashKDA kernels?

Install flash-linear-attention 0.5.0 or later and call chunk_kda under torch.inference_mode(). The dispatch to the kernel is automatic. Set FLA_FLASH_KDA=0 to fall back to the Triton path, and enable INFO logging to see whether the kernel was used or rejected and why.

### What head dimensions does the FlashKDA kernel support?

It currently requires K equal to V equal to 128. The tensor layout is batch, time, head, key or value, so the queries are at [B, T, H, K] and the values at [B, T, H, V].

### How are the FlashKDA correctness tests run?

With `bash tests/test.sh`, which runs tests/test_fwd.py. Those tests check an exact match against the torch reference and compare the result with flash-linear-attention, and the performance numbers are kept in BENCHMARK_H20.md and BENCHMARK_GB200.md.

## Sources

- [Issues](https://github.com/MoonshotAI/FlashKDA/issues)
- [License: MIT](https://github.com/MoonshotAI/FlashKDA/blob/master/LICENSE)
- [MoonshotAI/FlashKDA on GitHub](https://github.com/MoonshotAI/FlashKDA)
- [README](https://github.com/MoonshotAI/FlashKDA/blob/master/README.md)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/moonshotai-flashkda