The Problem

FlashKDA addresses the performance gap in attention mechanisms for sequence modeling, specifically implementing Kimi Delta Attention (KDA) — a variant that combines linear attention with gated recurrent state updates. Standard PyTorch implementations of such attention patterns are too slow for production use, and existing Triton-based implementations in flash-linear-attention leave performance on the table. FlashKDA delivers CUDA kernels that run KDA at hardware-limited speed on NVIDIA GPUs with SM90+ architecture.

What This Does

FlashKDA is a CUDA extension for PyTorch that implements high-performance KDA kernels. The core logic lives in csrc/flash_kda.cpp (the C++ binding layer) and csrc/smxx/fwd_kernel1.cuh/fwd_kernel2.cuh (the actual CUDA kernels). The Python entry point is flash_kda/__init__.py, which exposes a single fwd function matching the chunk_kda API from flash-linear-attention.

The repo includes a test suite (tests/) that verifies correctness against both a torch reference (tests/torch_ref.py) and the flash-linear-attention library, plus benchmark scripts (benchmarks/) that measure performance and generate comparison tables. The setup.py handles compilation, with CUDA architecture detection defaulting to the current device.

How It Is Wired

Execution starts at main in benchmarks/bench_fwd.py:145, which calls run_bench → run_case → bench_fn — this is the only path that leaves the process (via subprocess.run). The benchmark harness runs the kernel, parses output, and generates markdown comparison tables.

The highest-traffic function is fwd in flash_kda/__init__.py, called from 10 places across the codebase. It's the single public API — every test and benchmark routes through it. torch_ref (in tests/torch_ref.py) is the second most-called function with 6 call sites; it implements the float32 reference implementation used for correctness checks.

The module graph shows no circular dependencies. flash_kda/__init__.py is the hub (Ca=3, Ce=0), meaning it's imported by three files but imports nothing itself — a clean leaf dependency. The test files (tests/test_fwd.py, tests/test_fwd_full.py) and benchmarks all depend on it but not on each other.

How To Use It

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

For CI or multi-arch builds, set FLASH_KDA_CUDA_ARCHS=all before pip install. The package requires CUDA 12.9+, PyTorch 2.4+, and SM90+ hardware.

After install, FlashKDA auto-dispatches from flash-linear-attention's chunk_kda if flash-linear-attention >= 0.5.0 is installed. Call it under torch.inference_mode() with the documented parameters (q, k, v, g, beta, scale, etc.). Set FLA_FLASH_KDA=0 to fall back to Triton.

Tests run via bash tests/test.sh. The test suite covers exact-match correctness against the torch reference and comparison with flash-linear-attention.

Real-World Use

FlashKDA fits into any transformer-style model that uses KDA-style gated linear attention — for example, long-context language models. The integration with flash-linear-attention means you can swap the Triton backend for FlashKDA with zero code changes:

from fla.ops.kda import chunk_kda

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

Code Health & Issues

Static analysis found 5 issues (1 high, 4 medium) across 3 categories:

  • High — Deep nesting in tests/torch_ref.py, tests/test_fwd.py, tests/test_fwd_full.py: max indentation depth of 10 makes control flow hard to follow. Fix: flatten with guard clauses and early returns.
  • Medium — Broad exception handling in setup.py: bare or Exception-wide except swallows errors. Fix: catch specific exceptions.
  • Medium — High branching density in csrc/flash_kda.cpp: 63 branch points over 181 lines. Fix: decompose decision-heavy logic.

No CI/CD pipeline exists — there's no .github/ or CI config, meaning changes merge without automated build/test gates. The repo has a license but no Dockerfile or lockfile. No committed secrets were found.

The Bottom Line

FlashKDA is a focused, well-structured CUDA kernel project with a clean public API and solid test coverage. The main risks are the absence of CI and the deep nesting in test files. If you need fast KDA inference on H100-class hardware, this is worth adopting; if you're on older GPUs, it won't help you.