Moonshot AI just open-sourced a new tool called FlashKDA on GitHub, and it's designed to wring every last drop of performance out of Nvidia's latest Hopper-series GPUs—think the H100 and H20—for model inference. It's released under the permissive MIT license.
The tool is built for KDA, a new attention mechanism Moonshot introduced last year in its Kimi Linear paper. Here's the problem it solves: when large language models process long text, the computational cost of traditional attention scales quadratically with length. Linear attention mechanisms reduce that to a linear cost, and KDA is a refined version of that approach. The Kimi Linear model architecture alternates between three layers of KDA and one layer of traditional attention.
KDA already had an open-source implementation written in Triton, housed in the flash-linear-attention library (often abbreviated as fla). FlashKDA is a ground-up rewrite using Nvidia's low-level GPU library CUTLASS, specifically optimized to squeeze maximum performance from Hopper architecture cards.
Moonshot's own benchmarks on an H20 GPU show FlashKDA is 1.7 to 2.2 times faster than the Triton version for the same forward pass. The speedup is especially pronounced in scenarios with variable-length inputs processed in batches. It's worth noting, though, that the company only compared it to its own Triton implementation, not to other linear attention solutions out there.
For now, only the forward pass—the part used for running a model, or inference—has been open-sourced. You can't use it to train models yet; for that, you'll still need the original Triton version. The requirements are specific: you'll need a Hopper or newer GPU (SM90 architecture or above), CUDA 12.9 or higher, and PyTorch 2.4 or higher. FlashKDA has also been merged as a new backend into the upstream fla library (see pull request #852), so existing users can switch over by changing just one line of configuration.