Run Kimi Delta Attention kernels on GPU with automatic backend selection
If attention operations are a bottleneck during large language model inference, this repository reduces that time. FlashKDA is a collection of high-performance CUDA kernels for the Kimi Delta Attention (KDA) algorithm. It is designed to automatically call these fast kernels internally while you continue using existing libraries in a PyTorch environment.
v1 Design Background
This section explains the decisions made during kernel design. The repository provides CUTLASS-based high-performance KDA kernels, optimized specifically for GPU architectures SM90 and above. It runs in environments with CUDA 12.9 or higher and PyTorch 2.4 or higher.
Automatic Branching for Installation and Usage
Installation is performed via the pip install command, which detects the current CUDA device during the build and compiles for that architecture. To explicitly compile for all architectures, you must set the FLASH_KDA_CUDA_ARCHS=all environment variable. After installation, calling the chunk_kda function from the flash-linear-attention library automatically selects FlashKDA as the backend. It is recommended to call this within a torch.inference_mode() context, and you can revert to the existing Triton path by setting the FLA_FLASH_KDA=0 environment variable.
Kernel API and Constraints
The core function flash_kda.fwd accepts various tensors as input, including query, key, value, gate, and beta. Currently, the dimensions of K and V must be fixed at 128. You can pass cu_seqlens for variable-length batch processing, in which case the batch size B must be 1. Initial and final states can be provided as bf16 or fp32 tensors or set to None. If states are provided, the data types of the two tensors must match.
Accuracy Verification and Debugging
Accuracy tests are performed via tests/test_fwd.py, checking for exact matches with the PyTorch reference implementation and flash-linear-attention. Performance benchmark results are recorded in the BENCHMARK_H20.md file. For debugging, adding logging.basicConfig(level=logging.INFO) allows you to check via logs whether FlashKDA was selected or the reason for rejection.
Pre-Application Checkpoints
Since this repository requires SM90 or higher GPUs, CUDA 12.9, and PyTorch 2.4, you must first check your hardware and software environment. It cannot be applied to models where K and V dimensions are not 128. You must ensure that the flash-linear-attention version is 0.5.0 or higher, and if you need a build targeting a specific architecture, do not forget to set the environment variable. License information is not specified in the README, so separate verification is required before use.
