DOCSEARCHbyAlgolia
Submit Docs

FlashKDA (Flash Kimi Delta Attention)

repository·master·Indexed 19 days ago

https://github.com/moonshotai/flashkda

High-performance KDA kernels built on CUTLASS for efficient linear attention operations. Designed as a backend for flash-linear-attention's chunk_kda function, it supports SM90+ architectures, CUDA 12.9+, and PyTorch 2.4+. The v1 implementation utilizes a two-kernel fusion strategy (K1 token-parallel and K2 head-parallel) to maximize parallelism and reduce SM idling, featuring optimized chunk sizes and mixed-precision handling for numerical stability and speed.

Tokens
2.7K
Snippets
10
Records
15
Agent score
65%

What's inside FlashKDA

  1. How FlashKDA v1 kernel fusion works

    master

    FlashKDA v1 achieves high performance by partitioning the Kimi Delta Attention computation into two distinct kernels (K1 and K2) to maximize parallelism and avoid SM (Streaming Multiprocessor) idling.

    • K1 (Token-parallel): Operates on a grid of N × H × num_chunks. It handles the initial stages: g activation, L2 normalization, decay application, L / Mqk construction, and matrix inversion.
    • K2 (Head-parallel): Operates on a grid of N × H. It handles the recurrence and projection: chunk-by-chunk delta-rule recurrence, output projection, and running state accumulation.

    Splitting the pipeline into these two kernels provides at least a 15% end-to-end speedup compared to a single fused kernel design.

  2. FlashKDA v1 design parameters: Chunk Size and Precision

    master

    FlashKDA v1 uses specific design choices to balance performance and numerical stability:

    Chunk Size

    FlashKDA v1 uses CHUNK = 16. This differs from Flash Linear Attention (CHUNK = 64) and provides several benefits:

    • Numerical Stability: Keeps the range of exp(cumsum(g)) within bf16 precision (given a lower_bound of -5), avoiding complex rescaling.
    • Efficiency: 16 × 16 matrix inversion is computationally cheaper and can be done via Neumann-series expansion.
    • Hardware Mapping: Maps cleanly to SM80 MMA instructions on modern NVIDIA GPUs.

    Numerical Precision

    • Recurrent State: Stored in bf16 to halve the shared memory footprint and avoid fp32 → bf16 casts during GEMM. Accuracy is maintained by performing state updates with fp32 FMA instructions.
    • Matrix Inversion: The 16 × 16 inverse is computed in fp16. Since inverse elements are bounded within [-1, 1], fp16 is sufficient and avoids the overhead of bf16 MMA casts.
    • Sigmoid: Implemented via the PTX tanh.approx.f32 instruction for speed and precision.

  3. Install FlashKDA

    master

    To install FlashKDA, clone the repository, initialize submodules, and install via pip. By default, the build detects your current CUDA device and compiles for that architecture. For CI or wheel builds where you need to support multiple architectures, use the FLASH_KDA_CUDA_ARCHS environment variable.

    Requirements:

    • SM90 and above
    • CUDA 12.9 and above
    • PyTorch 2.4 and above
    git clone https://github.com/MoonshotAI/FlashKDA.git flash-kda
cd flash-kda
git submodule update --init --recursive
pip install -v --no-build-isolation .

  4. Set up IntelliSense for CUDA/C++ sources

    master

    Developers working on the CUDA/C++ source code can set up clangd for IntelliSense by running the provided setup script. This generates a .clangd file with correct repository paths and configures the global clangd config.yaml.

    bash setup_clangd.sh

  5. Run FlashKDA tests

    master

    To verify the installation and correctness of the kernels, run the test suite. The tests/test_fwd.py script specifically checks for correctness by performing an exact match against the PyTorch reference implementation and comparing results with flash-linear-attention.

    bash tests/test.sh

  6. Use FlashKDA as an FLA backend

    master

    FlashKDA is designed to be auto-dispatched from flash-linear-attention's chunk_kda function.

    Setup:

    1. Install flash-linear-attention >= 0.5.0 via pip install -U flash-linear-attention.
    2. Call chunk_kda within a torch.inference_mode() context.

    Configuration & Debugging:

    • Opt out: To bypass FlashKDA and fall back to the Triton path, set the environment variable FLA_FLASH_KDA=0.
    • Debug dispatch: To verify if FlashKDA is being used, add logging.basicConfig(level=logging.INFO) to your script. You should see [FLA Backend] kda.chunk_kda -> flashkda on a successful hit.
    import torch
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,
    )

  7. Generate KDA forward benchmarks

    master

    You can generate benchmark reports in Markdown format using the generate_benchmark_md.py script. Use the --device-label flag to specify the hardware being tested (e.g., Blackwell / GB200).

    python benchmarks/generate_benchmark_md.py -o BENCHMARK_GB200.md --device-label Blackwell / GB200

  8. Compile for all supported architectures

    master

    If you are building a wheel or running in a CI environment, you can explicitly compile for all supported architectures using the FLASH_KDA_CUDA_ARCHS=all environment variable. Other supported values include auto (default) or a comma-separated list of architectures (e.g., 90a,100a).

    FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .

  9. Generate KDA forward benchmarks for Hopper / H20

    master

    To generate the KDA forward benchmark report for Hopper (H20) hardware, run the following Python script. The benchmark uses specific settings: warmup=30, iters=200, and repeats=5.

    python benchmarks/generate_benchmark_hopper_h20_md.py

  10. Configure `fla_chunk_kda` for benchmarking

    master

    When benchmarking or using fla_chunk_kda, the following configuration parameters are used to control kernel behavior:

    • use_gate_in_kernel: Boolean
    • use_qk_l2norm_in_kernel: Boolean
    • use_beta_sigmoid_in_kernel: Boolean
    • lower_bound: Float (e.g., -5)
    • transpose_state_layout: Boolean

  11. Configure `fla_chunk_gated_delta_rule` for benchmarking

    master

    When benchmarking or using fla_chunk_gated_delta_rule, the following configuration parameters and shapes are applied:

    • Gate g: A scalar per-head gate of shape (1, T, H).
    • use_qk_l2norm_in_kernel: Boolean
    • transpose_state_layout: Boolean

  12. flash_kda.fwd API reference

    master

    The flash_kda.fwd function provides the core kernel interface for Flash Kimi Delta Attention.

    Constraints:

    • Currently requires K = V = 128.
    • initial_state / final_state accept None (stateless), bf16, or fp32 tensors. If both are provided, their dtypes must match.
    • Variable-length batching: If cu_seqlens is provided, B must be 1, T is the total length across all sequences, and initial_state / final_state must have shape [N, H, V, K].
    • Independent sequences: If cu_seqlens is None, each batch element is treated as an independent sequence, and the state shape must be [B, H, V, K].
    flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound,
              initial_state=None, final_state=None, cu_seqlens=None)