feat(kda): add KDA decode CuTe DSL kernel with per-K gating

Add KDA (Key-Driven Attention) decode support as a CuTe DSL kernel,
extending the GDN decode kernel from PR flashinfer-ai#2498 to support per-key-dimension
gating. KDA generalizes GDN's scalar gate (g in R^1) to per-K gating
(g in R^K), with the gate mapping naturally to the warp structure.

Changes:
- Extract shared gate-independent helpers from GDN kernel into
  flashinfer/gdn_kernels/_common.py (~290 lines), slimming
  gdn_decode_bf16_state.py. No GDN behavior change.
- Add HEAD_DIM=64 support to GDN dispatch (previously 128 only)
- Preserve lowBS_1chunk kernel variants for B<=4 (both GDN and KDA)
- New flashinfer/kda_kernels/ module with T=1-4 kernels for
  HEAD_DIM={64,128}, plus chunk_kda-compatible wrapper
- 80 KDA tests covering correctness, state updates, GDN reduction
- KDA decode benchmark

Tested on B200 (SM100) with CUDA 12.9. BF16 storage, FP32 compute.
GDN: 138/138 tests pass, no performance regression.
KDA: 80/80 tests pass.

AI-assisted (Claude Code)

Author: djmmoss

benchmarks/bench_gdn_decode.py

@@ -2695,7 +2695,7 @@ def main():
     parser.add_argument("--num-q-heads", type=int, default=16)
     parser.add_argument("--num-k-heads", type=int, default=16)
     parser.add_argument("--num-v-heads", type=int, default=32)
-    parser.add_argument("--head-size", type=int, default=128)
+    parser.add_argument("--head-size", type=int, default=128, choices=[64, 128])
     parser.add_argument(
         "--dtype", type=str, choices=["float16", "bfloat16"], default="bfloat16"
     )

benchmarks/bench_kda_decode.py

