KDA + GDN-2 bprop support + code cleanup

benchmark/linear_attention/Dockerfile

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-FROM nvcr.io/nvidia/pytorch:26.03-py3
+FROM nvcr.io/nvidia/pytorch:26.07-py3
 
 # Set working directory
 WORKDIR /workspace
@@ -26,5 +26,9 @@ RUN pip install nvidia-cutlass-dsl[cu13]==4.7.0 apache-tvm-ffi flash-linear-atte
 RUN git clone https://github.com/QwenLM/FlashQLA.git
 RUN pip install -v /workspace/FlashQLA
 
+# Install FlashKDA from source.
+RUN git clone https://github.com/MoonshotAI/FlashKDA.git
+RUN pip install -v /workspace/FlashKDA
+
 # Install the chart dependencies for plot_results.py
 RUN pip install pandas matplotlib seaborn

benchmark/linear_attention/README.md

@@ -2,7 +2,7 @@
 
 ## Introduction
 
-This directory contains benchmarking tools for linear attention operations (Gated DeltaNet and its variants) across various backends. The benchmarks target training use cases with support for forward and backward passes, grouped-value attention (GVA), and the per-sequence recurrent state ports (initial state in, final state out).
+This directory contains benchmarking tools for linear attention operations (GDN/KDA/GDN-2) across various backends. The benchmarks target training use cases with support for forward and backward passes.
 
 ## Contents
 
@@ -36,7 +36,7 @@ python benchmark_single_linear_attention.py \
     --la_backend cudnn --variant kda --data_type bfloat16 \
     --skip_ref --profile_pass bwd
 
-# cuDNN Frontend (GDN-2, forward only)
+# cuDNN Frontend (GDN-2, forward pass)
 python benchmark_single_linear_attention.py \
     --batch_size 1 --seqlen 8192 \
     --num_q_heads 16 --num_kv_heads 16 --head_dim 128 \
@@ -64,8 +64,14 @@ python benchmark_single_linear_attention.py \
     --la_backend flash_qla --variant gdn --data_type bfloat16 \
     --skip_ref --fwd_bwd
 
-# Recurrent state ports: seed with an initial state and request the final
-# state (its gradient feeds the backward pass)
+# FlashKDA comparison point (kda variant only, forward only, bf16)
+python benchmark_single_linear_attention.py \
+    --batch_size 1 --seqlen 8192 \
+    --num_q_heads 32 --num_kv_heads 32 --head_dim 128 \
+    --la_backend flash_kda --variant kda --data_type bfloat16 \
+    --skip_ref
+
+# Input initial state and dump state for every chunk
 python benchmark_single_linear_attention.py \
     --batch_size 1 --seqlen 8192 \
     --num_q_heads 8 --num_kv_heads 64 --head_dim 128 \
@@ -75,33 +81,15 @@ python benchmark_single_linear_attention.py \
 
 Run `python benchmark_single_linear_attention.py --help` for all options.
 
-Dropping `--skip_ref` validates the forward output against FLA (the same way the SDPA benchmark validates against FlashAttention 4).
+The `kda` and `gdn2` variants fuse q/k L2 normalization in-kernel on every backend; `gdn` runs unfused.
 
 ## Supported Backends
 
 | Backend | Description |
 |---------|-------------|
 | `cudnn` | cuDNN (native, via the cuDNN Frontend torch custom ops) |
-| `fla`   | FLA (flash-linear-attention, Triton) |
+| `fla`   | FLA (flash-linear-attention, Triton; `gdn`, `kda`, and `gdn2`) |
 | `flash_qla` | FlashQLA (TileLang fused GDN kernels, `gdn` variant only) |
+| `flash_kda` | FlashKDA (`kda` forward variant only) |
 
-The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere. Both passes run through autograd, exactly like a training step.
-
-## Supported Variants
-
-| Variant | Description |
-|---------|-------------|
-| `gdn`   | Gated DeltaNet: scalar per-token decay and write strength |
-| `kda`   | Kimi Delta Attention: per-key-channel decay |
-| `gdn2`  | Gated DeltaNet v2: channel-wise decay/erase/write gates (forward only, cuDNN only) |
-
-The benchmark runs `kda` and `gdn2` with the in-kernel q/k L2 normalization off (`use_qk_l2norm_in_kernel=False`) on every backend, for an apples-to-apples comparison.
-
-Recent `fla` releases dispatch `chunk_gated_delta_rule` to FlashQLA whenever `flash_qla` is importable; the benchmark sets `FLA_DISABLE_BACKEND_DISPATCH=1` (unless already set) so the `fla` backend always measures FLA's own Triton kernels and the two backends stay distinct.
-
-## Notes
-
-- Head convention: `--num_q_heads` counts the query/key heads and `--num_kv_heads` counts the value heads; the gates, output, and recurrent state live at `max(num_q_heads, num_kv_heads)` heads. Both grouping directions are supported for `gdn`: grouped-value attention (`num_kv_heads > num_q_heads`, v-heads grouped over q-heads) and GQA (`num_q_heads > num_kv_heads`, q-heads grouped over v-heads, e.g. `--num_q_heads 64 --num_kv_heads 8`). The two counts must be equal or one a multiple of the other; `kda` and `gdn2` support the GVA direction only, and so does the `flash_qla` backend.
-- The cuDNN ops use the THD (token-packed) layout internally; the benchmark expresses the dense batch as `cu_seqlens = [0, T, 2T, ...]`.
-- `--initial_state` provides a per-sequence fp32 recurrent state (its gradient is produced in the backward pass); `--store_on` requests the per-sequence final state from the forward pass and feeds its gradient in the backward pass. Both are once-per-kernel I/O ports (one `[head_dim_qk, head_dim_vo]` tile per sequence per state head).
-- Performance is measured with the torch profiler (device time of the matched kernels), with a 256 MB L2 flush before each timed iteration and the median reported. TFLOPS use the chunked-BMM FLOPs model documented in the script's `flops()`.
+The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere.