Add CUDA training kernels (fwd+bwd) for KDA

.gitignore

@@ -8,3 +8,4 @@ build/**
 **/__pycache__/**
 .clangd
 plot*.png
+build_dev/**

BENCHMARK_TRAIN_RTX5090.md

@@ -0,0 +1,47 @@
+# KDA training benchmark (fwd / fwd+bwd) (Blackwell / RTX 5090)
+
+- Generated: 2026-08-08
+
+- Command: `python benchmarks/generate_train_benchmark_md.py`
+
+- Benchmark settings: `warmup=30`, `iters=200`, `repeats=5`
+
+- `fla_chunk_kda` configuration: `use_gate_in_kernel=True`, `use_qk_l2norm_in_kernel=True`, post-sigmoid `beta`, `lower_bound=-5`, fp32 `initial_state`
+- `flash_kda_train` configuration: `flash_kda.train.chunk_kda_train_fwd`/`chunk_kda_train_bwd`, `use_gate_in_kernel=True`, post-sigmoid `beta`, `lower_bound=-5`, fp32 `initial_state`, `chunk_size=64`; q/k l2-normalized inside the timed region (matches `use_qk_l2norm_in_kernel`)
+
+## Stage-level breakdown (CUDA vs Triton, `B=2 T=16384 H=16 D=128`, 20 reps)
+
+Reproduce: `python benchmarks/bench_train_stages.py 2 16384 16 128`
+
+| stage | Triton (ms) | CUDA (ms) | ratio |
+|-------|------------:|----------:|------:|
+| gate_cumsum | 0.305 | 0.297 | 1.03x |
+| fwd_intra | 1.549 | 1.668 | 0.93x |
+| recompute_w_u | 0.910 | 0.950 | 0.96x |
+| fwd_h | 0.739 | 0.700 | 1.06x |
+| fwd_o | 0.712 | 0.694 | 1.03x |
+| bwd_dAv | 0.428 | 0.429 | 1.00x |
+| bwd_dhu | 1.127 | 0.899 | 1.25x |
+| bwd_wy_dqkg | 2.707 | 1.925 | 1.41x |
+| bwd_intra | 1.604 | 1.638 | 0.98x |
+| reverse_cumsum | 0.369 | 0.383 | 0.96x |
+| gate_bwd | 0.759 | 0.770 | 0.99x |
+| **TOTAL** | **11.207** | **10.353** | **1.08x** |
+
+The speedup comes mainly from the two heavy backward kernels (`bwd_dhu` 1.25x, `bwd_wy_dqkg` 1.41x, the latter bandwidth-saturated at ~1.47 TB/s). The stages still below 1.0x (`fwd_intra`, `recompute_w_u`) are at the measured bandwidth floor (0.94–0.97 TB/s; the Triton kernels sit at the same wall), so the remaining gap is memory-pattern bound, not scheduling slack.
+
+### `T=8192`, `H=96`, `D=128`
+
+| Case | `flash_kda_train` fwd (ms) | `fla_chunk_kda` fwd (ms) | fwd speedup | `flash_kda_train` fwd+bwd (ms) | `fla_chunk_kda` fwd+bwd (ms) | fwd+bwd speedup |
+|------|------------------:|------------------:|--------:|------------------:|------------------:|--------:|
+| Fixed | 6.1385 | 5.4185 | 0.88× | 21.2520 | 23.4677 | 1.10× |
+| Varlen, `seq_lens`=[1300, 547, 2048, 963, 271, 3063] | 6.1990 | 5.4671 | 0.88× | 21.3378 | 23.0464 | 1.08× |
+| Varlen, `seq_lens`=`1024 x 8` | 6.1546 | 5.4465 | 0.88× | 21.1676 | 23.0427 | 1.09× |
+
+### `T=8192`, `H=64`, `D=128`
+
+| Case | `flash_kda_train` fwd (ms) | `fla_chunk_kda` fwd (ms) | fwd speedup | `flash_kda_train` fwd+bwd (ms) | `fla_chunk_kda` fwd+bwd (ms) | fwd+bwd speedup |
+|------|------------------:|------------------:|--------:|------------------:|------------------:|--------:|
+| Fixed | 4.0432 | 3.5908 | 0.89× | 13.9893 | 15.3251 | 1.10× |
+| Varlen, `seq_lens`=[1300, 547, 2048, 963, 271, 3063] | 4.1186 | 3.5776 | 0.87× | 14.1548 | 15.0392 | 1.06× |
+| Varlen, `seq_lens`=`1024 x 8` | 4.0600 | 3.5983 | 0.89× | 13.9681 | 15.1306 | 1.08× |

benchmarks/bench_train.py

@@ -0,0 +1,206 @@
+import os
+import sys
+
+# Prefer a local flash-linear-attention checkout (FLA_REPO) so results are
+# measured against the intended Triton reference, not a stale site-packages copy.
+_FLA_REPO = os.environ.get("FLA_REPO", "/root/flash-linear-attention")
+if os.path.isdir(_FLA_REPO) and _FLA_REPO not in sys.path:
+    sys.path.insert(0, _FLA_REPO)
+
+# Pin the FLA baseline to the Triton path: with working dispatch, no_grad
+# chunk_kda calls would otherwise route to the flash_kda inference backend.
+os.environ.setdefault("FLA_FLASH_KDA", "0")
+os.environ.setdefault("FLA_FLASH_KDA_TRAIN", "0")
+
+import torch
+import torch.nn.functional as F
+import math
+
+from fla.modules.l2norm import l2norm_fwd
+from fla.ops.kda import chunk_kda
+from flash_kda.train import chunk_kda_train_bwd, chunk_kda_train_fwd, prepare_chunk_indices
+
+
+def bench_fn(fn, warmup, iters, repeats):
+    for _ in range(max(warmup, 1)):
+        fn()
+    torch.cuda.synchronize()
+
+    all_ms = []
+    for _ in range(repeats):
+        torch.cuda.synchronize()
+        starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+        ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+        for i in range(iters):
+            starts[i].record()
+            fn()
+            ends[i].record()
+        torch.cuda.synchronize()
+        all_ms.extend([s.elapsed_time(e) for s, e in zip(starts, ends)])
+
+    xs = sorted(float(x) for x in all_ms)
+    n = len(xs)
+    mean = sum(xs) / n if n else float("nan")
+    mn = xs[0] if n else float("nan")
+    mx = xs[-1] if n else float("nan")
+    return mean, mn, mx
+
+
+def run_case(seq_lens, H, D, warmup, iters, repeats):
+    device = torch.device("cuda")
+    LOWER_BOUND = -5.0
+    scale_float = 1.0 / math.sqrt(D)
+
+    varlen = len(seq_lens) > 1
+    T_total = sum(seq_lens)
+    N = len(seq_lens)
+
+    if varlen:
+        cu_seqlens = torch.tensor(
+            [0] + list(torch.cumsum(torch.tensor(seq_lens), dim=0).tolist()),
+            dtype=torch.long, device=device,
+        )
+        print(f"varlen shape=[{T_total},{H},{D}] seq_lens={seq_lens} warmup={warmup} iters={iters} repeats={repeats}")
+        extra = {"cu_seqlens": cu_seqlens}
+    else:
+        print(f"shape=[{T_total},{H},{D}] warmup={warmup} iters={iters} repeats={repeats}")
+        extra = {}
+
+    chunk_indices = None
+    if varlen:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, 64)
+        extra_train = {"cu_seqlens": cu_seqlens, "chunk_indices": chunk_indices}
+    else:
+        extra_train = {}
+
+    q = F.normalize(torch.randn((1, T_total, H, D), dtype=torch.float32, device=device), p=2, dim=-1).to(torch.bfloat16)
+    k = F.normalize(torch.randn((1, T_total, H, D), dtype=torch.float32, device=device), p=2, dim=-1).to(torch.bfloat16)
+    v = torch.randn((1, T_total, H, D), dtype=torch.bfloat16, device=device)
+    g = torch.randn((1, T_total, H, D), dtype=torch.bfloat16, device=device)
+    beta = torch.randn((1, T_total, H), dtype=torch.bfloat16, device=device)
+    A_log = torch.rand(H, dtype=torch.float32, device=device)
+    # fla chunk_kda expects a flat dt_bias of shape [H * D] (bwd returns it flat).
+    dt_bias = torch.rand(H * D, dtype=torch.float32, device=device)
+
+    initial_state = torch.randn(N, H, D, D, dtype=torch.float32, device=device)
+    # upstream chunk_kda hasn't implemented use_beta_sigmoid_in_kernel;
+    # both paths take post-sigmoid beta explicitly.
+    beta_sig = beta.sigmoid().contiguous()
+
+    do = torch.randn_like(v)
+    dht = torch.randn(N, H, D, D, dtype=torch.float32, device=device)
+
+    # l2norm_fwd matches the cost of fla's use_qk_l2norm_in_kernel and mirrors
+    # the FLA dispatch wrapper, which applies l2norm before the CUDA kernels.
+    def flash_fwd(qn, kn):
+        return chunk_kda_train_fwd(
+            q=qn, k=kn, v=v, g=g, beta=beta_sig, scale=scale_float,
+            initial_state=initial_state, output_final_state=True,
+            use_gate_in_kernel=True, A_log=A_log, dt_bias=dt_bias,
+            lower_bound=LOWER_BOUND, **extra_train,
+        )
+
+    def flash_step():
+        qn = l2norm_fwd(q)[0]
+        kn = l2norm_fwd(k)[0]
+        return flash_fwd(qn, kn)
+
+    # --- flash_kda train: fwd ---
+    mean, mn, mx = bench_fn(flash_step, warmup, iters, repeats)
+    print(f"  flash_kda_train fwd    : mean={mean:.4f} ms, min={mn:.4f} ms, max={mx:.4f} ms")
+
+    # --- flash_kda train: fwd+bwd ---
+    def flash_fwdbwd():
+        qn = l2norm_fwd(q)[0]
+        kn = l2norm_fwd(k)[0]
+        o, final_state, g_cumsum, Aqk, Akk = flash_fwd(qn, kn)
+        chunk_kda_train_bwd(
+            q=qn, k=kn, v=v, beta=beta_sig, Aqk=Aqk, Akk=Akk, scale=scale_float,
+            initial_state=initial_state, do=do, dht=dht,
+            g=g_cumsum, g_org=g,
+            use_gate_in_kernel=True, A_log=A_log, dt_bias=dt_bias,
+            lower_bound=LOWER_BOUND, **extra_train,
+        )
+
+    mean, mn, mx = bench_fn(flash_fwdbwd, warmup, iters, repeats)
+    print(f"  flash_kda_train fwdbwd : mean={mean:.4f} ms, min={mn:.4f} ms, max={mx:.4f} ms")
+
+    # --- fla chunk_kda: fwd ---
+    def run_chunk_kda_fwd():
+        with torch.no_grad():
+            chunk_kda(
+                q=q, k=k, v=v, g=g, beta=beta_sig,
+                scale=scale_float,
+                initial_state=initial_state,
+                output_final_state=True,
+                use_gate_in_kernel=True,
+                use_qk_l2norm_in_kernel=True,
+                A_log=A_log, dt_bias=dt_bias,
+                lower_bound=LOWER_BOUND,
+                **extra,
+            )
+
+    mean, mn, mx = bench_fn(run_chunk_kda_fwd, warmup, iters, repeats)
+    print(f"  fla_chunk_kda fwd      : mean={mean:.4f} ms, min={mn:.4f} ms, max={mx:.4f} ms")
+
+    # --- fla chunk_kda: fwd+bwd ---
+    qg = q.clone().requires_grad_(True)
+    kg = k.clone().requires_grad_(True)
+    vg = v.clone().requires_grad_(True)
+    gg = g.clone().requires_grad_(True)
+    bg = beta_sig.clone().requires_grad_(True)
+    h0g = initial_state.clone().requires_grad_(True)
+    A_log_g = A_log.clone().requires_grad_(True)
+    dt_bias_g = dt_bias.clone().requires_grad_(True)
+
+    def run_chunk_kda_fwdbwd():
+        o, ht = chunk_kda(
+            q=qg, k=kg, v=vg, g=gg, beta=bg,
+            scale=scale_float,
+            initial_state=h0g,
+            output_final_state=True,
+            use_gate_in_kernel=True,
+            use_qk_l2norm_in_kernel=True,
+            A_log=A_log_g, dt_bias=dt_bias_g,
+            lower_bound=LOWER_BOUND,
+            **extra,
+        )
+        ((o * do).sum() + (ht * dht).sum()).backward()
+
+    mean, mn, mx = bench_fn(run_chunk_kda_fwdbwd, warmup, iters, repeats)
+    print(f"  fla_chunk_kda fwdbwd   : mean={mean:.4f} ms, min={mn:.4f} ms, max={mx:.4f} ms")
+
+
+FIXED_CASES = [
+    [8192],
+]
+
+VARLEN_CASES = [
+    [1300, 547, 2048, 963, 271, 3063],
+    [1024] * 8,
+]
+
+
+def main():
+    import argparse
+    p = argparse.ArgumentParser()
+    p.add_argument("--warmup", type=int, default=30)
+    p.add_argument("--iters", type=int, default=200)
+    p.add_argument("--repeats", type=int, default=5)
+    p.add_argument("--mode", choices=["fixed", "varlen", "all"], default="all")
+    p.add_argument("--H", type=int, default=96)
+    p.add_argument("--D", type=int, default=128)
+    args = p.parse_args()
+
+    cases = []
+    if args.mode in ("fixed", "all"):
+        cases.extend(FIXED_CASES)
+    if args.mode in ("varlen", "all"):
+        cases.extend(VARLEN_CASES)
+
+    for seq_lens in cases:
+        run_case(seq_lens, args.H, args.D, args.warmup, args.iters, args.repeats)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/bench_train_stages.py

@@ -0,0 +1,129 @@
+"""Stage-level timing breakdown: CUDA pipeline vs Triton hosts, fwd and bwd.
+
+Times each pipeline stage with cuda events over N reps after warmup.
+Usage: python benchmarks/bench_train_stages.py [B T H D]
+"""
+
+import os
+import sys
+
+# Prefer a local flash-linear-attention checkout (FLA_REPO) so results are
+# measured against the intended Triton reference, not a stale site-packages copy.
+_FLA_REPO = os.environ.get("FLA_REPO", "/root/flash-linear-attention")
+if os.path.isdir(_FLA_REPO) and _FLA_REPO not in sys.path:
+    sys.path.insert(0, _FLA_REPO)
+
+import torch
+import torch.nn.functional as F
+
+
+import fla.ops.kda.chunk_bwd as tri_bwd
+import fla.ops.kda.chunk_fwd as tri_fwd
+from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu as tri_dhu
+from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h as tri_fwd_h
+from fla.ops.gla.chunk import chunk_gla_fwd_o_gk as tri_fwd_o
+from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra as tri_bwd_intra
+from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as tri_fwd_intra
+from fla.ops.kda.gate import kda_gate_chunk_cumsum as tri_gate_cumsum
+from fla.ops.kda.gate import kda_gate_bwd as tri_gate_bwd
+from fla.ops.kda.wy_fast import recompute_w_u_fwd as tri_recompute
+from fla.ops.utils import chunk_local_cumsum as tri_cumsum
+
+import flash_kda.train as ck
+
+B, T, H, D = (int(x) for x in sys.argv[1:5]) if len(sys.argv) > 4 else (2, 16384, 16, 128)
+REPS = 20
+device = "cuda"
+
+
+def bench(fn, reps=REPS):
+    for _ in range(3):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(True)
+    e = torch.cuda.Event(True)
+    s.record()
+    for _ in range(reps):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / reps
+
+
+torch.manual_seed(42)
+dtype = torch.bfloat16
+q = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32, device=device), p=2, dim=-1).to(dtype)
+k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32, device=device), p=2, dim=-1).to(dtype)
+v = torch.rand(B, T, H, D, dtype=dtype, device=device)
+g_raw = torch.randn(B, T, H, D, dtype=dtype, device=device)
+beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid()
+A_log = torch.log(torch.empty(H, dtype=torch.float32, device=device).uniform_(1, 16))
+dt_bias = torch.randn(H * D, dtype=torch.float32, device=device)
+h0 = torch.randn(B, H, D, D, dtype=torch.float32, device=device)
+do = torch.randn(B, T, H, D, dtype=dtype, device=device)
+dht = torch.randn(B, H, D, D, dtype=torch.float32, device=device)
+scale = D ** -0.5
+RCP_LN2 = 1.4426950408889634
+
+# shared intermediates from the Triton fwd (same inputs to both pipelines' later stages)
+g = tri_gate_cumsum(g=g_raw, A_log=A_log, dt_bias=dt_bias, scale=RCP_LN2, chunk_size=64, lower_bound=-5.0)
+w, u, qg, kg, Aqk, Akk = tri_fwd_intra(q=q, k=k, v=v, gk=g, beta=beta, scale=scale, safe_gate=True)
+if qg is None:
+    _, _, qg, _ = tri_recompute(k=k, v=v, beta=beta, A=Akk, gk=g, q=q)
+h, v_new, ht = tri_fwd_h(k=kg, w=w, u=u, gk=g, initial_state=h0, output_final_state=True)
+dAqk, dv = tri_bwd.chunk_kda_bwd_dAv(q=q, k=k, v=v_new, do=do, A=Aqk, scale=scale)
+dh, dh0, dv2 = tri_dhu(q=qg, k=kg, w=w, gk=g, h0=h0, dht=dht, do=do, dv=dv, scale=scale)
+dq0, dk0, dv3, db0, dg0, dAkk0 = tri_bwd.chunk_kda_bwd_wy_dqkg_fused(
+    q=q, k=k, v=v, v_new=v_new, g=g, beta=beta, A=Akk, h=h, do=do, dh=dh, dv=dv2, scale=scale)
+
+rows = []
+
+
+def add(name, tri_fn, cuda_fn):
+    t_tri = bench(tri_fn)
+    t_cuda = bench(cuda_fn)
+    rows.append((name, t_tri, t_cuda))
+
+
+add("gate_cumsum",
+    lambda: tri_gate_cumsum(g=g_raw, A_log=A_log, dt_bias=dt_bias, scale=RCP_LN2, chunk_size=64, lower_bound=-5.0),
+    lambda: ck.kda_gate_chunk_cumsum(g=g_raw, A_log=A_log, dt_bias=dt_bias, scale=RCP_LN2, chunk_size=64, lower_bound=-5.0))
+add("fwd_intra",
+    lambda: tri_fwd_intra(q=q, k=k, v=v, gk=g, beta=beta, scale=scale, safe_gate=True),
+    lambda: ck.chunk_kda_fwd_intra(q=q, k=k, v=v, gk=g, beta=beta, scale=scale, safe_gate=True))
+add("recompute_w_u",
+    lambda: tri_recompute(k=k, v=v, beta=beta, A=Akk, gk=g, q=q),
+    lambda: ck.recompute_w_u_fwd(k=k, v=v, beta=beta, A=Akk, gk=g, q=q))
+add("fwd_h",
+    lambda: tri_fwd_h(k=kg, w=w, u=u, gk=g, initial_state=h0, output_final_state=True),
+    lambda: ck.chunk_gated_delta_rule_fwd_h(k=kg, w=w, u=u, gk=g, initial_state=h0, output_final_state=True))
+add("fwd_o",
+    lambda: tri_fwd_o(q=q, v=v_new, g=g, A=Aqk, h=h, scale=scale),
+    lambda: ck.chunk_gla_fwd_o_gk(q=q, v=v_new, g=g, A=Aqk, h=h, scale=scale))
+add("bwd_dAv",
+    lambda: tri_bwd.chunk_kda_bwd_dAv(q=q, k=k, v=v_new, do=do, A=Aqk, scale=scale),
+    lambda: ck.chunk_kda_bwd_dAv(q=q, k=k, v=v_new, do=do, A=Aqk, scale=scale))
+add("bwd_dhu",
+    lambda: tri_dhu(q=qg, k=kg, w=w, gk=g, h0=h0, dht=dht, do=do, dv=dv, scale=scale),
+    lambda: ck.chunk_gated_delta_rule_bwd_dhu(q=qg, k=kg, w=w, gk=g, h0=h0, dht=dht, do=do, dv=dv, scale=scale))
+add("bwd_wy_dqkg",
+    lambda: tri_bwd.chunk_kda_bwd_wy_dqkg_fused(q=q, k=k, v=v, v_new=v_new, g=g, beta=beta, A=Akk, h=h, do=do, dh=dh, dv=dv2, scale=scale),
+    lambda: ck.chunk_kda_bwd_wy_dqkg_fused(q=q, k=k, v=v, v_new=v_new, g=g, beta=beta, A=Akk, h=h, do=do, dh=dh, dv=dv2, scale=scale))
+add("bwd_intra",
+    lambda: tri_bwd_intra(q=q, k=k, g=g, beta=beta, dAqk=dAqk, dAkk=dAkk0, dq=dq0, dk=dk0, db=db0, dg=dg0, safe_gate=True),
+    lambda: ck.chunk_kda_bwd_intra(q=q, k=k, g=g, beta=beta, dAqk=dAqk, dAkk=dAkk0, dq=dq0, dk=dk0, db=db0, dg=dg0, safe_gate=True))
+add("reverse_cumsum",
+    lambda: tri_cumsum(dg0, chunk_size=64, reverse=True),
+    lambda: ck.chunk_local_cumsum(dg0, chunk_size=64, reverse=True))
+add("gate_bwd",
+    lambda: tri_gate_bwd(g=g_raw, A_log=A_log, dt_bias=dt_bias, dyg=dg0, lower_bound=-5.0),
+    lambda: ck.kda_gate_bwd(g=g_raw, A_log=A_log, dt_bias=dt_bias, dyg=dg0, lower_bound=-5.0))
+
+print(f"\nshape B{B} T{T} H{H} D{D}, {REPS} reps")
+print(f"{'stage':16s} {'triton(ms)':>11s} {'cuda(ms)':>9s} {'ratio':>7s}")
+tot_t = tot_c = 0.0
+for name, t, c in rows:
+    print(f"{name:16s} {t:>11.3f} {c:>9.3f} {t/c:>6.2f}x")
+    tot_t += t
+    tot_c += c
+print(f"{'TOTAL':16s} {tot_t:>11.3f} {tot_c:>9.3f} {tot_t/tot_c:>6.2f}x")