@@ -0,0 +1,343 @@
+"""
+Copyright (c) 2025 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+"""
+KDA (Key-Driven Attention) Decode Benchmark
+
+Benchmarks the KDA CuTe DSL decode kernel with per-K-dimension gating.
+KDA differs from GDN by having gate g[B, T, HV, K] instead of a scalar gate.
+
+Usage:
+    python benchmarks/bench_kda_decode.py --batch-size 1 4 16 64 128 256
+    python benchmarks/bench_kda_decode.py --head-size 64 --batch-size 1 32 128
+    python benchmarks/bench_kda_decode.py --seq-len 1 2 3 4 --batch-size 1 32
+"""
+
+import argparse
+import numpy as np
+import torch
+
+from flashinfer.testing import bench_gpu_time
+
+# Import the KDA decode kernel
+try:
+    from flashinfer.kda_kernels.kda_decode_bf16_state import (
+        kda_gated_delta_rule as kda_decode,
+    )
+
+    KDA_DECODE_AVAILABLE = True
+except ImportError:
+    KDA_DECODE_AVAILABLE = False
+
+
+# ============================================================================
+# FLOPs and Bytes Calculation
+# ============================================================================
+
+
+def kda_decode_flops(
+    batch_size: int,
+    num_q_heads: int,
+    num_k_heads: int,
+    num_v_heads: int,
+    head_size: int,
+    seq_len: int = 1,
+) -> int:
+    """
+    Calculate FLOPs for KDA (Key-Driven Attention) decode.
+
+    8 * K * V FLOPs per token per head:
+    1. k @ state (prediction):    2 * K * V
+    2. k^T @ v_new (update):      2 * K * V
+    3. q @ state (output):        2 * K * V
+    4. Per-K gate application:    2 * K * V  (K*V element-wise multiply + K exp() calls)
+
+    Note: K = V = head_size for KDA.
+    """
+    num_o_heads = max(num_q_heads, num_v_heads)
+    total_flops = 8 * seq_len * batch_size * num_o_heads * head_size * head_size
+    return total_flops
+
+
+def kda_decode_bytes(
+    batch_size: int,
+    num_q_heads: int,
+    num_k_heads: int,
+    num_v_heads: int,
+    head_size: int,
+    dtype: torch.dtype,
+    seq_len: int = 1,
+) -> int:
+    """
+    Calculate memory bytes for KDA decode.
+
+    Includes:
+    - Q, K, V tensors: [B, T, H, K] - dtype
+    - G tensor (per-K gate): [B, T, HV, K] - dtype (extra vs GDN)
+    - Beta: [B, T, HV] - dtype
+    - State (read + write): [B, HV, V, K] - bf16 (2 bytes)
+    - Output: [B, T, HV, V] - dtype
+    """
+    num_o_heads = max(num_q_heads, num_v_heads)
+    elem_size = dtype.itemsize
+    state_dtype_bytes = 2  # BF16 state
+
+    # Input tensors
+    q_bytes = batch_size * seq_len * num_q_heads * head_size * elem_size
+    k_bytes = batch_size * seq_len * num_k_heads * head_size * elem_size
+    v_bytes = batch_size * seq_len * num_v_heads * head_size * elem_size
+
+    # Per-K gate: [B, T, HV, K] - the extra input vs GDN
+    g_bytes = batch_size * seq_len * num_o_heads * head_size * elem_size
+
+    # Beta: [B, T, HV]
+    beta_bytes = batch_size * seq_len * num_o_heads * elem_size
+
+    # Output: [B, T, HV, V]
+    o_bytes = batch_size * seq_len * num_o_heads * head_size * elem_size
+
+    # State: [B, HV, V, K] read + write
+    state_bytes = (
+        2 * batch_size * num_o_heads * head_size * head_size * state_dtype_bytes
+    )
+
+    total_bytes = (
+        q_bytes + k_bytes + v_bytes + g_bytes + beta_bytes + o_bytes + state_bytes
+    )
+    return total_bytes
+
+
+# ============================================================================
+# Benchmark Function
+# ============================================================================
+
+
+def bench_kda_decode(
+    batch_size: int,
+    seq_len: int,
+    num_q_heads: int,
+    num_k_heads: int,
+    num_v_heads: int,
+    head_size: int,
+    dtype: torch.dtype,
+    warmup_iters: int = 10,
+    bench_iters: int = 100,
+):
+    """Benchmark KDA decode kernel for T=1,2,3,4."""
+    if not KDA_DECODE_AVAILABLE:
+        raise RuntimeError("KDA decode kernel is not available")
+
+    assert seq_len in [1, 2, 3, 4], f"KDA decode supports T=1,2,3,4, got T={seq_len}"
+
+    num_o_heads = max(num_q_heads, num_v_heads)
+
+    # Create inputs
+    T = seq_len
+    q = torch.randn(batch_size, T, num_q_heads, head_size, dtype=dtype, device="cuda")
+    k = torch.randn(batch_size, T, num_k_heads, head_size, dtype=dtype, device="cuda")
+    v = torch.randn(batch_size, T, num_v_heads, head_size, dtype=dtype, device="cuda")
+
+    # KDA-specific: per-K log-space gate [B, T, HV, K]
+    g = torch.randn(batch_size, T, num_o_heads, head_size, dtype=dtype, device="cuda")
+
+    # Beta: [B, T, HV] (pre-sigmoided)
+    beta = torch.randn(batch_size, T, num_o_heads, dtype=dtype, device="cuda")
+
+    # Initial state: [B, HV, V, K] (K-last layout, BF16)
+    state = torch.randn(
+        batch_size,
+        num_o_heads,
+        head_size,
+        head_size,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+    # Scale factor
+    scale = 1.0 / (head_size**0.5)
+
+    # Benchmark with bench_gpu_time (CUPTI for accurate kernel timing)
+    kernel_times_ms = bench_gpu_time(
+        lambda: kda_decode(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            initial_state_source=state,
+            scale=scale,
+            use_qk_l2norm_in_kernel=True,
+        ),
+        enable_cupti=True,
+        dry_run_iters=warmup_iters,
+        repeat_iters=bench_iters,
+    )
+
+    # Calculate metrics
+    kernel_median_ms = np.median(kernel_times_ms)
+    flops = kda_decode_flops(
+        batch_size, num_q_heads, num_k_heads, num_v_heads, head_size, seq_len
+    )
+    bytes_accessed = kda_decode_bytes(
+        batch_size, num_q_heads, num_k_heads, num_v_heads, head_size, dtype, seq_len
+    )
+
+    kernel_tflops = flops / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0
+    kernel_tb_per_sec = (
+        bytes_accessed / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0
+    )
+
+    return {
+        "batch_size": batch_size,
+        "seq_len": seq_len,
+        "kernel_median_us": kernel_median_ms * 1000,
+        "kernel_tflops": kernel_tflops,
+        "kernel_tb_per_sec": kernel_tb_per_sec,
+    }
+
+
+# ============================================================================
+# Runner
+# ============================================================================
+
+
+def run_kda_decode_benchmark(args, dtype):
+    """Run KDA decode benchmark for T=1,2,3,4."""
+    if not KDA_DECODE_AVAILABLE:
+        print("Error: KDA decode kernel is not available.")
+        print("Make sure flashinfer.kda_kernels.kda_decode_bf16_state is importable.")
+        return
+
+    # Filter seq_len to only valid values (1,2,3,4)
+    valid_seq_lens = [t for t in args.seq_len if t in [1, 2, 3, 4]]
+    if not valid_seq_lens:
+        print("Error: --seq-len must include values from [1, 2, 3, 4]")
+        return
+
+    print("\n" + "=" * 100)
+    print(f"KDA Decode Benchmark (T={valid_seq_lens})")
+    print(
+        f"Config: q_heads={args.num_q_heads}, k_heads={args.num_k_heads}, "
+        f"v_heads={args.num_v_heads}, head_size={args.head_size}, "
+        f"dtype={args.dtype}"
+    )
+    print("=" * 100)
+    print()
+    print(f"{'batch':>6} {'T':>4} {'time(us)':>10} {'TFLOPS':>10} {'TB/s':>10}")
+    print("-" * 100)
+
+    all_results = []
+    for batch_size in args.batch_size:
+        for seq_len in valid_seq_lens:
+            try:
+                result = bench_kda_decode(
+                    batch_size=batch_size,
+                    seq_len=seq_len,
+                    num_q_heads=args.num_q_heads,
+                    num_k_heads=args.num_k_heads,
+                    num_v_heads=args.num_v_heads,
+                    head_size=args.head_size,
+                    dtype=dtype,
+                    warmup_iters=args.warmup,
+                    bench_iters=args.iters,
+                )
+                all_results.append(result)
+
+                print(
+                    f"{result['batch_size']:>6} {result['seq_len']:>4} "
+                    f"{result['kernel_median_us']:>10.2f} "
+                    f"{result['kernel_tflops']:>10.2f} "
+                    f"{result['kernel_tb_per_sec']:>10.2f}"
+                )
+            except Exception as e:
+                print(
+                    f"{batch_size:>6} {seq_len:>4} {'ERROR':>10} - {type(e).__name__}: {e}"
+                )
+
+    print("-" * 100)
+    print()
+
+    # Summary by T value
+    for t in valid_seq_lens:
+        t_results = [r for r in all_results if r["seq_len"] == t]
+        if t_results:
+            avg_time = np.mean([r["kernel_median_us"] for r in t_results])
+            avg_tflops = np.mean([r["kernel_tflops"] for r in t_results])
+            print(
+                f"T={t}: Average time={avg_time:.2f}us, Average TFLOPS={avg_tflops:.2f}"
+            )
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="KDA Decode Benchmark",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python benchmarks/bench_kda_decode.py --batch-size 1 4 16 64 128 256
+  python benchmarks/bench_kda_decode.py --head-size 64 --batch-size 1 32 128
+  python benchmarks/bench_kda_decode.py --seq-len 1 2 3 4 --batch-size 1 32
+""",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        nargs="+",
+        default=[1, 4, 16, 64, 128, 256],
+        help="Batch sizes to benchmark",
+    )
+    parser.add_argument("--num-q-heads", type=int, default=16)
+    parser.add_argument("--num-k-heads", type=int, default=16)
+    parser.add_argument("--num-v-heads", type=int, default=32)
+    parser.add_argument("--head-size", type=int, default=128, choices=[64, 128])
+    parser.add_argument(
+        "--dtype", type=str, choices=["float16", "bfloat16"], default="bfloat16"
+    )
+    parser.add_argument(
+        "--seq-len",
+        type=int,
+        nargs="+",
+        default=[1, 2, 3, 4],
+        help="Sequence lengths (T=1,2,3,4)",
+    )
+    parser.add_argument(
+        "--warmup",
+        type=int,
+        default=10,
+        help="Number of warmup iterations",
+    )
+    parser.add_argument(
+        "--iters",
+        type=int,
+        default=100,
+        help="Number of benchmark iterations",
+    )
+    args = parser.parse_args()
+
+    # Resolve dtype
+    dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16}
+    dtype = dtype_map[args.dtype]
+
+    run_kda_decode_benchmark(args, dtype)
+
+
+if __name__ == "__main__":
+    main()

flashinfer/gdn_decode.py

@@ -1004,13 +1004,13 @@ def gated_delta_rule_decode_pretranspose(
         f"Expected state shape [B={B}, HV={HV}, V={V}, K={K}], got {state.shape}"
     )
 
-    # Backend: gdn_decode_klast_bf16_state when bf16 state, T<=4, K-last layout, K=V=128
+    # Backend: gdn_decode_klast_bf16_state when bf16 state, T<=4, K-last layout, K=V in {64,128}
     use_gdn_decode_klast_bf16_state = (
         _GDN_DECODE_KLAST_BF16_STATE_AVAILABLE
         and state.dtype == torch.bfloat16
         and T in (1, 2, 3, 4)
-        and K == 128
-        and V == 128
+        and K == V
+        and K in (64, 128)
     )
     if use_gdn_decode_klast_bf16_state:
         assert q.dtype in (torch.float16, torch.bfloat16), (