LLM inference pipelines increasingly hinge on specialized kernels: Kimi Delta Attention (KDA) offers a fixed-size recurrent state and much lower KV cache growth than softmax attention, but those theoretical gains only matter when the attention kernel actually feeds tensor cores efficiently. FlashKDA fills that gap by shipping a CUTLASS/CUDA implementation tuned for Hopper (SM90+) hardware so KDA’s asymptotic advantages translate to real prefill throughput improvements.
What Sets It Apart
- CUTLASS-based CUDA kernel tuned for Hopper-specific features (TMA, WGMMA) — so what: manual scheduling and warp specialization squeeze more tensor-core utilization than the Triton reference on SM90 hardware.
- Measured prefill speedups of ~1.72x (fixed-length) up to ~2.22x (uniform ragged batches like seq_lens=1024x8) on H20 versus the flash-linear-attention Triton path — so what: higher throughput for mixed-length production prompts without code changes to callers that use FLA’s chunk_kda API.
- Native support for packed/variable-length batches via cu_seqlens and automatic dispatch from flash-linear-attention — so what: existing FLA-based inference servers can gain acceleration by installing FlashKDA alone.
- Practical constraints exposed up front (requires SM90+, K=V=128, CUDA 12.9+, PyTorch 2.4+) — so what: the kernel is production-ready on Hopper-class fleets but not applicable to Ampere/Ada Lovelace deployments without future porting work.
Who It's For and Tradeoffs
- Great fit if you run Kimi Linear / KDA workloads on Hopper GPUs and need higher prefill throughput or denser utilization across mixed-length prompts. It is especially useful for inference stacks already using flash-linear-attention since FlashKDA auto-dispatches as a backend.
- Look elsewhere if your fleet is Ampere (A100) or Ada Lovelace (L40/RTX 4090), if your model uses head dimensions other than K=V=128, or if you need acceleration for softmax attention or other linear-attention variants — those are outside FlashKDA’s current scope.
Where It Fits
FlashKDA sits at the inference-kernel layer: it is a drop-in performance backend for FLA’s KDA path, intended to turn KDA’s algorithmic memory and scaling benefits into measured throughput gains on production Hopper hardware. If you need a Triton-agnostic, CUTLASS-optimized kernel for KDA on SM90+, this is the pragmatic choice; for other hardware or attention types, continue using Triton or other implementations.
