The code that runs close to the graphics card is transitioning to the MIT license.
Moonshot AI's FlashKDA library doubles prefill speed on Hopper GPUs. Seamless integration is now available under the open-source MIT license.
A language model spends most of its computing time on a few highly repetitive operations, including the attention mechanism. Running them quickly requires code written as close as possible to the graphics card, known as kernels. Moonshot AI is releasing its own under the name FlashKDA, a CUTLASS-based implementation of the Kimi Delta Attention kernels, the in-house attention mechanism that powers its models.
The announced performance gain focuses on a specific phase. On an H20 card, the company claims a prefill—the initial ingestion of the context before the model starts responding—that is 1.72 to 2.22 times faster. The comparison is made against the reference implementation of flash-linear-attention, meaning against another code executing the same mechanism, rather than against classic attention. Benchmarks for the GB200 generation accompany the repository.
Integration is designed to be painless. Once installed, the library automatically replaces the usual flash-linear-attention path without any modification to the calling code, while an environment variable allows reverting to the Triton implementation. A logging mode indicates whether the substitution was successful, or for what reason it was bypassed. The provided tests verify the exact equality of outputs compared to the reference implementation.
The constraints are stated straightforwardly. The suite requires Hopper generation cards or newer, CUDA 12.9 and PyTorch 2.4 at a minimum, and currently only accepts key and value dimensions fixed at 128. The code is released under the MIT license.