diff --git a/.gitignore b/.gitignore index 66b35c4..96d9ec5 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ build/** **/__pycache__/** .clangd plot*.png +build_dev/** diff --git a/BENCHMARK_TRAIN_RTX5090.md b/BENCHMARK_TRAIN_RTX5090.md new file mode 100644 index 0000000..b773fd9 --- /dev/null +++ b/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× | diff --git a/benchmarks/bench_train.py b/benchmarks/bench_train.py new file mode 100644 index 0000000..f119efb --- /dev/null +++ b/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() diff --git a/benchmarks/bench_train_stages.py b/benchmarks/bench_train_stages.py new file mode 100644 index 0000000..e008472 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/benchmarks/generate_train_benchmark_md.py b/benchmarks/generate_train_benchmark_md.py new file mode 100644 index 0000000..2a249fe --- /dev/null +++ b/benchmarks/generate_train_benchmark_md.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +""" +Run ``bench_train.py`` twice (default ``H`` and ``--H 64``), parse stdout, and +write a training benchmark markdown report. + +Reports mean latency for ``flash_kda_train`` (CUDA training kernels) and +``fla_chunk_kda`` (FLA Triton), fwd and fwd+bwd, plus speedup +``fla_mean / flash_mean``. Generated date is UTC, day precision only +(YYYY-MM-DD). +""" +from __future__ import annotations + +import argparse +import ast +import datetime as _dt +import re +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +BENCH_TRAIN = Path(__file__).resolve().parent / "bench_train.py" +DEFAULT_OUT = REPO_ROOT / "BENCHMARK_TRAIN_RTX5090.md" +DEFAULT_DEVICE_LABEL = "Blackwell / RTX 5090" + +FLA_CHUNK_KDA_OPTIONS_MD = ( + "- `fla_chunk_kda` configuration: `use_gate_in_kernel=True`, " + "`use_qk_l2norm_in_kernel=True`, post-sigmoid `beta`, " + "`lower_bound=-5`, fp32 `initial_state`" +) + +RE_HEADER_FIXED = re.compile( + r"^shape=\[(\d+),(\d+),(\d+)\] warmup=(\d+) iters=(\d+) repeats=(\d+)\s*$" +) +RE_HEADER_VARLEN = re.compile( + r"^varlen shape=\[(\d+),(\d+),(\d+)\] seq_lens=(\[[^\]]+\]) " + r"warmup=(\d+) iters=(\d+) repeats=(\d+)\s*$" +) +RE_RESULT = re.compile( + r"^\s+(.+?)\s*:\s*mean=([\d.]+) ms, min=([\d.]+) ms, max=([\d.]+) ms\s*$" +) + + +def run_bench(extra_argv: list[str]) -> str: + cmd = [sys.executable, str(BENCH_TRAIN), *extra_argv] + proc = subprocess.run( + cmd, + cwd=str(REPO_ROOT), + capture_output=True, + text=True, + check=False, + ) + if proc.returncode != 0: + sys.stderr.write(proc.stderr or "") + sys.stderr.write(proc.stdout or "") + raise subprocess.CalledProcessError(proc.returncode, cmd, proc.stdout, proc.stderr) + return proc.stdout + + +def parse_stdout(text: str) -> list[dict]: + cases: list[dict] = [] + current: dict | None = None + + def new_case(kind, *, T, H, D, warmup, iters, repeats, seq_lens=None) -> dict: + c = { + "kind": kind, "T": T, "H": H, "D": D, + "warmup": warmup, "iters": iters, "repeats": repeats, + "flash_fwd_ms": None, "fla_fwd_ms": None, + "flash_fwdbwd_ms": None, "fla_fwdbwd_ms": None, + } + if seq_lens is not None: + c["seq_lens"] = seq_lens + return c + + for line in text.splitlines(): + m = RE_HEADER_VARLEN.match(line) + if m: + if current is not None: + cases.append(current) + t, h, d, seq_lens, w, it, rep = m.groups() + current = new_case("varlen", T=int(t), H=int(h), D=int(d), + warmup=int(w), iters=int(it), repeats=int(rep), seq_lens=seq_lens) + continue + + m = RE_HEADER_FIXED.match(line) + if m: + if current is not None: + cases.append(current) + t, h, d, w, it, rep = m.groups() + current = new_case("fixed", T=int(t), H=int(h), D=int(d), + warmup=int(w), iters=int(it), repeats=int(rep)) + continue + + m = RE_RESULT.match(line) + if m and current is not None: + name, mean, _mn, _mx = m.groups() + name = name.strip() + if name == "flash_kda_train fwd": + current["flash_fwd_ms"] = float(mean) + elif name == "fla_chunk_kda fwd": + current["fla_fwd_ms"] = float(mean) + elif name == "flash_kda_train fwdbwd": + current["flash_fwdbwd_ms"] = float(mean) + elif name == "fla_chunk_kda fwdbwd": + current["fla_fwdbwd_ms"] = float(mean) + + if current is not None: + cases.append(current) + return cases + + +def _fmt_seq_lens(seq_lens_str: str) -> str: + try: + xs = ast.literal_eval(seq_lens_str) + except (ValueError, SyntaxError): + return seq_lens_str + if not isinstance(xs, list) or not xs: + return seq_lens_str + if not all(isinstance(x, int) for x in xs): + return seq_lens_str + first = xs[0] + if len(xs) >= 2 and all(x == first for x in xs): + return f"{first} x {len(xs)}" + return seq_lens_str + + +def _case_detail(c: dict) -> str: + if c["kind"] == "fixed": + return "Fixed" + seq = _fmt_seq_lens(c["seq_lens"]) + if seq.startswith("["): + return f"Varlen, `seq_lens`={seq}" + return f"Varlen, `seq_lens`=`{seq}`" + + +def _fmt_ms(x: float) -> str: + return f"{x:.4f}" + + +def _fmt_speedup(flash: float, fla: float) -> str: + if flash <= 0: + return "—" + return f"{fla / flash:.2f}×" + + +def _complete_cases(raw: list[dict]) -> list[dict]: + return [ + c for c in raw + if all(c[k] is not None for k in ( + "flash_fwd_ms", "fla_fwd_ms", "flash_fwdbwd_ms", "fla_fwdbwd_ms")) + ] + + +def _render_table_block(cases: list[dict]) -> list[str]: + lines = [ + "| 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 |", + "|------|------------------:|------------------:|--------:|" + "------------------:|------------------:|--------:|", + ] + for c in cases: + cell = _case_detail(c).replace("|", "\\|") + lines.append( + f"| {cell} | {_fmt_ms(c['flash_fwd_ms'])} | {_fmt_ms(c['fla_fwd_ms'])} |" + f" {_fmt_speedup(c['flash_fwd_ms'], c['fla_fwd_ms'])} |" + f" {_fmt_ms(c['flash_fwdbwd_ms'])} | {_fmt_ms(c['fla_fwdbwd_ms'])} |" + f" {_fmt_speedup(c['flash_fwdbwd_ms'], c['fla_fwdbwd_ms'])} |" + ) + lines.append("") + return lines + + +def render_markdown(sections, generated_at, generator_cmd, device_label) -> str: + title = "# KDA training benchmark (fwd / fwd+bwd)" + if device_label: + title += f" ({device_label})" + + lines = [title, "", f"- Generated: {generated_at}", ""] + + if not sections: + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + lines.append(f"- Command: `{generator_cmd}`") + lines.append("") + + first_cases = next((c for c in sections if c), None) + c0 = first_cases[0] if first_cases else None + if c0 is not None: + lines.append( + f"- Benchmark settings: `warmup={c0['warmup']}`, `iters={c0['iters']}`, " + f"`repeats={c0['repeats']}`" + ) + lines.append("") + lines.append(FLA_CHUNK_KDA_OPTIONS_MD) + lines.append( + "- `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`)" + ) + lines.append("") + + for cases in sections: + if not cases: + continue + c0 = cases[0] + lines.append(f"### `T={c0['T']}`, `H={c0['H']}`, `D={c0['D']}`") + lines.append("") + lines.extend(_render_table_block(cases)) + + return "\n".join(lines).rstrip() + "\n" + + +def main() -> None: + p = argparse.ArgumentParser( + description="Run bench_train.py and write a benchmark markdown report." + ) + p.add_argument("-o", "--output", type=Path, default=DEFAULT_OUT, + help=f"Output markdown path (default: {DEFAULT_OUT})") + p.add_argument("--device-label", default=DEFAULT_DEVICE_LABEL, + help=f"Device/platform label for the report title (default: {DEFAULT_DEVICE_LABEL!r})") + args, bench_extra = p.parse_known_args() + + def _fmt_generator_cmd(extra: list[str]) -> str: + cmd = "python benchmarks/generate_train_benchmark_md.py" + if args.output != DEFAULT_OUT: + cmd += f" -o {args.output}" + if args.device_label != DEFAULT_DEVICE_LABEL: + cmd += f" --device-label {args.device_label}" + tail = " ".join(extra) + return f"{cmd} {tail}".strip() if tail else cmd + + def _argv_with_h(argv: list[str], h: int) -> list[str]: + out: list[str] = [] + i = 0 + while i < len(argv): + a = argv[i] + if a == "--H" and i + 1 < len(argv): + i += 2 + continue + if a.startswith("--H="): + i += 1 + continue + out.append(a) + i += 1 + out.extend(["--H", str(h)]) + return out + + stdout_a = run_bench(list(bench_extra)) + stdout_b = run_bench(_argv_with_h(bench_extra, 64)) + cases_a = _complete_cases(parse_stdout(stdout_a)) + cases_b = _complete_cases(parse_stdout(stdout_b)) + + sections = [cases_a, cases_b] + + if not cases_a or not cases_b: + sys.stderr.write( + "Warning: missing complete benchmark rows for one or both runs.\n" + ) + + generated = _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%d") + md = render_markdown(sections, generated, _fmt_generator_cmd(bench_extra), args.device_label) + out_path = args.output.resolve() + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(md, encoding="utf-8") + print(f"Wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/csrc/train/binding.cpp b/csrc/train/binding.cpp new file mode 100644 index 0000000..4d53e73 --- /dev/null +++ b/csrc/train/binding.cpp @@ -0,0 +1,91 @@ +#include + +#include "train_ops.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("kda_gate_chunk_cumsum", &kda_gate_chunk_cumsum, "KDA gate activation + chunk cumsum (CUDA)", + py::arg("g"), py::arg("A_log"), py::arg("dt_bias"), + py::arg("out"), py::arg("scale"), py::arg("has_scale"), + py::arg("lower_bound"), py::arg("use_lower_bound"), + py::arg("chunk_size"), py::arg("cu_seqlens"), py::arg("chunk_indices")); + m.def("chunk_local_cumsum", &chunk_local_cumsum, "Chunk-local (reverse) cumsum (CUDA)", + py::arg("g"), py::arg("out"), py::arg("scale"), py::arg("has_scale"), + py::arg("reverse"), py::arg("chunk_size"), + py::arg("cu_seqlens"), py::arg("chunk_indices")); + m.def("kda_gate_bwd", &kda_gate_bwd, "KDA gate backward (CUDA)", + py::arg("g"), py::arg("A_log"), py::arg("dt_bias"), + py::arg("dyg"), py::arg("dg"), py::arg("dA_partial"), + py::arg("lower_bound"), py::arg("use_lower_bound")); + + m.def("chunk_kda_fwd_intra_sub_chunk", &chunk_kda_fwd_intra_sub_chunk, + "KDA intra sub-chunk diagonal blocks (safe_gate path)", + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("beta"), + py::arg("Aqk"), py::arg("Akkd"), py::arg("scale"), py::arg("chunk_size"), + py::arg("cu_seqlens") = py::none(), py::arg("chunk_indices") = py::none()); + m.def("chunk_kda_fwd_intra_token_parallel", &chunk_kda_fwd_intra_token_parallel, + "KDA intra token-parallel diagonal blocks (non safe_gate path)", + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("beta"), + py::arg("Aqk"), py::arg("Akkd"), py::arg("scale"), py::arg("chunk_size"), + py::arg("cu_seqlens") = py::none()); + m.def("chunk_kda_fwd_inter_solve_fused", &chunk_kda_fwd_inter_solve_fused, + "KDA intra off-diagonal blocks + merged tril solve", + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("beta"), + py::arg("Aqk"), py::arg("Akkd"), py::arg("Akk"), + py::arg("scale"), py::arg("chunk_size"), py::arg("safe_gate"), + py::arg("cu_seqlens") = py::none(), py::arg("chunk_indices") = py::none()); + + m.def("recompute_w_u_fwd", &recompute_w_u_fwd, "Recompute w/u/qg/kg for KDA forward (CUDA)", + py::arg("k"), py::arg("v"), py::arg("beta"), py::arg("A"), py::arg("gk"), + py::arg("q") = py::none(), + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_indices") = py::none()); + + m.def("chunk_gla_fwd_o_gk", &chunk_gla_fwd_o_gk, "Chunked GLA/KDA forward output (CUDA)", + py::arg("q"), py::arg("v"), py::arg("g"), py::arg("A"), py::arg("h"), + py::arg("scale"), + py::arg("state_v_first") = false, + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_size") = 64, + py::arg("chunk_indices") = py::none()); + + m.def("chunk_gated_delta_rule_fwd_h", &chunk_gated_delta_rule_fwd_h, + "KDA chunked state forward h (CUDA)", + py::arg("kg"), py::arg("w"), py::arg("u"), py::arg("gk"), + py::arg("initial_state") = py::none(), + py::arg("output_final_state") = false, + py::arg("chunk_size") = 64, + py::arg("state_v_first") = false, + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_offsets") = py::none(), + py::arg("nt_total") = 0); + m.def("chunk_gated_delta_rule_bwd_dhu", &chunk_gated_delta_rule_bwd_dhu, + "KDA chunked state backward dhu (CUDA)", + py::arg("qg"), py::arg("kg"), py::arg("w"), py::arg("gk"), + py::arg("do_"), py::arg("dv"), + py::arg("h0") = py::none(), + py::arg("dht") = py::none(), + py::arg("scale") = 1.0, + py::arg("chunk_size") = 64, + py::arg("state_v_first") = false, + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_offsets") = py::none(), + py::arg("nt_total") = 0); + + m.def("chunk_kda_bwd_dAv", &chunk_kda_bwd_dAv, "KDA backward dAqk + intra dv (CUDA)", + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("do_"), py::arg("A"), + py::arg("scale"), py::arg("cu_seqlens") = py::none(), + py::arg("chunk_indices") = py::none(), py::arg("chunk_size") = 64); + m.def("chunk_kda_bwd_wy_dqkg_fused", &chunk_kda_bwd_wy_dqkg_fused, + "KDA fused backward dq/dk/dv2/dg/db/dAkk (CUDA)", + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("v_new"), + py::arg("g"), py::arg("beta"), py::arg("A"), py::arg("h"), + py::arg("do"), py::arg("dh"), py::arg("dv"), + py::arg("scale"), py::arg("state_v_first"), + py::arg("cu_seqlens"), py::arg("chunk_indices"), py::arg("chunk_size")); + m.def("chunk_kda_bwd_intra", &chunk_kda_bwd_intra, "KDA backward intra-chunk dq/dk/db/dg (CUDA)", + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("beta"), + py::arg("dAqk"), py::arg("dAkk"), + py::arg("dq"), py::arg("dk"), py::arg("db"), py::arg("dg"), + py::arg("safe_gate"), py::arg("cu_seqlens") = py::none(), + py::arg("chunk_indices") = py::none(), py::arg("chunk_size") = 64); +} diff --git a/csrc/train/bwd_dav.cu b/csrc/train/bwd_dav.cu new file mode 100644 index 0000000..787dddd --- /dev/null +++ b/csrc/train/bwd_dav.cu @@ -0,0 +1,357 @@ +// KDA backward: dAqk (attention gradient matrix) and the intra-chunk part of dv. +// Replicates fla/ops/kda/chunk_bwd.py::chunk_kda_bwd_kernel_dAv. +// +// Per (chunk, head) block: +// dAqk[i,j] = scale * sum_v do[i,v] * v_new[j,v] for i >= j (fp32) +// dv[i,:] = sum_j tril(Aqk)[j,i] * do[j,:] (stored in do dtype) +// All MMAs run on tensor cores (SM80 16x8x16 bf16/fp16 atoms, fp32 accumulators). +// +// Data movement: gmem tiles are staged with 16B cp.async (masked rows zero-filled +// via src-size 0), smem rows padded +8 halves against bank conflicts, MMA operand +// fragments load with ldmatrix (x4 for row-major tiles, .trans for strided +// views), and gmem stores go through smem staging for 16B coalesced writes. +// GEMM/accumulation order and rounding points are unchanged from v1. + +#include +#include +#include + +#include + +#include "common.cuh" + +namespace kda_impl { + +using namespace cute; + +constexpr int kBT = 64; // chunk size (KDA bwd is only ever run with 64) +constexpr int kBV = 64; // V tile, matches fla's CONST_TILING on non-Hopper +constexpr int kThreads = 128; +constexpr int kPad = 8; // smem row padding in halves (16B) +constexpr int kCP = kBT + kPad; // padded row stride of every 64-wide tile + +template struct MmaAtom; +template <> struct MmaAtom { using type = SM80_16x8x16_F32BF16BF16F32_TN; }; +template <> struct MmaAtom { using type = SM80_16x8x16_F32F16F16F32_TN; }; + +__device__ __forceinline__ void cp_async16(void* dst, void const* src, bool full) { + uint32_t s = cute::cast_smem_ptr_to_uint(dst); + int sz = full ? 16 : 0; // src-size 0 zero-fills, used for masked rows/cols + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" + :: "r"(s), "l"(src), "r"(sz) : "memory"); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory"); +} + +// s[r][c] = g[r*row_stride + c] via 16B cp.async; rows past rows_valid and col +// chunks at/past cols_valid zero-filled. Requires 16B-aligned rows (row_stride +// and cols_valid multiples of 8 halves). +template +__device__ __forceinline__ void stage_tile(T* s, T const* g, int64_t row_stride, + int rows_valid, int cols_valid, int tid) { + constexpr int kCG = kBT / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + bool full = (r < rows_valid) && (c < cols_valid); + cp_async16(s + r * kCP + c, full ? g + (int64_t)r * row_stride + c : g, full); + } +} + +// scalar fallback for irregular V (not a multiple of 8) +template +__device__ __forceinline__ void stage_tile_scalar(T* s, T const* g, int64_t row_stride, + int rows_valid, int cols_valid, int tid) { + for (int idx = tid; idx < kBT * kBT; idx += kThreads) { + int r = idx / kBT, c = idx % kBT; + s[r * kCP + c] = (r < rows_valid && c < cols_valid) ? g[(int64_t)r * row_stride + c] : T(0.0f); + } +} + +// Grid: (NT, B*HV). +template +__global__ void __launch_bounds__(kThreads) chunk_kda_bwd_dav_kernel( + T const* __restrict__ A, + T const* __restrict__ v, + T const* __restrict__ do_, + T* __restrict__ dv, + float* __restrict__ dA, + float scale, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int T_len, int HV, int V +) { + int64_t const i_t = blockIdx.x; + int64_t const i_bh = blockIdx.y; + int64_t const i_hv = i_bh % HV; + + int64_t bos, t0; + int64_t seq_len = T_len; + if (IS_VARLEN) { + int64_t const i_n = chunk_indices[i_t * 2]; + int64_t const i_tl = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + seq_len = cu_seqlens[i_n + 1] - bos; + t0 = i_tl * kBT; + } else { + bos = (i_bh / HV) * (int64_t)T_len; + t0 = i_t * kBT; + } + int64_t const rem = seq_len - t0; + int const rows_valid = int(rem < (int64_t)kBT ? rem : (int64_t)kBT); + if (rows_valid <= 0) return; + + T const* A_base = A + (bos * HV + i_hv) * kBT; + T const* v_base = v + (bos * HV + i_hv) * V; + T const* do_base = do_ + (bos * HV + i_hv) * V; + T* dv_base = dv + (bos * HV + i_hv) * V; + float* dA_base = dA + (bos * HV + i_hv) * kBT; + + // dv fragments are staged through sV (dead by then); the fp32 dA tile + // overlays all three tiles after the V loop. + __shared__ union { + struct { + T sDo[kBT * kCP]; // do tile: [r][v] + T sV[kBV * kCP]; // v_new tile: [j][v] + T sA[kBT * kCP]; // Aqk tile: [r][a], tril-masked in smem + } t; + float sDA[kBT * kBT]; + } sm; + + int const tid = threadIdx.x; + bool const vec_ok = (V % 8) == 0; + + // A tile: natural [r][a] staging, then apply the tril/seq column mask in smem + // (v1 staged the transposed view with strided scalar gmem reads instead). + if (vec_ok) { + stage_tile(sm.t.sA, A_base + t0 * (int64_t)HV * kBT, (int64_t)HV * kBT, rows_valid, kBT, tid); + cp_async_commit(); + cp_async_wait<0>(); + } else { + stage_tile_scalar(sm.t.sA, A_base + t0 * (int64_t)HV * kBT, (int64_t)HV * kBT, rows_valid, kBT, tid); + } + __syncthreads(); + for (int idx = tid; idx < kBT * kBT; idx += kThreads) { + int r = idx / kBT, a = idx % kBT; + if (a > r || t0 + a >= seq_len) sm.t.sA[r * kCP + a] = T(0.0f); + } + + using Atom = typename MmaAtom::type; + auto mma = make_tiled_mma(Atom{}, Layout>{}, Tile<_64, _64, _16>{}); + auto thr_mma = mma.get_thread_slice(tid); + + Copy_Atom ldsm_n; + Copy_Atom ldsm_t; + auto s2r_a = make_tiled_copy_A(ldsm_n, mma); // row-major (M,K) tiles + auto s2r_at = make_tiled_copy_A(ldsm_t, mma); // strided (M,K) views + auto s2r_b = make_tiled_copy_B(ldsm_n, mma); // row-major (N,K) tiles + auto s2r_bt = make_tiled_copy_B(ldsm_t, mma); // strided (N,K) views + auto thr_s2r_a = s2r_a.get_thread_slice(tid); + auto thr_s2r_at = s2r_at.get_thread_slice(tid); + auto thr_s2r_b = s2r_b.get_thread_slice(tid); + auto thr_s2r_bt = s2r_bt.get_thread_slice(tid); + + Tensor sDo_rm = make_tensor(make_smem_ptr(sm.t.sDo), Layout, Stride, _1>>{}); + Tensor sV_rm = make_tensor(make_smem_ptr(sm.t.sV), Layout, Stride, _1>>{}); + Tensor sA_rm = make_tensor(make_smem_ptr(sm.t.sA), Layout, Stride, _1>>{}); + // strided views over the same tiles + Tensor sA_st = make_tensor(make_smem_ptr(sm.t.sA), Layout, Stride<_1, Int>>{}); // (M=a, K=r) + Tensor sDo_st = make_tensor(make_smem_ptr(sm.t.sDo), Layout, Stride<_1, Int>>{}); // (N=v, K=r) + + Tensor cC = make_identity_tensor(Shape<_64, _64>{}); + Tensor tCcC = thr_mma.partition_C(cC); + + Tensor fdA = thr_mma.make_fragment_C(tCcC); + clear(fdA); + + // acc += sA_view[64(M),64(K)] @ sB_view[64(N),64(K)]^T, K sliced 16-wide; + // copies/mmas run k-block by k-block in ascending order (v1 numerics). + auto gemm_rm_rm = [&](auto& acc, auto const& sA_t, auto const& sB_t) { + Tensor tCrA = thr_mma.partition_fragment_A(sA_t); + Tensor tCrB = thr_mma.partition_fragment_B(sB_t); + Tensor tXsA = thr_s2r_a.partition_S(sA_t); + Tensor tXsB = thr_s2r_b.partition_S(sB_t); + Tensor tXrA = thr_s2r_a.retile_D(tCrA); + Tensor tXrB = thr_s2r_b.retile_D(tCrB); + constexpr int KB = decltype(size<2>(tXsA))::value; + CUTE_UNROLL + for (int kb = 0; kb < KB; ++kb) { + copy(s2r_a, tXsA(_, _, kb), tXrA(_, _, kb)); + copy(s2r_b, tXsB(_, _, kb), tXrB(_, _, kb)); + gemm(mma, tCrA(_, _, kb), tCrB(_, _, kb), acc); + } + }; + auto gemm_st_st = [&](auto& acc, auto const& sA_t, auto const& sB_t) { + Tensor tCrA = thr_mma.partition_fragment_A(sA_t); + Tensor tCrB = thr_mma.partition_fragment_B(sB_t); + Tensor tXsA = thr_s2r_at.partition_S(sA_t); + Tensor tXsB = thr_s2r_bt.partition_S(sB_t); + Tensor tXrA = thr_s2r_at.retile_D(tCrA); + Tensor tXrB = thr_s2r_bt.retile_D(tCrB); + constexpr int KB = decltype(size<2>(tXsA))::value; + CUTE_UNROLL + for (int kb = 0; kb < KB; ++kb) { + copy(s2r_at, tXsA(_, _, kb), tXrA(_, _, kb)); + copy(s2r_bt, tXsB(_, _, kb), tXrB(_, _, kb)); + gemm(mma, tCrA(_, _, kb), tCrB(_, _, kb), acc); + } + }; + + for (int v0 = 0; v0 < V; v0 += kBV) { + int const cols_valid = min(kBV, V - v0); + if (vec_ok) { + stage_tile(sm.t.sDo, do_base + t0 * (int64_t)HV * V + v0, (int64_t)HV * V, rows_valid, cols_valid, tid); + stage_tile(sm.t.sV, v_base + t0 * (int64_t)HV * V + v0, (int64_t)HV * V, rows_valid, cols_valid, tid); + cp_async_commit(); + cp_async_wait<0>(); + } else { + stage_tile_scalar(sm.t.sDo, do_base + t0 * (int64_t)HV * V + v0, (int64_t)HV * V, rows_valid, cols_valid, tid); + stage_tile_scalar(sm.t.sV, v_base + t0 * (int64_t)HV * V + v0, (int64_t)HV * V, rows_valid, cols_valid, tid); + } + __syncthreads(); + + // fdA[t, j] += sum_v do[t,v] * v_new[j,v] + gemm_rm_rm(fdA, sDo_rm, sV_rm); + + // fdv[a, v] = sum_r tril(Aqk)^T[a,r] * do[r,v] + Tensor fdv = thr_mma.make_fragment_C(tCcC); + clear(fdv); + gemm_st_st(fdv, sA_st, sDo_st); + __syncthreads(); // sV is dead now; reuse it as the dv staging tile + + for (int i = 0; i < size(fdv); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + sm.t.sV[m * kCP + n] = T(fdv(i)); + } + __syncthreads(); + if (vec_ok) { + constexpr int kCG = kBV / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + if (r < rows_valid && c < cols_valid) { + *reinterpret_cast(dv_base + (t0 + r) * (int64_t)HV * V + v0 + c) = + *reinterpret_cast(sm.t.sV + r * kCP + c); + } + } + } else { + for (int idx = tid; idx < kBT * kBV; idx += kThreads) { + int r = idx / kBV, vv = idx % kBV; + if (r < rows_valid && vv < cols_valid) { + dv_base[(t0 + r) * (int64_t)HV * V + v0 + vv] = sm.t.sV[r * kCP + vv]; + } + } + } + __syncthreads(); + } + + // dA[t, j] = (t >= j) ? fdA * scale : 0, fp32, staged for 16B stores + for (int i = 0; i < size(fdA); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + sm.sDA[m * kBT + n] = (m >= n) ? fdA(i) * scale : 0.0f; + } + __syncthreads(); + { + constexpr int kCG = kBT / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 4; + if (r < rows_valid) { + *reinterpret_cast(dA_base + (t0 + r) * (int64_t)HV * kBT + c) = + *reinterpret_cast(sm.sDA + r * kBT + c); + } + } + } +} + +template +void launch_dav( + T const* A, T const* v, T const* do_, T* dv, float* dA, + float scale, + int64_t const* cu_seqlens, int64_t const* chunk_indices, + int64_t NT, int64_t B, int64_t T_len, int64_t HV, int64_t V, + cudaStream_t stream +) { + dim3 grid((unsigned)NT, (unsigned)(B * HV)); + dim3 block(kThreads); + if (cu_seqlens) { + chunk_kda_bwd_dav_kernel<<>>( + A, v, do_, dv, dA, scale, cu_seqlens, chunk_indices, (int)T_len, (int)HV, (int)V); + } else { + chunk_kda_bwd_dav_kernel<<>>( + A, v, do_, dv, dA, scale, nullptr, nullptr, (int)T_len, (int)HV, (int)V); + } +} + +} // namespace kda_impl + +using kda_impl::launch_dav; +using kda_impl::kBT; + +std::tuple chunk_kda_bwd_dAv( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor do_, + torch::Tensor A, + double scale, + std::optional cu_seqlens, + std::optional chunk_indices, + int64_t chunk_size +) { + TORCH_CHECK(chunk_size == kBT, "chunk_kda_bwd_dAv only supports chunk_size 64"); + TORCH_CHECK(do_.is_cuda() && do_.is_contiguous(), "do must be contiguous CUDA tensor"); + TORCH_CHECK(v.is_cuda() && v.is_contiguous() && v.sizes() == do_.sizes()); + TORCH_CHECK(A.is_cuda() && A.is_contiguous() && A.scalar_type() == do_.scalar_type()); + TORCH_CHECK(do_.dim() == 4, "do must be [B, T, HV, V]"); + + int64_t B = do_.size(0), T = do_.size(1), HV = do_.size(2), V = do_.size(3); + TORCH_CHECK(A.dim() == 4 && A.size(0) == B && A.size(1) == T && A.size(2) == HV && A.size(3) == kBT); + + int64_t const* cu_ptr = nullptr; + int64_t const* ci_ptr = nullptr; + int64_t NT = (T + kBT - 1) / kBT; + if (cu_seqlens.has_value()) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_indices.has_value(), "chunk_indices must be provided with cu_seqlens"); + auto const& cu = cu_seqlens.value(); + auto const& ci = chunk_indices.value(); + TORCH_CHECK(cu.is_cuda() && cu.is_contiguous() && cu.dtype() == torch::kLong); + TORCH_CHECK(ci.is_cuda() && ci.is_contiguous() && ci.dtype() == torch::kLong); + TORCH_CHECK(ci.dim() == 2 && ci.size(1) == 2); + cu_ptr = cu.data_ptr(); + ci_ptr = ci.data_ptr(); + NT = ci.size(0); + } + + auto dA = torch::empty({B, T, HV, kBT}, do_.options().dtype(torch::kFloat32)); + auto dv = torch::empty_like(do_); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + if (do_.scalar_type() == at::kBFloat16) { + launch_dav( + reinterpret_cast(A.data_ptr()), + reinterpret_cast(v.data_ptr()), + reinterpret_cast(do_.data_ptr()), + reinterpret_cast(dv.data_ptr()), + dA.data_ptr(), (float)scale, cu_ptr, ci_ptr, NT, B, T, HV, V, stream); + } else if (do_.scalar_type() == at::kHalf) { + launch_dav( + reinterpret_cast(A.data_ptr()), + reinterpret_cast(v.data_ptr()), + reinterpret_cast(do_.data_ptr()), + reinterpret_cast(dv.data_ptr()), + dA.data_ptr(), (float)scale, cu_ptr, ci_ptr, NT, B, T, HV, V, stream); + } else { + TORCH_CHECK(false, "chunk_kda_bwd_dAv: unsupported dtype ", do_.scalar_type()); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dA, dv}; +} diff --git a/csrc/train/bwd_dav_binding.cpp b/csrc/train/bwd_dav_binding.cpp new file mode 100644 index 0000000..a26f70c --- /dev/null +++ b/csrc/train/bwd_dav_binding.cpp @@ -0,0 +1,23 @@ +#include + +// fla/ops/kda/chunk_bwd.py::chunk_kda_bwd_dAv. +// q, k are accepted for signature parity with the Triton host (the kernel does not read them). +// v is v_new, A is Aqk ([B,T,HV,64], same dtype as do). Returns (dA fp32 [B,T,HV,64], dv like do). +std::tuple chunk_kda_bwd_dAv( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor do_, + torch::Tensor A, + double scale, + std::optional cu_seqlens, + std::optional chunk_indices, + int64_t chunk_size +); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("chunk_kda_bwd_dAv", &chunk_kda_bwd_dAv, "KDA backward dAqk + intra dv (CUDA)", + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("do_"), py::arg("A"), + py::arg("scale"), py::arg("cu_seqlens") = py::none(), + py::arg("chunk_indices") = py::none(), py::arg("chunk_size") = 64); +} diff --git a/csrc/train/bwd_intra.cu b/csrc/train/bwd_intra.cu new file mode 100644 index 0000000..a279946 --- /dev/null +++ b/csrc/train/bwd_intra.cu @@ -0,0 +1,777 @@ +// KDA backward: intra-chunk (sub-chunk) gradients accumulated into dq/dk/db/dg. +// Replicates fla/ops/kda/chunk_intra.py::chunk_kda_bwd_kernel_intra. +// +// Each block handles one (k-slice i_k, sub-chunk i_i, chunk i_t, batch*head) tile of +// shape [BC=16, BK=64]. The kernel has three parts, matching the Triton source: +// (a) contributions from previous sub-chunks of the chunk (query/key side) +// (b) the diagonal sub-chunk, i<=j masked side +// (c) the reverse (key side) contributions, incl. the i>=j masked diagonal side +// All MMAs run on tensor cores (SM80 16x8x8 tf32 atoms, fp32 accumulators); the Triton +// reference computes these dots from fp32 inputs with tf32 precision. +// +// BK=64 with 256 threads (8 warps, one 1x8 tiled mma) halves the dAqk/dAkk gmem +// traffic per output element versus BK=32 and keeps all staging loads/stores on +// 16B vectors. Per-element math and accumulation order are unchanged from the +// BK=32 version (each output element is computed by the identical op sequence). + +#include +#include +#include + +#include +#include + +#include "common.cuh" + +namespace kda_impl { + +using namespace cute; + +constexpr int kBT = 64; // chunk size (KDA bwd is only ever run with 64) +constexpr int kBC = 16; // sub-chunk size +constexpr int kBK = 64; // K tile +constexpr int kThreads = 256; + +__device__ __forceinline__ cutlass::tfloat32_t to_tf32(float x) { + return cutlass::tfloat32_t(x); +} + +__device__ __forceinline__ uint4 pack_tf32(float a, float b, float c, float d) { + return make_uint4(to_tf32(a).storage, to_tf32(b).storage, to_tf32(c).storage, to_tf32(d).storage); +} + +// --------------------------------------------------------------------------- +// staging helpers (vector path requires K % kBK == 0 so every row segment is +// in-bounds and 16B-aligned; the scalar fallbacks keep arbitrary K working) + +// fp32 gmem tile [kBC][kBK] -> fp32 smem, rows past rows_valid / cols past +// cols_valid zero-filled (vec path requires the full row segment in-bounds). +__device__ __forceinline__ void stage_f32(float* s, float const* g, int64_t row_stride, + int rows_valid, int cols_valid, bool vec_ok, int tid) { + if (vec_ok) { + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int r = idx / kCG; + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (r < rows_valid) v = *reinterpret_cast(g + (int64_t)r * row_stride + (idx % kCG) * 4); + *reinterpret_cast(s + idx * 4) = v; + } + } else { + for (int idx = tid; idx < kBC * kBK; idx += kThreads) { + int r = idx / kBK, c = idx % kBK; + s[idx] = (r < rows_valid && c < cols_valid) ? g[(int64_t)r * row_stride + c] : 0.f; + } + } +} + +// bf16/fp16 gmem tile [kBC][kBK] -> fp32 smem (converted), zero-filled past +// rows_valid / cols_valid. +template +__device__ __forceinline__ void stage_t_f32(float* s, T const* g, int64_t row_stride, + int rows_valid, int cols_valid, bool vec_ok, int tid) { + if (vec_ok) { + constexpr int kCG = kBK / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + float vals[8] = {}; + if (r < rows_valid) { + uint4 raw = *reinterpret_cast(g + (int64_t)r * row_stride + c); + T const* h = reinterpret_cast(&raw); + CUTE_UNROLL + for (int j = 0; j < 8; ++j) vals[j] = to_f32(h[j]); + } + float4* d = reinterpret_cast(s + r * kBK + c); + d[0] = make_float4(vals[0], vals[1], vals[2], vals[3]); + d[1] = make_float4(vals[4], vals[5], vals[6], vals[7]); + } + } else { + for (int idx = tid; idx < kBC * kBK; idx += kThreads) { + int r = idx / kBK, c = idx % kBK; + s[idx] = (r < rows_valid && c < cols_valid) ? to_f32(g[(int64_t)r * row_stride + c]) : 0.f; + } + } +} + +// dA tile [kBC][kBC] fp32 gmem -> tf32 smem, natural orientation: +// s[r][jj] = dA[r * row_stride + col0 + jj], rows past rows_valid zero-filled. +__device__ __forceinline__ void stage_dA(cutlass::tfloat32_t* sA, float const* dA, + int64_t row_stride, int col0, int rows_valid, int tid) { + CUTE_UNROLL + for (int idx = tid; idx < kBC * (kBC / 4); idx += kThreads) { + int r = idx / (kBC / 4), c = (idx % (kBC / 4)) * 4; + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (r < rows_valid) v = *reinterpret_cast(dA + (int64_t)r * row_stride + col0 + c); + *reinterpret_cast(sA + r * kBC + c) = pack_tf32(v.x, v.y, v.z, v.w); + } +} + +// dA tile transposed: s[r][jj] = dA[jj * row_stride + col0 + r], so the gmem +// reads coalesce along r (the tile's column dim). Rows (jj) past rows_valid +// are zero-filled. +__device__ __forceinline__ void stage_dA_T(cutlass::tfloat32_t* sA, float const* dA, + int64_t row_stride, int col0, int rows_valid, int tid) { + CUTE_UNROLL + for (int idx = tid; idx < kBC * (kBC / 4); idx += kThreads) { + int jj = idx / (kBC / 4), r = (idx % (kBC / 4)) * 4; + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (jj < rows_valid) v = *reinterpret_cast(dA + (int64_t)jj * row_stride + col0 + r); + sA[r * kBC + jj] = to_tf32(v.x); + sA[(r + 1) * kBC + jj] = to_tf32(v.y); + sA[(r + 2) * kBC + jj] = to_tf32(v.z); + sA[(r + 3) * kBC + jj] = to_tf32(v.w); + } +} + +// Grid: (NK*NC, NT, B*HV). +template +__global__ void __launch_bounds__(kThreads, 3) chunk_kda_bwd_intra_kernel( + T const* __restrict__ q, + T const* __restrict__ k, + float const* __restrict__ g, + float const* __restrict__ beta, + float const* __restrict__ dAqk, + float const* __restrict__ dAkk, + float const* __restrict__ dq, + float* __restrict__ dq2, + float const* __restrict__ dk, + float* __restrict__ dk2, + float const* __restrict__ dg, + float* __restrict__ dg2, + float* __restrict__ db2, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int64_t allT, + int T_len, int H, int HV, int K, int NC +) { + int const i_kc = blockIdx.x; + int64_t const i_t = blockIdx.y; + int64_t const i_bh = blockIdx.z; + int const i_k = i_kc / NC, i_i = i_kc % NC; + int64_t const i_hv = i_bh % HV; + int64_t const i_h = i_hv / (HV / H); + + int64_t bos, t_chunk; + int64_t seq_len = T_len; + if (IS_VARLEN) { + int64_t const i_n = chunk_indices[i_t * 2]; + int64_t const i_tl = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + seq_len = cu_seqlens[i_n + 1] - bos; + t_chunk = i_tl * kBT; + } else { + bos = (i_bh / HV) * (int64_t)T_len; + t_chunk = i_t * kBT; + } + int const seq = (int)seq_len; + + int const i_ti = (int)t_chunk + i_i * kBC; + if (i_ti >= seq) return; + + int64_t const HK = (int64_t)H * K; + int64_t const HVK = (int64_t)HV * K; + int64_t const HVBT = (int64_t)HV * kBT; + + T const* q_base = q + (bos * H + i_h) * K; + T const* k_base = k + (bos * H + i_h) * K; + float const* g_base = g + (bos * HV + i_hv) * K; + float const* beta_base = beta + bos * HV + i_hv; + float const* dAqk_base = dAqk + (bos * HV + i_hv) * kBT; + float const* dAkk_base = dAkk + (bos * HV + i_hv) * kBT; + float const* dq_base = dq + (bos * HV + i_hv) * K; + float* dq2_base = dq2 + (bos * HV + i_hv) * K; + float const* dk_base = dk + (bos * HV + i_hv) * K; + float* dk2_base = dk2 + (bos * HV + i_hv) * K; + float const* dg_base = dg + (bos * HV + i_hv) * K; + float* dg2_base = dg2 + (bos * HV + i_hv) * K; + float* db2_base = db2 + (i_kc / NC * allT + bos) * HV + i_hv; + + __shared__ __align__(16) float s_g[kBC * kBK]; // this sub-chunk's g rows (fp32) + __shared__ __align__(16) float s_q[kBC * kBK]; + __shared__ __align__(16) float s_k[kBC * kBK]; + __shared__ float s_beta[kBC]; + __shared__ __align__(16) float s_dq2[kBC * kBK]; // fp32 staging tiles for elementwise phases + __shared__ __align__(16) float s_dk2[kBC * kBK]; + __shared__ __align__(16) float s_dg2[kBC * kBK]; + // s_dq2 is dead after the dq2/beta-scale phase below, so the (c)-loop B2 + // operand and the final dkt staging tile overlap it (saves 8KB -> 3 CTAs/SM). + __shared__ __align__(16) cutlass::tfloat32_t s_A[kBC * kBC]; // MMA A operand (M=c, K=j) + __shared__ __align__(16) cutlass::tfloat32_t s_A2[kBC * kBC]; + __shared__ __align__(16) cutlass::tfloat32_t s_B[kBC * kBK]; // MMA B operand, stored [j][n], viewed (N=n, K=j) + float* s_dkt = s_dq2; + cutlass::tfloat32_t* s_B2 = reinterpret_cast(s_dq2); + __shared__ float s_colA[kBC], s_colB[kBC]; // scalar-loop staging (non-safe paths) + __shared__ float s_rowq[kBK], s_rowk[kBK], s_rowg[kBK]; + + int const tid = threadIdx.x; + int const col0 = i_k * kBK; + bool const vec_ok = (K % kBK) == 0; + int const rows_valid_c = min(kBC, seq - i_ti); // i_ti < seq, so >= 1 + int const cols_valid = min(kBK, K - col0); + + stage_f32(s_g, g_base + (int64_t)i_ti * HVK + col0, HVK, rows_valid_c, cols_valid, vec_ok, tid); + stage_t_f32(s_q, q_base + (int64_t)i_ti * HK + col0, HK, rows_valid_c, cols_valid, vec_ok, tid); + stage_t_f32(s_k, k_base + (int64_t)i_ti * HK + col0, HK, rows_valid_c, cols_valid, vec_ok, tid); + if (tid < kBC) { + int t = i_ti + tid; + s_beta[tid] = (t < seq) ? beta_base[(int64_t)t * HV] : 0.0f; + } + __syncthreads(); + + // 8 warps split the N=64 dimension of the [16,64] output tile. + auto mma = make_tiled_mma(SM80_16x8x8_F32TF32TF32F32_TN{}, Layout>{}); + auto thr_mma = mma.get_thread_slice(tid); + + Tensor sA_t = make_tensor(make_smem_ptr(s_A), Layout, Stride<_16, _1>>{}); + Tensor sA2_t = make_tensor(make_smem_ptr(s_A2), Layout, Stride<_16, _1>>{}); + Tensor sB_t = make_tensor(make_smem_ptr(s_B), Layout, _16>, Stride<_1, Int>>{}); + Tensor sB2_t = make_tensor(make_smem_ptr(s_B2), Layout, _16>, Stride<_1, Int>>{}); + + // Identity tensors give the (m,k)/(n,k)/(m,n) coordinate of each fragment element; + // fragments themselves are shaped from the fp32 smem staging tiles. + Tensor cA = make_identity_tensor(Shape<_16, _16>{}); + Tensor cB = make_identity_tensor(Shape, _16>{}); + Tensor cC = make_identity_tensor(Shape<_16, Int>{}); + Tensor tCcA = thr_mma.partition_A(cA); + Tensor tCcB = thr_mma.partition_B(cB); + Tensor tCcC = thr_mma.partition_C(cC); + + Tensor sC_t = make_tensor(make_smem_ptr(s_dq2), Layout>, Stride, _1>>{}); + + Tensor tCrA = thr_mma.partition_fragment_A(sA_t); + Tensor tCrA2 = thr_mma.partition_fragment_A(sA2_t); + Tensor tCrB = thr_mma.partition_fragment_B(sB_t); + Tensor tCrB2 = thr_mma.partition_fragment_B(sB2_t); + + Tensor fdq2 = thr_mma.partition_fragment_C(sC_t); + Tensor fdk2 = thr_mma.partition_fragment_C(sC_t); + Tensor fdkt = thr_mma.partition_fragment_C(sC_t); + Tensor ftmp = thr_mma.partition_fragment_C(sC_t); + Tensor ftmp2 = thr_mma.partition_fragment_C(sC_t); + clear(fdq2); + clear(fdk2); + clear(fdkt); + + auto load_frag = [](auto& frag, auto const& coords, auto const& s_t) { + for (int i = 0; i < size(frag); ++i) + frag(i) = s_t(get<0>(coords(i)), get<1>(coords(i))); + }; + + // (a) contributions from previous sub-chunks: dq2/dk2 += dA[c,j] @ (k_j * 2^(gn-g_j)) + if (i_i > 0) { + for (int i_j = 0; i_j < i_i; ++i_j) { + int const j0 = (int)t_chunk + i_j * kBC; + int const rows_valid_j = min(kBC, seq - j0); + stage_dA(s_A, dAqk_base + (int64_t)i_ti * HVBT + i_j * kBC, HVBT, 0, rows_valid_c, tid); + stage_dA(s_A2, dAkk_base + (int64_t)i_ti * HVBT + i_j * kBC, HVBT, 0, rows_valid_c, tid); + if (vec_ok) { + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int jj = idx / kCG, n = (idx % kCG) * 4; + int tj = j0 + jj; + float kv[4] = {}, gv[4] = {}; + if (tj < seq) { + uint2 rawk = *reinterpret_cast(k_base + (int64_t)tj * HK + col0 + n); + T const* hk = reinterpret_cast(&rawk); + float4 g4 = *reinterpret_cast(g_base + (int64_t)tj * HVK + col0 + n); + kv[0] = to_f32(hk[0]); kv[1] = to_f32(hk[1]); kv[2] = to_f32(hk[2]); kv[3] = to_f32(hk[3]); + gv[0] = g4.x; gv[1] = g4.y; gv[2] = g4.z; gv[3] = g4.w; + } + float4 gn = *reinterpret_cast(s_g + n); // s_g row 0 is g[i_ti] (gn) + *reinterpret_cast(s_B + jj * kBK + n) = pack_tf32( + kv[0] * exp2f(gn.x - gv[0]), kv[1] * exp2f(gn.y - gv[1]), + kv[2] * exp2f(gn.z - gv[2]), kv[3] * exp2f(gn.w - gv[3])); + } + } else { + for (int idx = tid; idx < kBC * kBK; idx += kThreads) { + int jj = idx / kBK, c = idx % kBK; + int tj = j0 + jj, col = col0 + c; + bool valid = (tj < seq) && (col < K); + float kv = valid ? to_f32(k_base[(int64_t)tj * HK + col]) : 0.0f; + float gv = valid ? g_base[(int64_t)tj * HVK + col] : 0.0f; + s_B[idx] = to_tf32(kv * exp2f(s_g[c] - gv)); + } + } + __syncthreads(); + load_frag(tCrA, tCcA, sA_t); + load_frag(tCrA2, tCcA, sA2_t); + load_frag(tCrB, tCcB, sB_t); + cute::gemm(mma, tCrA, tCrB, fdq2); + cute::gemm(mma, tCrA2, tCrB, fdk2); + __syncthreads(); + } + for (int i = 0; i < size(fdq2); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + float e = exp2f(s_g[m * kBK + n] - s_g[n]); + fdq2(i) *= e; + fdk2(i) *= e; + } + } + + // (b) diagonal sub-chunk, i <= j masked side + int const mid = min(kBC / 2, seq - i_ti - 1); // safe_gate midpoint reference row + if (SAFE_GATE) { + { + // masked dA loads: keep = (r >= jj) && both tokens valid + CUTE_UNROLL + for (int idx = tid; idx < kBC * (kBC / 4); idx += kThreads) { + int r = idx / (kBC / 4), c = (idx % (kBC / 4)) * 4; + float4 v1 = make_float4(0.f, 0.f, 0.f, 0.f), v2 = make_float4(0.f, 0.f, 0.f, 0.f); + if (i_ti + r < seq) { + v1 = *reinterpret_cast(dAqk_base + (int64_t)(i_ti + r) * HVBT + i_i * kBC + c); + v2 = *reinterpret_cast(dAkk_base + (int64_t)(i_ti + r) * HVBT + i_i * kBC + c); + } + float va1[4] = {v1.x, v1.y, v1.z, v1.w}, va2[4] = {v2.x, v2.y, v2.z, v2.w}; + CUTE_UNROLL + for (int j = 0; j < 4; ++j) { + bool keep = (r >= c + j) && (i_ti + r < seq) && (i_ti + c + j < seq); + va1[j] = keep ? va1[j] : 0.f; + va2[j] = keep ? va2[j] : 0.f; + } + *reinterpret_cast(s_A + r * kBC + c) = pack_tf32(va1[0], va1[1], va1[2], va1[3]); + *reinterpret_cast(s_A2 + r * kBC + c) = pack_tf32(va2[0], va2[1], va2[2], va2[3]); + } + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int j = idx / kCG, c = (idx % kCG) * 4; + float4 kj = *reinterpret_cast(s_k + j * kBK + c); + float4 gj = *reinterpret_cast(s_g + j * kBK + c); + float4 gm = *reinterpret_cast(s_g + mid * kBK + c); + float o[4] = {0.f, 0.f, 0.f, 0.f}; + if (i_ti + j < seq) { + o[0] = kj.x * exp2f(-(gj.x - gm.x)); + o[1] = kj.y * exp2f(-(gj.y - gm.y)); + o[2] = kj.z * exp2f(-(gj.z - gm.z)); + o[3] = kj.w * exp2f(-(gj.w - gm.w)); + } + *reinterpret_cast(s_B + j * kBK + c) = pack_tf32(o[0], o[1], o[2], o[3]); + } + } + __syncthreads(); + load_frag(tCrA, tCcA, sA_t); + load_frag(tCrA2, tCcA, sA2_t); + load_frag(tCrB, tCcB, sB_t); + clear(ftmp); + clear(ftmp2); + cute::gemm(mma, tCrA, tCrB, ftmp); + cute::gemm(mma, tCrA2, tCrB, ftmp2); + for (int i = 0; i < size(fdq2); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + float e = (i_ti + m < seq) ? exp2f(s_g[m * kBK + n] - s_g[mid * kBK + n]) : 0.0f; + fdq2(i) += ftmp(i) * e; + fdk2(i) += ftmp2(i) * e; + } + __syncthreads(); + } else { + int const jmax = min(kBC, seq - i_ti); + for (int j = 0; j < jmax; ++j) { + if (tid < kBC) { + int t = i_ti + tid; + int64_t off = (int64_t)t * HVBT + i_i * kBC + j; + s_colA[tid] = (t < seq) ? dAqk_base[off] : 0.0f; + s_colB[tid] = (t < seq) ? dAkk_base[off] : 0.0f; + } else if (tid < kBC + kBK) { + int c = tid - kBC; + int col = col0 + c; + s_rowk[c] = (col < K) ? to_f32(k_base[(int64_t)(i_ti + j) * HK + col]) : 0.0f; + s_rowg[c] = (col < K) ? g_base[(int64_t)(i_ti + j) * HVK + col] : 0.0f; + } + __syncthreads(); + for (int i = 0; i < size(fdq2); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + if (m >= j) { + float e = exp2f(s_g[m * kBK + n] - s_rowg[n]); + fdq2(i) += s_colA[m] * s_rowk[n] * e; + fdk2(i) += s_colB[m] * s_rowk[n] * e; + } + } + __syncthreads(); + } + } + + // db is the row-sum of dk2*k BEFORE the beta scaling; dq2/dg2 use pre-add values. + for (int i = 0; i < size(fdq2); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + s_dq2[m * kBK + n] = fdq2(i); + s_dk2[m * kBK + n] = fdk2(i); + } + __syncthreads(); + if (tid < kBC) { + float acc = 0.0f; + CUTE_UNROLL + for (int c4 = 0; c4 < kBK / 4; ++c4) { + float4 dv = *reinterpret_cast(s_dk2 + tid * kBK + c4 * 4); + float4 kv = *reinterpret_cast(s_k + tid * kBK + c4 * 4); + acc += dv.x * kv.x; + acc += dv.y * kv.y; + acc += dv.z * kv.z; + acc += dv.w * kv.w; + } + int t = i_ti + tid; + if (t < seq) db2_base[(int64_t)t * HV] = acc; + } + __syncthreads(); // db must read dk2 before the beta scaling below overwrites it + if (vec_ok) { + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 4; + int t = i_ti + r; + float4 dq2v = *reinterpret_cast(s_dq2 + r * kBK + c); + float4 qv = *reinterpret_cast(s_q + r * kBK + c); + float4 dk2v = *reinterpret_cast(s_dk2 + r * kBK + c); + *reinterpret_cast(s_dg2 + r * kBK + c) = + make_float4(qv.x * dq2v.x, qv.y * dq2v.y, qv.z * dq2v.z, qv.w * dq2v.w); + if (t < seq) { + int64_t off = (int64_t)t * HVK + col0 + c; + float4 dqv = *reinterpret_cast(dq_base + off); + *reinterpret_cast(dq2_base + off) = make_float4( + dq2v.x + dqv.x, dq2v.y + dqv.y, dq2v.z + dqv.z, dq2v.w + dqv.w); + } + float b = s_beta[r]; + *reinterpret_cast(s_dk2 + r * kBK + c) = + make_float4(dk2v.x * b, dk2v.y * b, dk2v.z * b, dk2v.w * b); + } + } else { + for (int idx = tid; idx < kBC * kBK; idx += kThreads) { + int r = idx / kBK, c = idx % kBK; + int t = i_ti + r, col = col0 + c; + bool valid = (t < seq) && (col < K); + s_dg2[idx] = s_q[idx] * s_dq2[idx]; + if (valid) dq2_base[(int64_t)t * HVK + col] = s_dq2[idx] + dq_base[(int64_t)t * HVK + col]; + s_dk2[idx] *= s_beta[r]; + } + } + __syncthreads(); + + // (c) reverse (key side) contributions. + int const NC_eff = min(NC, (seq - (int)t_chunk + kBC - 1) / kBC); + if (i_i < NC_eff - 1) { + int const r3 = min(i_ti + kBC, seq) - 1 - i_ti; // row of this sub-chunk's last valid token + for (int i_j = i_i + 1; i_j < NC_eff; ++i_j) { + int const j0 = (int)t_chunk + i_j * kBC; + int const rows_valid_j = min(kBC, seq - j0); + stage_dA_T(s_A, dAqk_base + (int64_t)j0 * HVBT + i_i * kBC, HVBT, 0, rows_valid_j, tid); + stage_dA_T(s_A2, dAkk_base + (int64_t)j0 * HVBT + i_i * kBC, HVBT, 0, rows_valid_j, tid); + if (vec_ok) { + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int jj = idx / kCG, n = (idx % kCG) * 4; + int tj = j0 + jj; + float e[4] = {}, qv[4] = {}, kv[4] = {}, bv = 0.f; + if (tj < seq) { + float4 g4 = *reinterpret_cast(g_base + (int64_t)tj * HVK + col0 + n); + uint2 rawq = *reinterpret_cast(q_base + (int64_t)tj * HK + col0 + n); + uint2 rawk = *reinterpret_cast(k_base + (int64_t)tj * HK + col0 + n); + T const* hq = reinterpret_cast(&rawq); + T const* hk = reinterpret_cast(&rawk); + float4 gr3 = *reinterpret_cast(s_g + r3 * kBK + n); + e[0] = exp2f(g4.x - gr3.x); e[1] = exp2f(g4.y - gr3.y); + e[2] = exp2f(g4.z - gr3.z); e[3] = exp2f(g4.w - gr3.w); + qv[0] = to_f32(hq[0]); qv[1] = to_f32(hq[1]); qv[2] = to_f32(hq[2]); qv[3] = to_f32(hq[3]); + kv[0] = to_f32(hk[0]); kv[1] = to_f32(hk[1]); kv[2] = to_f32(hk[2]); kv[3] = to_f32(hk[3]); + bv = beta_base[(int64_t)tj * HV]; + } + *reinterpret_cast(s_B + jj * kBK + n) = pack_tf32( + qv[0] * e[0], qv[1] * e[1], qv[2] * e[2], qv[3] * e[3]); + *reinterpret_cast(s_B2 + jj * kBK + n) = pack_tf32( + kv[0] * bv * e[0], kv[1] * bv * e[1], kv[2] * bv * e[2], kv[3] * bv * e[3]); + } + } else { + for (int idx = tid; idx < kBC * kBK; idx += kThreads) { + int jj = idx / kBK, c = idx % kBK; + int tj = j0 + jj, col = col0 + c; + bool valid = (tj < seq) && (col < K); + float e = 0.0f, qv = 0.0f, kv = 0.0f, bv = 0.0f; + if (valid) { + e = exp2f(g_base[(int64_t)tj * HVK + col] - s_g[r3 * kBK + c]); + qv = to_f32(q_base[(int64_t)tj * HK + col]); + kv = to_f32(k_base[(int64_t)tj * HK + col]); + bv = beta_base[(int64_t)tj * HV]; + } + s_B[idx] = to_tf32(qv * e); + s_B2[idx] = to_tf32(kv * bv * e); + } + } + __syncthreads(); + load_frag(tCrA, tCcA, sA_t); + load_frag(tCrA2, tCcA, sA2_t); + load_frag(tCrB, tCcB, sB_t); + load_frag(tCrB2, tCcB, sB2_t); + cute::gemm(mma, tCrA, tCrB, fdkt); + cute::gemm(mma, tCrA2, tCrB2, fdkt); + __syncthreads(); + } + for (int i = 0; i < size(fdkt); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + fdkt(i) *= exp2f(s_g[r3 * kBK + n] - s_g[m * kBK + n]); + } + } + + // (c) diagonal sub-chunk, i >= j masked side + if (SAFE_GATE) { + { + // masked transposed dA loads: keep = (r <= jj) && both tokens valid; + // gmem coalesces along r + CUTE_UNROLL + for (int idx = tid; idx < kBC * (kBC / 4); idx += kThreads) { + int jj = idx / (kBC / 4), r = (idx % (kBC / 4)) * 4; + float4 v1 = make_float4(0.f, 0.f, 0.f, 0.f), v2 = make_float4(0.f, 0.f, 0.f, 0.f); + if (i_ti + jj < seq) { + v1 = *reinterpret_cast(dAqk_base + (int64_t)(i_ti + jj) * HVBT + i_i * kBC + r); + v2 = *reinterpret_cast(dAkk_base + (int64_t)(i_ti + jj) * HVBT + i_i * kBC + r); + } + float va1[4] = {v1.x, v1.y, v1.z, v1.w}, va2[4] = {v2.x, v2.y, v2.z, v2.w}; + CUTE_UNROLL + for (int j = 0; j < 4; ++j) { + bool keep = (r + j <= jj) && (i_ti + r + j < seq) && (i_ti + jj < seq); + va1[j] = keep ? va1[j] : 0.f; + va2[j] = keep ? va2[j] : 0.f; + } + s_A[r * kBC + jj] = to_tf32(va1[0]); + s_A[(r + 1) * kBC + jj] = to_tf32(va1[1]); + s_A[(r + 2) * kBC + jj] = to_tf32(va1[2]); + s_A[(r + 3) * kBC + jj] = to_tf32(va1[3]); + s_A2[r * kBC + jj] = to_tf32(va2[0]); + s_A2[(r + 1) * kBC + jj] = to_tf32(va2[1]); + s_A2[(r + 2) * kBC + jj] = to_tf32(va2[2]); + s_A2[(r + 3) * kBC + jj] = to_tf32(va2[3]); + } + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int j = idx / kCG, c = (idx % kCG) * 4; + float4 qv = *reinterpret_cast(s_q + j * kBK + c); + float4 kv = *reinterpret_cast(s_k + j * kBK + c); + float4 gj = *reinterpret_cast(s_g + j * kBK + c); + float4 gm = *reinterpret_cast(s_g + mid * kBK + c); + float e[4] = {0.f, 0.f, 0.f, 0.f}; + if (i_ti + j < seq) { + e[0] = exp2f(gj.x - gm.x); + e[1] = exp2f(gj.y - gm.y); + e[2] = exp2f(gj.z - gm.z); + e[3] = exp2f(gj.w - gm.w); + } + float b = s_beta[j]; + *reinterpret_cast(s_B + j * kBK + c) = pack_tf32( + qv.x * e[0], qv.y * e[1], qv.z * e[2], qv.w * e[3]); + *reinterpret_cast(s_B2 + j * kBK + c) = pack_tf32( + kv.x * b * e[0], kv.y * b * e[1], kv.z * b * e[2], kv.w * b * e[3]); + } + } + __syncthreads(); + load_frag(tCrA, tCcA, sA_t); + load_frag(tCrA2, tCcA, sA2_t); + load_frag(tCrB, tCcB, sB_t); + load_frag(tCrB2, tCcB, sB2_t); + clear(ftmp); + clear(ftmp2); + cute::gemm(mma, tCrA, tCrB, ftmp); + cute::gemm(mma, tCrA2, tCrB2, ftmp2); + for (int i = 0; i < size(fdkt); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + float en = (i_ti + m < seq) ? exp2f(-(s_g[m * kBK + n] - s_g[mid * kBK + n])) : 0.0f; + fdkt(i) += (ftmp(i) + ftmp2(i)) * en; + } + __syncthreads(); + } else { + int const jmax = min(kBC, seq - i_ti); + for (int j = 0; j < jmax; ++j) { + if (tid < kBC) { + int64_t off = (int64_t)(i_ti + j) * HVBT + i_i * kBC + tid; + s_colA[tid] = dAqk_base[off]; + s_colB[tid] = dAkk_base[off]; + } else if (tid < kBC + kBK) { + int c = tid - kBC; + int col = col0 + c; + s_rowq[c] = (col < K) ? to_f32(q_base[(int64_t)(i_ti + j) * HK + col]) : 0.0f; + s_rowk[c] = (col < K) ? to_f32(k_base[(int64_t)(i_ti + j) * HK + col]) : 0.0f; + s_rowg[c] = (col < K) ? g_base[(int64_t)(i_ti + j) * HVK + col] : 0.0f; + } + __syncthreads(); + float const bj = s_beta[j]; + for (int i = 0; i < size(fdkt); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + if (m <= j) { + float e = exp2f(s_rowg[n] - s_g[m * kBK + n]); + fdkt(i) += s_colA[m] * s_rowq[n] * e + s_colB[m] * s_rowk[n] * bj * e; + } + } + __syncthreads(); + } + } + + // Epilogue: dk2 = beta*dk2 + dk_in + dkt; dg2 = q*dq2 + (beta*dk2 - dkt)*k + dg_in + for (int i = 0; i < size(fdkt); ++i) { + int m = get<0>(tCcC(i)), n = get<1>(tCcC(i)); + s_dkt[m * kBK + n] = fdkt(i); + } + __syncthreads(); + if (vec_ok) { + constexpr int kCG = kBK / 4; + CUTE_UNROLL + for (int idx = tid; idx < kBC * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 4; + int t = i_ti + r; + if (t >= seq) continue; + int64_t off = (int64_t)t * HVK + col0 + c; + float4 dg2v = *reinterpret_cast(s_dg2 + r * kBK + c); + float4 dk2v = *reinterpret_cast(s_dk2 + r * kBK + c); + float4 dktv = *reinterpret_cast(s_dkt + r * kBK + c); + float4 kv = *reinterpret_cast(s_k + r * kBK + c); + float4 dgv_in = *reinterpret_cast(dg_base + off); + float4 dkv_in = *reinterpret_cast(dk_base + off); + *reinterpret_cast(dg2_base + off) = make_float4( + dg2v.x + (dk2v.x - dktv.x) * kv.x + dgv_in.x, + dg2v.y + (dk2v.y - dktv.y) * kv.y + dgv_in.y, + dg2v.z + (dk2v.z - dktv.z) * kv.z + dgv_in.z, + dg2v.w + (dk2v.w - dktv.w) * kv.w + dgv_in.w); + *reinterpret_cast(dk2_base + off) = make_float4( + dk2v.x + dkv_in.x + dktv.x, dk2v.y + dkv_in.y + dktv.y, + dk2v.z + dkv_in.z + dktv.z, dk2v.w + dkv_in.w + dktv.w); + } + } else { + for (int idx = tid; idx < kBC * kBK; idx += kThreads) { + int r = idx / kBK, c = idx % kBK; + int t = i_ti + r, col = col0 + c; + bool valid = (t < seq) && (col < K); + if (!valid) continue; + int64_t off = (int64_t)t * HVK + col; + float dgv = s_dg2[idx] + (s_dk2[idx] - s_dkt[idx]) * s_k[idx] + dg_base[off]; + float dkv = s_dk2[idx] + dk_base[off] + s_dkt[idx]; + dg2_base[off] = dgv; + dk2_base[off] = dkv; + } + } +} + +template +void launch_intra( + T const* q, T const* k, float const* g, float const* beta, + float const* dAqk, float const* dAkk, + float const* dq, float* dq2, float const* dk, float* dk2, + float const* dg, float* dg2, float* db2, + int64_t const* cu_seqlens, int64_t const* chunk_indices, + int64_t allT, int64_t NT, int64_t B, int64_t T_len, + int H, int HV, int K, int NC, bool safe_gate, + cudaStream_t stream +) { + int const NK = (K + kBK - 1) / kBK; + dim3 grid((unsigned)(NK * NC), (unsigned)NT, (unsigned)(B * HV)); + dim3 block(kThreads); + + #define LAUNCH_INTRA(IS_VARLEN, SAFE_GATE) \ + chunk_kda_bwd_intra_kernel<<>>( \ + q, k, g, beta, dAqk, dAkk, dq, dq2, dk, dk2, dg, dg2, db2, \ + cu_seqlens, chunk_indices, allT, (int)T_len, H, HV, K, NC) + + if (cu_seqlens) { + if (safe_gate) { LAUNCH_INTRA(true, true); } else { LAUNCH_INTRA(true, false); } + } else { + if (safe_gate) { LAUNCH_INTRA(false, true); } else { LAUNCH_INTRA(false, false); } + } + #undef LAUNCH_INTRA +} + +} // namespace kda_impl + +using kda_impl::launch_intra; +using kda_impl::kBT; +using kda_impl::kBC; +using kda_impl::kBK; + +// fla/ops/kda/chunk_intra.py::chunk_kda_bwd_intra host wrapper. +// dq/dk/db/dg are the upstream (fused kernel) gradients; the kernel accumulates the +// intra-chunk parts and the host wrapper reduces db2 over the NK dim and adds db, +// exactly like the Triton host. Returns (dq2, dk2, db_out, dg2), all fp32. +std::tuple chunk_kda_bwd_intra( + torch::Tensor q, + torch::Tensor k, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor dAqk, + torch::Tensor dAkk, + torch::Tensor dq, + torch::Tensor dk, + torch::Tensor db, + torch::Tensor dg, + bool safe_gate, + std::optional cu_seqlens, + std::optional chunk_indices, + int64_t chunk_size +) { + TORCH_CHECK(chunk_size == kBT, "chunk_kda_bwd_intra only supports chunk_size 64"); + TORCH_CHECK(q.is_cuda() && q.is_contiguous() && q.dim() == 4, "q must be [B, T, H, K]"); + TORCH_CHECK(k.is_cuda() && k.is_contiguous() && k.sizes() == q.sizes()); + TORCH_CHECK(k.scalar_type() == q.scalar_type()); + + int64_t B = k.size(0), T = k.size(1), H = k.size(2), K = k.size(3); + int64_t HV = g.size(2); + TORCH_CHECK(HV % H == 0, "HV must be a multiple of H"); + TORCH_CHECK(g.is_cuda() && g.is_contiguous() && g.dtype() == torch::kFloat32); + TORCH_CHECK(g.dim() == 4 && g.size(0) == B && g.size(1) == T && g.size(3) == K); + for (auto const& t : {beta, db}) { + TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.dtype() == torch::kFloat32); + TORCH_CHECK(t.dim() == 3 && t.size(0) == B && t.size(1) == T && t.size(2) == HV); + } + for (auto const& t : {dAqk, dAkk}) { + TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.dtype() == torch::kFloat32); + TORCH_CHECK(t.dim() == 4 && t.size(0) == B && t.size(1) == T && t.size(2) == HV && t.size(3) == kBT); + } + for (auto const& t : {dq, dk, dg}) { + TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.dtype() == torch::kFloat32); + TORCH_CHECK(t.sizes() == g.sizes()); + } + + int64_t const* cu_ptr = nullptr; + int64_t const* ci_ptr = nullptr; + int64_t NT = (T + kBT - 1) / kBT; + if (cu_seqlens.has_value()) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_indices.has_value(), "chunk_indices must be provided with cu_seqlens"); + auto const& cu = cu_seqlens.value(); + auto const& ci = chunk_indices.value(); + TORCH_CHECK(cu.is_cuda() && cu.is_contiguous() && cu.dtype() == torch::kLong); + TORCH_CHECK(ci.is_cuda() && ci.is_contiguous() && ci.dtype() == torch::kLong); + TORCH_CHECK(ci.dim() == 2 && ci.size(1) == 2); + cu_ptr = cu.data_ptr(); + ci_ptr = ci.data_ptr(); + NT = ci.size(0); + } + + int64_t const NK = (K + kBK - 1) / kBK; + int64_t const NC = kBT / kBC; + auto dq2 = torch::empty_like(dq); + auto dk2 = torch::empty_like(dk); + auto dg2 = torch::empty_like(dg); + auto db2 = torch::empty({NK, B, T, HV}, dq.options().dtype(torch::kFloat32)); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + if (k.scalar_type() == at::kBFloat16) { + launch_intra( + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + g.data_ptr(), beta.data_ptr(), dAqk.data_ptr(), dAkk.data_ptr(), + dq.data_ptr(), dq2.data_ptr(), dk.data_ptr(), dk2.data_ptr(), + dg.data_ptr(), dg2.data_ptr(), db2.data_ptr(), + cu_ptr, ci_ptr, B * T, NT, B, T, (int)H, (int)HV, (int)K, (int)NC, safe_gate, stream); + } else if (k.scalar_type() == at::kHalf) { + launch_intra( + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + g.data_ptr(), beta.data_ptr(), dAqk.data_ptr(), dAkk.data_ptr(), + dq.data_ptr(), dq2.data_ptr(), dk.data_ptr(), dk2.data_ptr(), + dg.data_ptr(), dg2.data_ptr(), db2.data_ptr(), + cu_ptr, ci_ptr, B * T, NT, B, T, (int)H, (int)HV, (int)K, (int)NC, safe_gate, stream); + } else { + TORCH_CHECK(false, "chunk_kda_bwd_intra: unsupported dtype ", k.scalar_type()); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + auto db_out = db2.sum(0).add_(db); + return {dq2, dk2, db_out, dg2}; +} diff --git a/csrc/train/bwd_intra_binding.cpp b/csrc/train/bwd_intra_binding.cpp new file mode 100644 index 0000000..2114b76 --- /dev/null +++ b/csrc/train/bwd_intra_binding.cpp @@ -0,0 +1,31 @@ +#include + +// fla/ops/kda/chunk_intra.py::chunk_kda_bwd_intra. +// dq/dk/db/dg are the upstream (fused kernel) gradients; the kernel accumulates the +// intra-chunk parts and the host wrapper reduces db2 over the NK dim and adds db, +// exactly like the Triton host. Returns (dq2, dk2, db_out, dg2), all fp32. +std::tuple chunk_kda_bwd_intra( + torch::Tensor q, + torch::Tensor k, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor dAqk, + torch::Tensor dAkk, + torch::Tensor dq, + torch::Tensor dk, + torch::Tensor db, + torch::Tensor dg, + bool safe_gate, + std::optional cu_seqlens, + std::optional chunk_indices, + int64_t chunk_size +); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("chunk_kda_bwd_intra", &chunk_kda_bwd_intra, "KDA backward intra-chunk dq/dk/db/dg (CUDA)", + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("beta"), + py::arg("dAqk"), py::arg("dAkk"), + py::arg("dq"), py::arg("dk"), py::arg("db"), py::arg("dg"), + py::arg("safe_gate"), py::arg("cu_seqlens") = py::none(), + py::arg("chunk_indices") = py::none(), py::arg("chunk_size") = 64); +} diff --git a/csrc/train/bwd_wy_dqkg.cu b/csrc/train/bwd_wy_dqkg.cu new file mode 100644 index 0000000..445113c --- /dev/null +++ b/csrc/train/bwd_wy_dqkg.cu @@ -0,0 +1,615 @@ +// Fused KDA backward kernel producing dq/dk/dv2/dg/db/dAkk for one chunk. +// Replicates fla/ops/kda/chunk_bwd.py::chunk_kda_bwd_kernel_wy_dqkg_fused +// (host wrapper at chunk_bwd.py:366-431). +// +// One CTA handles one (chunk, batch*head) pair: BT=64 token rows, tiled over +// K in BK=64 blocks and V in BV=64 blocks. All GEMMs run on tensor cores via +// CuTe SM80 16x8x16 mma atoms with fp32 accumulators, matching Triton's +// input-dtype mma + fp32 accumulate semantics. + +#include +#include +#include + +#include + +#include "common.cuh" + +namespace { + +constexpr int kBT = 64; // chunk size +constexpr int kBK = 64; // K tile +constexpr int kBV = 64; // V tile +constexpr int kThreads = 256; // 8 warps, one 4x2 tiled mma; keeps the four + // fp32 accumulators register-resident (16 vals + // per thread each, no local-memory spills) +constexpr int kPad = 8; // smem row padding in elements (16B) against bank conflicts +constexpr int kCP = 64 + kPad; // padded row stride of every 64-wide tile + +// --------------------------------------------------------------------------- +// cp.async helpers + +__device__ __forceinline__ void cp_async16(void* dst, void const* src, bool full) { + uint32_t s = cute::cast_smem_ptr_to_uint(dst); + int sz = full ? 16 : 0; // src-size 0 zero-fills, used for masked rows + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" + :: "r"(s), "l"(src), "r"(sz) : "memory"); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory"); +} + +template +struct FusedSmem { + static constexpr int kTile = kBT * kCP; + // persistent + T A_t[kTile]; // Akk transposed: A_t[c][r] = Akk[t0+r][c] + // g/k staging for the current i_k, persistent across the V loop so the + // loads overlap the V-loop staging/compute + union { + T raw[3 * kTile]; + struct { + float g[kBT * 68]; // fp32 [t][k], padded to 68 floats (16B) per row + T k[kTile]; + } gk; + } sp; + float beta[kBT]; + float db[kBT]; + float dgk[64]; + float gn[64]; + union { + // V-loop tiles. h/dh are stored as [k][v] (B-operand K-major). + struct { + T do_[kTile]; T v_new[kTile]; T dv[kTile]; T h[kTile]; T dh[kTile]; + T v[kTile]; // only used when i_k == 0 + } vl; + // epilogue tiles: kg/dw as [t][k], dwT as [k][t]; q is staged into dw + // after the epilogue GEMMs. Physically overlaps the vl tiles, which + // are all dead by this phase. + struct { T kg[kTile]; T dw[kTile]; T dwT[kTile]; } ep; + // dA postprocess: masked dA and the first product transposed + struct { T dAm[kTile]; T c1T[kTile]; } pp; + }; +}; + +// s[r][c] = g[r*row_stride + c] via 16B cp.async, rows past rows_valid +// zero-filled. Requires 16B-aligned rows (row_stride and base offset multiples +// of 8 elements). +template +__device__ __forceinline__ void stage_tile(T* s, T const* g, T const* g_safe, int64_t row_stride, + int rows_valid, int tid) { + constexpr int kCG = 64 / 8; + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + bool full = r < rows_valid; + cp_async16(s + r * kCP + c, full ? g + (int64_t)r * row_stride + c : g_safe, full); + } +} + +// h/dh state tile stored as s[kk][vv]; k/v dims are always fully valid. +// state_v_first=false: gmem is [K, V], s[kk][vv] = g[(k0+kk)*ld + v0+vv], cp.async +// state_v_first=true: gmem is [V, K], s[kk][vv] = g[(v0+vv)*ld + k0+kk], scalar +template +__device__ __forceinline__ void stage_state_tile(T* s, T const* g, int64_t ld, int64_t k0, int64_t v0, + bool state_v_first, int tid) { + if (!state_v_first) { + constexpr int kCG = 64 / 8; + for (int idx = tid; idx < 64 * kCG; idx += kThreads) { + int kk = idx / kCG, vv = (idx % kCG) * 8; + cp_async16(s + kk * kCP + vv, g + (k0 + kk) * ld + v0 + vv, true); + } + } else { + for (int idx = tid; idx < 64 * 64; idx += kThreads) { + int kk = idx >> 6, vv = idx & 63; + s[kk * kCP + vv] = g[(v0 + vv) * ld + k0 + kk]; + } + } +} + +template +__global__ void __launch_bounds__(kThreads, 1) chunk_kda_bwd_wy_dqkg_fused_kernel( + T const* __restrict__ q_g, T const* __restrict__ k_g, + T const* __restrict__ v_g, T const* __restrict__ v_new_g, + float const* __restrict__ g_g, float const* __restrict__ beta_g, + T const* __restrict__ A_g, T const* __restrict__ h_g, + T const* __restrict__ do_g, T const* __restrict__ dh_g, T const* __restrict__ dv_g, + float* __restrict__ dq_g, float* __restrict__ dk_g, T* __restrict__ dv2_g, + float* __restrict__ dg_g, float* __restrict__ db_g, float* __restrict__ dA_g, + int64_t const* __restrict__ cu_seqlens, int64_t const* __restrict__ chunk_indices, + float scale, int T_len, int H, int HV, + bool state_v_first, bool is_varlen +) { + using namespace cute; + + constexpr int NK = K / kBK; + constexpr int NV = V / kBV; + const int tid = threadIdx.x; + + int64_t i_t = blockIdx.x; + const int i_bh = blockIdx.y; + const int i_b = i_bh / HV; + const int i_hv = i_bh % HV; + const int i_h = i_hv / (HV / H); + + int64_t i_tg, bos; + int Tl; + if (is_varlen) { + i_tg = i_t; + int64_t i_n = chunk_indices[i_t * 2]; + i_t = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + Tl = int(cu_seqlens[i_n + 1] - bos); + } else { + int NT = (T_len + kBT - 1) / kBT; + i_tg = (int64_t)i_b * NT + i_t; + bos = (int64_t)i_b * T_len; + Tl = T_len; + } + const int t0 = int(i_t) * kBT; + const int t_end = Tl < t0 + kBT ? Tl : t0 + kBT; + const int last = t_end - 1; // local row of the last valid token + const int rows_valid = t_end - t0; // valid rows in this chunk + + // base offsets (int64 everywhere, matching the Triton kernel) + const int64_t qk_base = (bos * H + i_h) * (int64_t)K; + const int64_t hvK_base = (bos * HV + i_hv) * (int64_t)K; + const int64_t hvV_base = (bos * HV + i_hv) * (int64_t)V; + const int64_t h_base = (i_tg * HV + i_hv) * (int64_t)K * V; + const int64_t A_base = (bos * HV + i_hv) * (int64_t)kBT; + const int64_t beta_base = bos * HV + i_hv; + const int64_t qk_row = (int64_t)H * K; + const int64_t hvK_row = (int64_t)HV * K; + const int64_t hvV_row = (int64_t)HV * V; + const int64_t A_row = (int64_t)HV * kBT; + + extern __shared__ __align__(128) unsigned char smem_raw[]; + using Smem = FusedSmem; + Smem& sm = *reinterpret_cast(smem_raw); + + // preload beta, db init, and Akk (transposed; the natural orientation is + // read through a strided-B view where needed) + if (tid < kBT) { + sm.db[tid] = 0.f; + sm.beta[tid] = (tid < rows_valid) ? beta_g[beta_base + (int64_t)(t0 + tid) * HV] : 0.f; + } + { + T const* Ap = A_g + A_base + (int64_t)t0 * A_row; + for (int idx = tid; idx < kBT * 8; idx += kThreads) { + int r = idx / 8, c = (idx % 8) * 8; + T vals[8]; + if (r < rows_valid) { + *reinterpret_cast(vals) = + *reinterpret_cast(Ap + (int64_t)r * A_row + c); + } else { + CUTE_UNROLL + for (int j = 0; j < 8; ++j) vals[j] = T(0.f); + } + CUTE_UNROLL + for (int j = 0; j < 8; ++j) sm.A_t[(c + j) * kCP + r] = vals[j]; + } + } + + using MmaOp = std::conditional_t, + SM80_16x8x16_F32BF16BF16F32_TN, + SM80_16x8x16_F32F16F16F32_TN>; + auto tiled_mma = make_tiled_mma(MMA_Atom{}, Layout>{}); + auto thr_mma = tiled_mma.get_thread_slice(tid); + Tensor cC = make_identity_tensor(Shape, _64>{}); + Tensor tCcC = thr_mma.partition_C(cC); + + // acc += sA[64,K-dim] @ sB[64,K-dim]^T with all tiles K-major in smem. + // Operand fragments load via ldmatrix; copies/mma run k-block by k-block to + // bound register pressure. + Copy_Atom ldsm_n; + // x2 variant for B operands: the per-warp B tile per k-block is too small + // for the x4 atom (same constraint as chunk_h.cu). + Copy_Atom ldsm_n2; + auto s2r_a = make_tiled_copy_A(ldsm_n, tiled_mma); + auto s2r_b = make_tiled_copy_B(ldsm_n2, tiled_mma); + auto thr_s2r_a = s2r_a.get_thread_slice(tid); + auto thr_s2r_b = s2r_b.get_thread_slice(tid); + auto gemm_sm = [&](auto& acc, T const* sA, T const* sB) { + Tensor sAt = make_tensor(make_smem_ptr(sA), Layout, Stride, _1>>{}); + Tensor sBt = make_tensor(make_smem_ptr(sB), Layout, Stride, _1>>{}); + Tensor tCrA = thr_mma.partition_fragment_A(sAt); + Tensor tCrB = thr_mma.partition_fragment_B(sBt); + Tensor tXsA = thr_s2r_a.partition_S(sAt); // (CPY, M, K) + Tensor tXsB = thr_s2r_b.partition_S(sBt); // (CPY, N, K) + Tensor tXrA = thr_s2r_a.retile_D(tCrA); + Tensor tXrB = thr_s2r_b.retile_D(tCrB); + constexpr int KB = decltype(size<2>(tXsA))::value; + CUTE_UNROLL + for (int kb = 0; kb < KB; ++kb) { + copy(s2r_a, tXsA(_, _, kb), tXrA(_, _, kb)); + copy(s2r_b, tXsB(_, _, kb), tXrB(_, _, kb)); + gemm(tiled_mma, tCrA(_, _, kb), tCrB(_, _, kb), acc); + } + }; + // Same, but the B operand is a contraction-strided view of a natural + // [K][N] tile: B[N][K] with stride (1, kCP), loaded with ldsm .trans. + Copy_Atom ldsm_t2; + auto s2r_bt = make_tiled_copy_B(ldsm_t2, tiled_mma); + auto thr_s2r_bt = s2r_bt.get_thread_slice(tid); + auto gemm_sm_bt = [&](auto& acc, T const* sA, T const* sB) { + Tensor sAt = make_tensor(make_smem_ptr(sA), Layout, Stride, _1>>{}); + Tensor sBt = make_tensor(make_smem_ptr(sB), Layout, Stride<_1, Int>>{}); + Tensor tCrA = thr_mma.partition_fragment_A(sAt); + Tensor tCrB = thr_mma.partition_fragment_B(sBt); + Tensor tXsA = thr_s2r_a.partition_S(sAt); + Tensor tXsB = thr_s2r_bt.partition_S(sBt); + Tensor tXrA = thr_s2r_a.retile_D(tCrA); + Tensor tXrB = thr_s2r_bt.retile_D(tCrB); + constexpr int KB = decltype(size<2>(tXsA))::value; + CUTE_UNROLL + for (int kb = 0; kb < KB; ++kb) { + copy(s2r_a, tXsA(_, _, kb), tXrA(_, _, kb)); + copy(s2r_bt, tXsB(_, _, kb), tXrB(_, _, kb)); + gemm(tiled_mma, tCrA(_, _, kb), tCrB(_, _, kb), acc); + } + }; + + Tensor acc_dA = partition_fragment_C(tiled_mma, Shape, _64>{}); + clear(acc_dA); + + for (int i_k = 0; i_k < NK; ++i_k) { + __syncthreads(); // separates this iteration's tile writes from the previous phase + if (tid < 64) { + sm.dgk[tid] = 0.f; + sm.gn[tid] = g_g[hvK_base + (int64_t)last * hvK_row + i_k * kBK + tid]; + } + // stage g/k for the epilogue phases; the loads overlap the V loop + { + float const* gsrc = g_g + hvK_base + (int64_t)t0 * hvK_row + i_k * kBK; + for (int idx = tid; idx < kBT * 16; idx += kThreads) { + int r = idx / 16, c = (idx % 16) * 4; + bool full = r < rows_valid; + cp_async16(sm.sp.gk.g + r * 68 + c, + full ? gsrc + (int64_t)r * hvK_row + c : g_g, full); + } + T const* ksrc = k_g + qk_base + (int64_t)t0 * qk_row + i_k * kBK; + stage_tile(sm.sp.gk.k, ksrc, k_g, qk_row, rows_valid, tid); + } + Tensor acc_dq = partition_fragment_C(tiled_mma, Shape, _64>{}); + Tensor acc_dk = partition_fragment_C(tiled_mma, Shape, _64>{}); + Tensor acc_dw = partition_fragment_C(tiled_mma, Shape, _64>{}); + clear(acc_dq); + clear(acc_dk); + clear(acc_dw); + + for (int i_v = 0; i_v < NV; ++i_v) { + const int64_t v0 = i_v * kBV; + T const* do_p = do_g + hvV_base + (int64_t)t0 * hvV_row + v0; + stage_tile(sm.vl.do_, do_p, do_g, hvV_row, rows_valid, tid); + T const* vn_p = v_new_g + hvV_base + (int64_t)t0 * hvV_row + v0; + stage_tile(sm.vl.v_new, vn_p, v_new_g, hvV_row, rows_valid, tid); + T const* dv_p = dv_g + hvV_base + (int64_t)t0 * hvV_row + v0; + stage_tile(sm.vl.dv, dv_p, dv_g, hvV_row, rows_valid, tid); + stage_state_tile(sm.vl.h, h_g + h_base, state_v_first ? K : V, i_k * kBK, v0, state_v_first, tid); + stage_state_tile(sm.vl.dh, dh_g + h_base, state_v_first ? K : V, i_k * kBK, v0, state_v_first, tid); + if (i_k == 0) { + T const* v_p = v_g + hvV_base + (int64_t)t0 * hvV_row + v0; + stage_tile(sm.vl.v, v_p, v_g, hvV_row, rows_valid, tid); + } + cp_async_commit(); + cp_async_wait<0>(); + __syncthreads(); + + // dgk[k] += sum_v h[k][v] * dh[k][v] (fp32) + { + int kk = tid & 63, part = tid >> 6; + constexpr int kSpan = 64 / (kThreads / 64); + float s = 0.f; + for (int vv = part * kSpan; vv < part * kSpan + kSpan; ++vv) + s += to_f32(sm.vl.h[kk * kCP + vv]) * to_f32(sm.vl.dh[kk * kCP + vv]); + atomicAdd(&sm.dgk[kk], s); + } + + gemm_sm(acc_dq, sm.vl.do_, sm.vl.h); // dq += do @ h + gemm_sm(acc_dk, sm.vl.v_new, sm.vl.dh); // dk += v_new @ dh + gemm_sm(acc_dw, sm.vl.dv, sm.vl.h); // dw += dv @ h + + if (i_k == 0) { + gemm_sm(acc_dA, sm.vl.dv, sm.vl.v); // dA += dv @ v^T + Tensor acc_dvb = partition_fragment_C(tiled_mma, Shape, _64>{}); + clear(acc_dvb); + gemm_sm_bt(acc_dvb, sm.A_t, sm.vl.dv); // dvb = Akk^T @ dv + T* dv2_p = dv2_g + hvV_base + (int64_t)t0 * hvV_row; + // in-thread partial sums per fragment row half (e bit1): one + // smem atomic per row instead of one per element + float db_acc[2] = {0.f, 0.f}; + CUTE_UNROLL + for (int e = 0; e < size(acc_dvb); e += 2) { + auto crd = tCcC(e); + int i = get<0>(crd), vv = get<1>(crd); + float dvb0 = acc_dvb(e), dvb1 = acc_dvb(e + 1); + db_acc[(e >> 1) & 1] += dvb0 * to_f32(sm.vl.v[i * kCP + vv]) + + dvb1 * to_f32(sm.vl.v[i * kCP + vv + 1]); + if (i < rows_valid) { + // adjacent-column pair -> one 4B store + T pair[2] = {T(dvb0 * sm.beta[i]), T(dvb1 * sm.beta[i])}; + *reinterpret_cast(dv2_p + (int64_t)i * hvV_row + v0 + vv) = + *reinterpret_cast(pair); + } + } + CUTE_UNROLL + for (int hh = 0; hh < 2; ++hh) + atomicAdd(&sm.db[get<0>(tCcC(2 * hh))], db_acc[hh]); + } + __syncthreads(); // tiles are reusable from the next i_v iteration + } + + // decay + dw/kg tiles for the dA/dkgb GEMMs (g/k already staged) + if (tid < 64) sm.dgk[tid] *= exp2f(sm.gn[tid]); + { + for (int e = 0; e < size(acc_dq); ++e) { + auto crd = tCcC(e); + int i = get<0>(crd), kk = get<1>(crd); + bool rv = i < rows_valid; + float g_ik = sm.sp.gk.g[i * 68 + kk]; + float k_ik = to_f32(sm.sp.gk.k[i * kCP + kk]); + float e2g = exp2f(g_ik); + acc_dq(e) *= e2g * scale; + acc_dk(e) *= rv ? exp2f(sm.gn[kk] - g_ik) : 0.f; + float dw_v = -acc_dw(e); + sm.ep.dw[i * kCP + kk] = T(dw_v); + sm.ep.dwT[kk * kCP + i] = T(dw_v); + sm.ep.kg[i * kCP + kk] = T(k_ik * e2g); + } + } + __syncthreads(); + gemm_sm(acc_dA, sm.ep.dw, sm.ep.kg); // dA += (-dw) @ kg^T + Tensor acc_dkgb = partition_fragment_C(tiled_mma, Shape, _64>{}); + clear(acc_dkgb); + gemm_sm(acc_dkgb, sm.A_t, sm.ep.dwT); // dkgb = Akk^T @ (-dw) + __syncthreads(); // dw/dwT/kg consumed; dw slot is dead + + // stage q into the dead dw slot, overlapped with the db/dgk phase + { + T const* qsrc = q_g + qk_base + (int64_t)t0 * qk_row + i_k * kBK; + stage_tile(sm.ep.dw, qsrc, q_g, qk_row, rows_valid, tid); + cp_async_commit(); + } + + // db += sum_k dkgb * kg; dgk[k] += sum_i k * dk (dk before the dkgb term) + // in-thread partial sums: e bit1 selects the row half (2 rows), e bit0 + // and bits 2+ select the column (8 cols) -> one atomic per row/col. + { + float db_acc[2] = {0.f, 0.f}; + float dgk_acc[8] = {}; + CUTE_UNROLL + for (int e = 0; e < size(acc_dkgb); ++e) { + auto crd = tCcC(e); + int i = get<0>(crd), kk = get<1>(crd); + float g_ik = sm.sp.gk.g[i * 68 + kk]; + float k_ik = to_f32(sm.sp.gk.k[i * kCP + kk]); + db_acc[(e >> 1) & 1] += acc_dkgb(e) * (k_ik * exp2f(g_ik)); + dgk_acc[(e & 1) + 2 * (e >> 2)] += k_ik * acc_dk(e); + } + CUTE_UNROLL + for (int hh = 0; hh < 2; ++hh) + atomicAdd(&sm.db[get<0>(tCcC(2 * hh))], db_acc[hh]); + CUTE_UNROLL + for (int cc = 0; cc < 8; ++cc) + atomicAdd(&sm.dgk[get<1>(tCcC((cc & 1) + 4 * (cc >> 1)))], dgk_acc[cc]); + } + cp_async_wait<0>(); + __syncthreads(); + + // dg/dk composition, store dq/dk/dg (adjacent-column pairs as float2) + { + float* dq_p = dq_g + hvK_base + (int64_t)t0 * hvK_row + i_k * kBK; + float* dk_p = dk_g + hvK_base + (int64_t)t0 * hvK_row + i_k * kBK; + float* dg_p = dg_g + hvK_base + (int64_t)t0 * hvK_row + i_k * kBK; + CUTE_UNROLL + for (int e = 0; e < size(acc_dq); e += 2) { + auto crd = tCcC(e); + int i = get<0>(crd), kk = get<1>(crd); + if (i >= rows_valid) continue; + float2 fq, fk, fg; + CUTE_UNROLL + for (int j = 0; j < 2; ++j) { + float g_ik = sm.sp.gk.g[i * 68 + kk + j]; + float e2g = exp2f(g_ik); + float k_ik = to_f32(sm.sp.gk.k[i * kCP + kk + j]); + float kg = k_ik * e2g; + float kdk = k_ik * acc_dk(e + j); + float dg_v = to_f32(sm.ep.dw[i * kCP + kk + j]) * acc_dq(e + j) - kdk + + (t0 + i == last ? sm.dgk[kk + j] : 0.f) + + kg * acc_dkgb(e + j) * sm.beta[i]; + float dk_v = acc_dk(e + j) + acc_dkgb(e + j) * e2g * sm.beta[i]; + (&fq.x)[j] = acc_dq(e + j); + (&fk.x)[j] = dk_v; + (&fg.x)[j] = dg_v; + } + int64_t off = (int64_t)i * hvK_row + kk; + *reinterpret_cast(dq_p + off) = fq; + *reinterpret_cast(dk_p + off) = fk; + *reinterpret_cast(dg_p + off) = fg; + } + } + } + + // dA postprocess: strict lower mask, column beta, then dAkk = -A^T (dA . beta) A^T + __syncthreads(); + for (int e = 0; e < size(acc_dA); ++e) { + auto crd = tCcC(e); + int i = get<0>(crd), j = get<1>(crd); + bool m = (i > j) && (i < rows_valid) && (j < rows_valid); + sm.pp.dAm[i * kCP + j] = T(m ? acc_dA(e) * sm.beta[j] : 0.f); + } + __syncthreads(); + Tensor acc1 = partition_fragment_C(tiled_mma, Shape, _64>{}); + clear(acc1); + gemm_sm_bt(acc1, sm.pp.dAm, sm.A_t); // dAm @ Akk^T (A_t[c][j] = Akk[j][c]) + __syncthreads(); + for (int e = 0; e < size(acc1); ++e) { + auto crd = tCcC(e); + int i = get<0>(crd), j = get<1>(crd); + sm.pp.c1T[j * kCP + i] = T(acc1(e)); + } + __syncthreads(); + Tensor acc2 = partition_fragment_C(tiled_mma, Shape, _64>{}); + clear(acc2); + gemm_sm(acc2, sm.A_t, sm.pp.c1T); // Akk^T @ (...) + { + float* dA_p = dA_g + A_base + (int64_t)t0 * A_row; + for (int e = 0; e < size(acc2); ++e) { + auto crd = tCcC(e); + int i = get<0>(crd), j = get<1>(crd); + if (i < rows_valid) { + bool m = (i > j) && (j < rows_valid); + dA_p[(int64_t)i * A_row + j] = m ? -acc2(e) : 0.f; + } + } + } + if (tid < kBT && tid < rows_valid) + db_g[beta_base + (int64_t)(t0 + tid) * HV] = sm.db[tid]; +} + +template +void launch_fused( + T const* q, T const* k, T const* v, T const* v_new, + float const* g, float const* beta, T const* A, T const* h, + T const* do_, T const* dh, T const* dv, + float* dq, float* dk, T* dv2, float* dg, float* db, float* dA, + int64_t const* cu_seqlens, int64_t const* chunk_indices, + float scale, int64_t B, int64_t T_len, int64_t H, int64_t HV, int64_t NT, + bool state_v_first, cudaStream_t stream +) { + auto* kern = chunk_kda_bwd_wy_dqkg_fused_kernel; + static bool configured = false; + if (!configured) { + cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, + int(sizeof(FusedSmem))); + configured = true; + } + dim3 grid(unsigned(NT), unsigned(B * HV)); + kern<<), stream>>>( + q, k, v, v_new, g, beta, A, h, do_, dh, dv, + dq, dk, dv2, dg, db, dA, + cu_seqlens, chunk_indices, scale, int(T_len), int(H), int(HV), + state_v_first, cu_seqlens != nullptr + ); +} + +template +void dispatch_kv( + torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& v_new, torch::Tensor const& g, torch::Tensor const& beta, + torch::Tensor const& A, torch::Tensor const& h, torch::Tensor const& do_, + torch::Tensor const& dh, torch::Tensor const& dv, + torch::Tensor& dq, torch::Tensor& dk, torch::Tensor& dv2, + torch::Tensor& dg, torch::Tensor& db, torch::Tensor& dA, + int64_t const* cu_seqlens, int64_t const* chunk_indices, + float scale, int64_t B, int64_t T_len, int64_t H, int64_t HV, int64_t NT, + bool state_v_first, cudaStream_t stream +) { + int64_t K = k.size(3), V = v.size(3); + #define LAUNCH_KV(KK, VV) \ + launch_fused( \ + reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), \ + reinterpret_cast(v.data_ptr()), reinterpret_cast(v_new.data_ptr()), \ + g.data_ptr(), beta.data_ptr(), \ + reinterpret_cast(A.data_ptr()), reinterpret_cast(h.data_ptr()), \ + reinterpret_cast(do_.data_ptr()), reinterpret_cast(dh.data_ptr()), \ + reinterpret_cast(dv.data_ptr()), \ + dq.data_ptr(), dk.data_ptr(), reinterpret_cast(dv2.data_ptr()), \ + dg.data_ptr(), db.data_ptr(), dA.data_ptr(), \ + cu_seqlens, chunk_indices, scale, B, T_len, H, HV, NT, state_v_first, stream) + if (K == 128 && V == 128) { LAUNCH_KV(128, 128); } + else if (K == 64 && V == 64) { LAUNCH_KV(64, 64); } + else if (K == 128 && V == 64) { LAUNCH_KV(128, 64); } + else if (K == 64 && V == 128) { LAUNCH_KV(64, 128); } + else { TORCH_CHECK(false, "unsupported K/V: ", K, "/", V, " (must be 64 or 128)"); } + #undef LAUNCH_KV +} + +} // namespace + +// fla/ops/kda/chunk_bwd.py::chunk_kda_bwd_wy_dqkg_fused host wrapper. +// Returns (dq, dk, dv2, db, dg, dAkk) in the same order as fla. +std::vector chunk_kda_bwd_wy_dqkg_fused( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor v_new, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor A, + torch::Tensor h, + torch::Tensor do_, + torch::Tensor dh, + torch::Tensor dv, + double scale, + bool state_v_first, + std::optional cu_seqlens, + std::optional chunk_indices, + int64_t chunk_size +) { + TORCH_CHECK(k.is_cuda() && k.is_contiguous(), "k must be contiguous CUDA tensor"); + TORCH_CHECK(k.dim() == 4, "k must be [B, T, H, K]"); + TORCH_CHECK(chunk_size == kBT, "only chunk_size=64 is supported"); + int64_t B = k.size(0), T_len = k.size(1), H = k.size(2), K = k.size(3); + int64_t HV = v.size(2), V = v.size(3); + TORCH_CHECK(K % kBK == 0 && V % kBV == 0, "K and V must be multiples of 64"); + TORCH_CHECK(HV % H == 0, "HV must be a multiple of H"); + TORCH_CHECK(q.scalar_type() == k.scalar_type() && v.scalar_type() == k.scalar_type() && + v_new.scalar_type() == k.scalar_type() && do_.scalar_type() == k.scalar_type() && + dh.scalar_type() == k.scalar_type() && dv.scalar_type() == k.scalar_type() && + A.scalar_type() == k.scalar_type() && h.scalar_type() == k.scalar_type(), + "q/k/v/v_new/do/dh/dv/A/h must share the same dtype"); + TORCH_CHECK(k.scalar_type() == at::kBFloat16 || k.scalar_type() == at::kHalf, + "only bf16/fp16 are supported"); + TORCH_CHECK(g.is_cuda() && g.is_contiguous() && g.scalar_type() == at::kFloat, "g must be fp32"); + TORCH_CHECK(beta.is_cuda() && beta.is_contiguous() && beta.scalar_type() == at::kFloat, "beta must be fp32"); + for (auto const* t : {&q, &v, &v_new, &A, &h, &do_, &dh, &dv}) { + TORCH_CHECK(t->is_cuda() && t->is_contiguous(), "all inputs must be contiguous CUDA tensors"); + } + + bool is_varlen = cu_seqlens.has_value(); + int64_t NT; + int64_t const* cu_ptr = nullptr; + int64_t const* ci_ptr = nullptr; + if (is_varlen) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_indices.has_value(), "chunk_indices must be provided with cu_seqlens"); + auto const& cu = cu_seqlens.value(); + auto const& ci = chunk_indices.value(); + TORCH_CHECK(cu.scalar_type() == torch::kLong && cu.is_cuda() && cu.is_contiguous()); + TORCH_CHECK(ci.scalar_type() == torch::kLong && ci.is_cuda() && ci.is_contiguous()); + TORCH_CHECK(ci.dim() == 2 && ci.size(1) == 2); + cu_ptr = cu.data_ptr(); + ci_ptr = ci.data_ptr(); + NT = ci.size(0); + } else { + NT = (T_len + kBT - 1) / kBT; + } + + auto opts_f = g.options().dtype(at::kFloat); + torch::Tensor dq = torch::empty({B, T_len, HV, K}, opts_f); + torch::Tensor dk = torch::empty({B, T_len, HV, K}, opts_f); + torch::Tensor dv2 = torch::empty_like(v); + torch::Tensor dg = torch::empty_like(g, opts_f); + torch::Tensor db = torch::empty_like(beta, opts_f); + torch::Tensor dA = torch::empty_like(A, opts_f); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + if (k.scalar_type() == at::kBFloat16) { + dispatch_kv( + q, k, v, v_new, g, beta, A, h, do_, dh, dv, dq, dk, dv2, dg, db, dA, + cu_ptr, ci_ptr, float(scale), B, T_len, H, HV, NT, state_v_first, stream); + } else { + dispatch_kv( + q, k, v, v_new, g, beta, A, h, do_, dh, dv, dq, dk, dv2, dg, db, dA, + cu_ptr, ci_ptr, float(scale), B, T_len, H, HV, NT, state_v_first, stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dq, dk, dv2, db, dg, dA}; +} diff --git a/csrc/train/bwd_wy_dqkg_binding.cpp b/csrc/train/bwd_wy_dqkg_binding.cpp new file mode 100644 index 0000000..0220cc3 --- /dev/null +++ b/csrc/train/bwd_wy_dqkg_binding.cpp @@ -0,0 +1,35 @@ +#include + +#include +#include + +// Defined in bwd_wy_dqkg.cu. fla/ops/kda/chunk_bwd.py::chunk_kda_bwd_wy_dqkg_fused. +// Returns (dq, dk, dv2, db, dg, dAkk). +std::vector chunk_kda_bwd_wy_dqkg_fused( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor v_new, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor A, + torch::Tensor h, + torch::Tensor do_, + torch::Tensor dh, + torch::Tensor dv, + double scale, + bool state_v_first, + std::optional cu_seqlens, + std::optional chunk_indices, + int64_t chunk_size +); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("chunk_kda_bwd_wy_dqkg_fused", &chunk_kda_bwd_wy_dqkg_fused, + "KDA fused backward dq/dk/dv2/dg/db/dAkk (CUDA)", + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("v_new"), + py::arg("g"), py::arg("beta"), py::arg("A"), py::arg("h"), + py::arg("do"), py::arg("dh"), py::arg("dv"), + py::arg("scale"), py::arg("state_v_first"), + py::arg("cu_seqlens"), py::arg("chunk_indices"), py::arg("chunk_size")); +} diff --git a/csrc/train/chunk_h.cu b/csrc/train/chunk_h.cu new file mode 100644 index 0000000..b07d61a --- /dev/null +++ b/csrc/train/chunk_h.cu @@ -0,0 +1,1288 @@ +// KDA chunked state-recurrence kernels: forward h and backward dhu. +// Replicates fla/ops/common/chunk_delta_h.py restricted to the KDA call shape: +// USE_G=False, USE_GK=True, SAVE_NEW_VALUE=True, chunk_size=64, bf16 operands. +// +// Each thread block owns one (sequence, value head, BV-wide V tile) and walks the +// sequence's chunks serially, mirroring the Triton program. For K <= 128 the +// pipelined kernels (kda_*_pipe_kernel) keep the [K, BV] fp32 state in registers +// as the accumulator fragment of the state-update MMA across the whole chunk +// loop; per chunk the bf16 state copy is staged to smem once (paired 32-bit +// stores, padded rows against bank conflicts) and the operand fragments are +// loaded with ldmatrix. All gmem tiles move through 16B cp.async into padded +// smem with one-chunk-ahead prefetch (phase-shifted single buffering), and all +// gmem outputs are staged through smem for fully coalesced 16B stores. +// For K > 128 the pipelined backward no longer fits the 99KB opt-in dynamic +// smem limit of consumer Blackwell (sm_120), so the legacy smem-state kernels +// below still serve those shapes. + +#include +#include + +#include + +#include + +#include "common.cuh" + +namespace { + +using cute::Int; +using cute::Layout; +using cute::Shape; +using cute::Stride; +using cute::Tensor; +using cute::_1; +using cute::_4; +using cute::get; +using cute::make_coord; +using cute::make_identity_tensor; +using cute::make_layout; +using cute::make_shape; +using cute::make_smem_ptr; +using cute::make_stride; +using cute::make_tensor; +using cute::make_tiled_copy_A; +using cute::make_tiled_copy_B; +using cute::make_tiled_mma; + +using BF16 = cutlass::bfloat16_t; + +constexpr int kBT = 64; // chunk size +constexpr int kThreads = 128; +constexpr int kPad = 8; // smem row padding in elements (16B) against bank conflicts + +using MmaAtom = cute::SM80_16x8x16_F32BF16BF16F32_TN; +using LdsmN = cute::Copy_Atom; +using LdsmT = cute::Copy_Atom; +// x2 variant for B operands: with Layout> the per-warp B tile is +// 8(N) x 16(K) = 4 vals/thread, too small for the x4 atom. +using LdsmT2 = cute::Copy_Atom; + +// --------------------------------------------------------------------------- +// cp.async helpers + +__device__ __forceinline__ void cp_async16(BF16* dst, BF16 const* src, bool full) { + uint32_t s = cute::cast_smem_ptr_to_uint(dst); + int sz = full ? 16 : 0; // src-size 0 zero-fills, used for masked rows/columns + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" + :: "r"(s), "l"(src), "r"(sz) : "memory"); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory"); +} + +__device__ __forceinline__ uint32_t pack_bf16(float lo, float hi) { + BF16 a(lo), b(hi); + return uint32_t(a.storage) | (uint32_t(b.storage) << 16); +} + +// --------------------------------------------------------------------------- +// pipelined-kernel shared memory + +template +struct FwdPipeSmem { + static constexpr int CP = Kp + kPad; + static constexpr int VP = BV + kPad; + alignas(128) BF16 s_w[kBT * CP]; // [BT][Kp] natural + alignas(128) BF16 s_kg[kBT * CP]; // [BT][Kp] natural + alignas(128) BF16 s_u[kBT * VP]; // [BT][BV] natural + alignas(128) BF16 s_hb[Kp * VP]; // [Kp][BV] bf16 state copy + alignas(128) BF16 s_vn[kBT * VP]; // [BT][BV] v_new +}; + +template +struct BwdPipeSmem { + static constexpr int CP = Kp + kPad; + static constexpr int VP = BV + kPad; + alignas(128) BF16 s_kg[kBT * CP]; // [BT][Kp] natural + alignas(128) BF16 s_qg[kBT * CP]; // [BT][Kp] natural + alignas(128) BF16 s_w[kBT * CP]; // [BT][Kp] natural + alignas(128) BF16 s_do[kBT * VP]; // [BT][BV] natural + alignas(128) BF16 s_dv[kBT * VP]; // [BT][BV] natural + alignas(128) BF16 s_dhb[Kp * VP]; // [Kp][BV] bf16 state-gradient copy + alignas(128) BF16 s_dv2[kBT * VP]; // [BT][BV] holds NEGATED dv2 (see kernel) +}; + +// --------------------------------------------------------------------------- +// pipelined-kernel staging helpers (all threads of the block cooperate) + +// Queue a [R][C] gmem tile (row stride `row_stride`) into padded smem [R][CP] +// with cp.async, zero-filling rows >= r_valid and columns >= c_valid via +// src-size 0. Falls back to scalar stores when 16B alignment does not hold. +template +__device__ void stage_tile_async(BF16* dst, BF16 const* src, BF16 const* src_safe, + int64_t row_stride, int r_valid, int c_valid, + bool vec_ok, int tid) { + constexpr int kCG = C / 8; + if (vec_ok && (c_valid % 8) == 0) { + for (int idx = tid; idx < R * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + bool full = (r < r_valid) && (c < c_valid); + BF16 const* g = full ? src + (int64_t)r * row_stride + c : src_safe; + cp_async16(dst + r * CP + c, g, full); + } + } else { + for (int idx = tid; idx < R * C; idx += kThreads) { + int r = idx / C, c = idx % C; + dst[r * CP + c] = (r < r_valid && c < c_valid) + ? src[(int64_t)r * row_stride + c] : BF16(0.f); + } + } +} + +// Store a padded smem tile [R][CP] to gmem (row stride `row_stride`), masked to +// r < r_valid and c < c_valid, with 16B vectorized accesses when aligned. +// NEG flips the bf16 sign bit (exact negation) on the way out. +template +__device__ void store_tile_gmem(OutT* dst, BF16 const* ssm, int64_t row_stride, + int r_valid, int c_valid, bool vec_ok, int tid) { + constexpr int kCG = C / 8; + if (vec_ok && (c_valid % 8) == 0) { + for (int idx = tid; idx < R * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + if (r >= r_valid) break; + if (c + 8 <= c_valid) { + uint4 v = *reinterpret_cast(ssm + r * CP + c); + if (NEG) { + v.x ^= 0x80008000u; v.y ^= 0x80008000u; + v.z ^= 0x80008000u; v.w ^= 0x80008000u; + } + *reinterpret_cast(dst + (int64_t)r * row_stride + c) = v; + } else { + for (int j = 0; j < 8 && c + j < c_valid; ++j) { + float val = to_f32(ssm[r * CP + c + j]); + dst[(int64_t)r * row_stride + c + j] = OutT(NEG ? -val : val); + } + } + } + } else { + for (int idx = tid; idx < R * C; idx += kThreads) { + int r = idx / C, c = idx % C; + if (r < r_valid && c < c_valid) { + float val = to_f32(ssm[r * CP + c]); + dst[(int64_t)r * row_stride + c] = OutT(NEG ? -val : val); + } + } + } +} + +// --------------------------------------------------------------------------- +// kernel params + +struct FwdParams { + BF16 const* kg; + BF16 const* w; + BF16 const* u; + float const* gk; + BF16* h; + BF16* v_new; + float const* h0; + float* ht; + int64_t const* cu_seqlens; + int64_t const* chunk_offsets; + int T, H, HV, K, V, NV; + bool varlen, state_v_first; +}; + +struct BwdParams { + BF16 const* qg; + BF16 const* kg; + BF16 const* w; + float const* gk; + BF16 const* do_; + BF16 const* dv; + BF16* dv2; + BF16* dh; + float* dh0; + float const* dht; + int64_t const* cu_seqlens; + int64_t const* chunk_offsets; + float scale; + int T, H, HV, K, V, NV; + bool varlen, state_v_first; +}; + +// --------------------------------------------------------------------------- +// pipelined forward: per chunk, h[i_t] = S; v_new = u - w @ S; +// S = diag(exp2(gk_last)) S + kg^T @ v_new +// +// S lives in the accumulator fragment of the update MMA across the chunk loop. +// Per chunk: prefetch (w, u, kg) of chunk i+1 with cp.async; the update GEMM +// accumulates directly into the state fragment after the elementwise decay. + +template +__global__ void __launch_bounds__(kThreads) kda_fwd_h_pipe_kernel(FwdParams const p) { + constexpr int CP = Kp + kPad; + constexpr int VP = BV + kPad; + constexpr int kMM = Kp / 64; // MMA_M of the state-update GEMM + extern __shared__ __align__(128) unsigned char smem_raw[]; + auto& sm = *reinterpret_cast*>(smem_raw); + const int tid = threadIdx.x; + + const int i_v = blockIdx.x % p.NV; + const int64_t i_nh = blockIdx.x / p.NV; + const int64_t i_n = i_nh / p.HV; + const int i_hv = int(i_nh % p.HV); + const int i_hq = i_hv / (p.HV / p.H); + + int64_t bos; + int T, NT; + int64_t boh; + if (p.varlen) { + bos = p.cu_seqlens[i_n]; + T = int(p.cu_seqlens[i_n + 1] - bos); + NT = (T + kBT - 1) / kBT; + boh = p.chunk_offsets[i_n]; + } else { + bos = i_n * p.T; + T = p.T; + NT = (T + kBT - 1) / kBT; + boh = i_n * NT; + } + + BF16 const* kg = p.kg + (bos * p.H + i_hq) * (int64_t)p.K; + BF16 const* w = p.w + (bos * p.HV + i_hv) * (int64_t)p.K; + BF16 const* u = p.u + (bos * p.HV + i_hv) * (int64_t)p.V; + float const* gk = p.gk + (bos * p.HV + i_hv) * (int64_t)p.K; + BF16* h = p.h + (boh * p.HV + i_hv) * (int64_t)p.K * p.V; + BF16* v_new = p.v_new + (bos * p.HV + i_hv) * (int64_t)p.V; + const int v0 = i_v * BV; + + const bool vec_k = (p.K % 8) == 0; + const bool vec_v = (p.V % 8) == 0; + const int cv = min(BV, p.V - v0); + + auto mma = make_tiled_mma(MmaAtom{}, Layout>{}); + auto thr_mma = mma.get_thread_slice(tid); + + Tensor sW = make_tensor(make_smem_ptr(sm.s_w), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + Tensor sU = make_tensor(make_smem_ptr(sm.s_u), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + Tensor sVN = make_tensor(make_smem_ptr(sm.s_vn), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + // GEMM operand views: the state copy as B[N=BV, K=Kp], kg as A[M=Kp, K=BT], + // v_new as B[N=BV, K=BT], all with the contraction dim strided (ldsm .trans). + Tensor sHBb = make_tensor(make_smem_ptr(sm.s_hb), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + Tensor sKGa = make_tensor(make_smem_ptr(sm.s_kg), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + Tensor sVNb = make_tensor(make_smem_ptr(sm.s_vn), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + + Tensor state = cute::partition_fragment_C(mma, Shape, Int>{}); + Tensor tCcS = thr_mma.partition_C(make_identity_tensor(make_shape(Int{}, Int{}))); + Tensor tCcV = thr_mma.partition_C(make_identity_tensor(make_shape(Int{}, Int{}))); + + // initial state, fragment-direct + { + float const* h0 = p.h0 == nullptr ? nullptr : p.h0 + i_nh * (int64_t)p.K * p.V; + for (int i = 0; i < int(size(state)); ++i) { + int k = get<0>(tCcS(i)), n = get<1>(tCcS(i)); + int vg = v0 + n; + float val = 0.f; + if (h0 != nullptr && k < p.K && vg < p.V) { + val = h0[p.state_v_first ? (int64_t)vg * p.K + k : (int64_t)k * p.V + vg]; + } + state(i) = val; + } + } + // distinct state rows this thread holds: index i -> row j = (i>>1 & 1) + 2*((i>>2) % kMM) + int row_k[2 * kMM]; + for (int m = 0; m < kMM; ++m) { + for (int half = 0; half < 2; ++half) { + row_k[m * 2 + half] = get<0>(tCcS(make_coord(make_coord(0, half), m, 0))); + } + } + + auto stage_wu = [&](int i_t) { + int tv = i_t < NT ? min(kBT, T - i_t * kBT) : 0; + int64_t t0 = (int64_t)i_t * kBT; + BF16 const* wsrc = i_t < NT ? w + t0 * (int64_t)p.HV * p.K : w; + BF16 const* usrc = i_t < NT ? u + t0 * (int64_t)p.HV * p.V + v0 : u; + stage_tile_async(sm.s_w, wsrc, w, (int64_t)p.HV * p.K, tv, p.K, vec_k, tid); + stage_tile_async(sm.s_u, usrc, u, (int64_t)p.HV * p.V, tv, cv, vec_v, tid); + cp_async_commit(); + }; + auto stage_kg = [&](int i_t) { + int tv = i_t < NT ? min(kBT, T - i_t * kBT) : 0; + int64_t t0 = (int64_t)i_t * kBT; + BF16 const* ksrc = i_t < NT ? kg + t0 * (int64_t)p.H * p.K : kg; + stage_tile_async(sm.s_kg, ksrc, kg, (int64_t)p.H * p.K, tv, p.K, vec_k, tid); + cp_async_commit(); + }; + + stage_wu(0); + stage_kg(0); + + for (int i_t = 0; i_t < NT; ++i_t) { + const int t_valid = min(kBT, T - i_t * kBT); + const int64_t t0 = (int64_t)i_t * kBT; + + cp_async_wait<1>(); // w, u of this chunk (kg may still be in flight) + __syncthreads(); + + // decay factors for the state update at the end of this chunk + const int last = min((i_t + 1) * kBT, T) - 1; + float const* gk_last = gk + (int64_t)last * p.HV * p.K; + float gval[2 * kMM]; + for (int j = 0; j < 2 * kMM; ++j) { + gval[j] = row_k[j] < p.K ? exp2f(gk_last[row_k[j]]) : 1.0f; + } + + // bf16 state copy: B operand of the v_new GEMM and the h[i_t] output + for (int i = 0; i < int(size(state)); i += 2) { + int k = get<0>(tCcS(i)), n = get<1>(tCcS(i)); + *reinterpret_cast(&sm.s_hb[k * VP + n]) = pack_bf16(state(i), state(i + 1)); + } + __syncthreads(); + + // v_new = w @ h operand fragments + Tensor tCrW = thr_mma.partition_fragment_A(sW); + Tensor tCrHB = thr_mma.partition_fragment_B(sHBb); + auto s2r_w = make_tiled_copy_A(LdsmN{}, mma); + auto thr_w = s2r_w.get_thread_slice(tid); + copy(s2r_w, thr_w.partition_S(sW), thr_w.retile_D(tCrW)); + auto s2r_hb = make_tiled_copy_B(LdsmT2{}, mma); + auto thr_hb = s2r_hb.get_thread_slice(tid); + copy(s2r_hb, thr_hb.partition_S(sHBb), thr_hb.retile_D(tCrHB)); + + // h[i_t] = state at chunk entry (bf16) + BF16* h_t = h + (int64_t)i_t * p.HV * p.K * p.V; + if (p.state_v_first) { + for (int idx = tid; idx < Kp * BV; idx += kThreads) { + int k = idx / BV, n = idx % BV, vg = v0 + n; + if (k < p.K && vg < p.V) h_t[(int64_t)vg * p.K + k] = sm.s_hb[k * VP + n]; + } + } else { + store_tile_gmem(h_t + v0, sm.s_hb, p.V, p.K, cv, vec_v, tid); + } + + Tensor acc = cute::partition_fragment_C(mma, Shape, Int>{}); + clear(acc); + gemm(thr_mma, tCrW, tCrHB, acc); + + // v_new = u - w @ h; u is zero-filled beyond t_valid, so rows past the + // sequence end yield exactly 0 (required by the state-update GEMM). + Tensor tCsU = thr_mma.partition_C(sU); + for (int i = 0; i < int(size(acc)); i += 2) { + int m = get<0>(tCcV(i)), n = get<1>(tCcV(i)); + float vlo = to_f32(tCsU(i)) - acc(i); + float vhi = to_f32(tCsU(i + 1)) - acc(i + 1); + *reinterpret_cast(&sm.s_vn[m * VP + n]) = pack_bf16(vlo, vhi); + } + __syncthreads(); // s_w consumed by the GEMM; s_vn complete + + stage_wu(i_t + 1); + store_tile_gmem(v_new + t0 * (int64_t)p.HV * p.V + v0, + sm.s_vn, (int64_t)p.HV * p.V, t_valid, cv, vec_v, tid); + + cp_async_wait<1>(); // kg of this chunk (the new w/u group may fly) + __syncthreads(); + + // h = diag(exp2(gk_last)) h + kg^T @ v_new, accumulated into the state + Tensor tCrKG = thr_mma.partition_fragment_A(sKGa); + Tensor tCrVN = thr_mma.partition_fragment_B(sVNb); + auto s2r_kg = make_tiled_copy_A(LdsmT{}, mma); + auto thr_kg = s2r_kg.get_thread_slice(tid); + copy(s2r_kg, thr_kg.partition_S(sKGa), thr_kg.retile_D(tCrKG)); + auto s2r_vn = make_tiled_copy_B(LdsmT2{}, mma); + auto thr_vn = s2r_vn.get_thread_slice(tid); + copy(s2r_vn, thr_vn.partition_S(sVNb), thr_vn.retile_D(tCrVN)); + + for (int i = 0; i < int(size(state)); ++i) { + state(i) *= gval[((i >> 1) & 1) + 2 * ((i >> 2) % kMM)]; + } + gemm(thr_mma, tCrKG, tCrVN, state); + __syncthreads(); // s_kg and s_vn consumed + stage_kg(i_t + 1); + } + + if (p.ht != nullptr) { + float* ht = p.ht + i_nh * (int64_t)p.K * p.V; + for (int i = 0; i < int(size(state)); ++i) { + int k = get<0>(tCcS(i)), n = get<1>(tCcS(i)), vg = v0 + n; + if (k < p.K && vg < p.V) { + ht[p.state_v_first ? (int64_t)vg * p.K + k : (int64_t)k * p.V + vg] = state(i); + } + } + } +} + +// --------------------------------------------------------------------------- +// pipelined backward: per chunk (in reverse), dh[i_t] = S; +// dv2 = dv + kg @ S; S = S * exp2(gk_last) + scale * qg^T @ do - w^T @ dv2 +// +// S_dv2 holds NEGATED dv2 (sign flip is exact in bf16) so the w^T @ dv2 term +// accumulates directly into the state fragment; the gmem dv2 store flips the +// sign back. The qg^T @ do term uses a temporary fragment for the scale. + +template +__global__ void __launch_bounds__(kThreads) kda_bwd_dhu_pipe_kernel(BwdParams const p) { + constexpr int CP = Kp + kPad; + constexpr int VP = BV + kPad; + constexpr int kMM = Kp / 64; + extern __shared__ __align__(128) unsigned char smem_raw[]; + auto& sm = *reinterpret_cast*>(smem_raw); + const int tid = threadIdx.x; + + const int i_v = blockIdx.x % p.NV; + const int64_t i_nh = blockIdx.x / p.NV; + const int64_t i_n = i_nh / p.HV; + const int i_hv = int(i_nh % p.HV); + const int i_hq = i_hv / (p.HV / p.H); + + int64_t bos; + int T, NT; + int64_t boh; + if (p.varlen) { + bos = p.cu_seqlens[i_n]; + T = int(p.cu_seqlens[i_n + 1] - bos); + NT = (T + kBT - 1) / kBT; + boh = p.chunk_offsets[i_n]; + } else { + bos = i_n * p.T; + T = p.T; + NT = (T + kBT - 1) / kBT; + boh = i_n * NT; + } + + BF16 const* qg = p.qg + (bos * p.H + i_hq) * (int64_t)p.K; + BF16 const* kg = p.kg + (bos * p.H + i_hq) * (int64_t)p.K; + BF16 const* w = p.w + (bos * p.HV + i_hv) * (int64_t)p.K; + float const* gk = p.gk + (bos * p.HV + i_hv) * (int64_t)p.K; + BF16 const* do_ = p.do_ + (bos * p.HV + i_hv) * (int64_t)p.V; + BF16 const* dv = p.dv + (bos * p.HV + i_hv) * (int64_t)p.V; + BF16* dv2 = p.dv2 + (bos * p.HV + i_hv) * (int64_t)p.V; + BF16* dh = p.dh + (boh * p.HV + i_hv) * (int64_t)p.K * p.V; + const int v0 = i_v * BV; + + const bool vec_k = (p.K % 8) == 0; + const bool vec_v = (p.V % 8) == 0; + const int cv = min(BV, p.V - v0); + + auto mma = make_tiled_mma(MmaAtom{}, Layout>{}); + auto thr_mma = mma.get_thread_slice(tid); + + Tensor sKG = make_tensor(make_smem_ptr(sm.s_kg), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + Tensor sDV = make_tensor(make_smem_ptr(sm.s_dv), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + // GEMM operand views (all contraction-strided, ldsm .trans), except kg which + // is the K-major A operand of the dv2 GEMM. + Tensor sDHBb = make_tensor(make_smem_ptr(sm.s_dhb), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + Tensor sQGa = make_tensor(make_smem_ptr(sm.s_qg), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + Tensor sWa = make_tensor(make_smem_ptr(sm.s_w), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + Tensor sDOb = make_tensor(make_smem_ptr(sm.s_do), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + Tensor sDV2b = make_tensor(make_smem_ptr(sm.s_dv2), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, Int{}))); + + Tensor state = cute::partition_fragment_C(mma, Shape, Int>{}); + Tensor tCcS = thr_mma.partition_C(make_identity_tensor(make_shape(Int{}, Int{}))); + Tensor tCcV = thr_mma.partition_C(make_identity_tensor(make_shape(Int{}, Int{}))); + + { + float const* dht = p.dht == nullptr ? nullptr : p.dht + i_nh * (int64_t)p.K * p.V; + for (int i = 0; i < int(size(state)); ++i) { + int k = get<0>(tCcS(i)), n = get<1>(tCcS(i)); + int vg = v0 + n; + float val = 0.f; + if (dht != nullptr && k < p.K && vg < p.V) { + val = dht[p.state_v_first ? (int64_t)vg * p.K + k : (int64_t)k * p.V + vg]; + } + state(i) = val; + } + } + int row_k[2 * kMM]; + for (int m = 0; m < kMM; ++m) { + for (int half = 0; half < 2; ++half) { + row_k[m * 2 + half] = get<0>(tCcS(make_coord(make_coord(0, half), m, 0))); + } + } + + auto stage_kgdv = [&](int i_t) { + bool ok = i_t >= 0; + int tv = ok ? min(kBT, T - i_t * kBT) : 0; + int64_t t0 = (int64_t)i_t * kBT; + BF16 const* ksrc = ok ? kg + t0 * (int64_t)p.H * p.K : kg; + BF16 const* dvsrc = ok ? dv + t0 * (int64_t)p.HV * p.V + v0 : dv; + stage_tile_async(sm.s_kg, ksrc, kg, (int64_t)p.H * p.K, tv, p.K, vec_k, tid); + stage_tile_async(sm.s_dv, dvsrc, dv, (int64_t)p.HV * p.V, tv, cv, vec_v, tid); + cp_async_commit(); + }; + auto stage_qgdo = [&](int i_t) { + bool ok = i_t >= 0; + int tv = ok ? min(kBT, T - i_t * kBT) : 0; + int64_t t0 = (int64_t)i_t * kBT; + BF16 const* qsrc = ok ? qg + t0 * (int64_t)p.H * p.K : qg; + BF16 const* dosrc = ok ? do_ + t0 * (int64_t)p.HV * p.V + v0 : do_; + stage_tile_async(sm.s_qg, qsrc, qg, (int64_t)p.H * p.K, tv, p.K, vec_k, tid); + stage_tile_async(sm.s_do, dosrc, do_, (int64_t)p.HV * p.V, tv, cv, vec_v, tid); + cp_async_commit(); + }; + auto stage_w = [&](int i_t) { + bool ok = i_t >= 0; + int tv = ok ? min(kBT, T - i_t * kBT) : 0; + int64_t t0 = (int64_t)i_t * kBT; + BF16 const* wsrc = ok ? w + t0 * (int64_t)p.HV * p.K : w; + stage_tile_async(sm.s_w, wsrc, w, (int64_t)p.HV * p.K, tv, p.K, vec_k, tid); + cp_async_commit(); + }; + + stage_kgdv(NT - 1); + stage_qgdo(NT - 1); + stage_w(NT - 1); + + for (int i_t = NT - 1; i_t >= 0; --i_t) { + const int t_valid = min(kBT, T - i_t * kBT); + const int64_t t0 = (int64_t)i_t * kBT; + + cp_async_wait<2>(); // kg, dv of this chunk + __syncthreads(); + + const int last = min((i_t + 1) * kBT, T) - 1; + float const* gk_last = gk + (int64_t)last * p.HV * p.K; + float gval[2 * kMM]; + for (int j = 0; j < 2 * kMM; ++j) { + gval[j] = row_k[j] < p.K ? exp2f(gk_last[row_k[j]]) : 1.0f; + } + + // bf16 state-gradient copy: B operand of the dv2 GEMM and the dh[i_t] output + for (int i = 0; i < int(size(state)); i += 2) { + int k = get<0>(tCcS(i)), n = get<1>(tCcS(i)); + *reinterpret_cast(&sm.s_dhb[k * VP + n]) = pack_bf16(state(i), state(i + 1)); + } + __syncthreads(); + + // dv2 = kg @ dh operand fragments + Tensor tCrKG = thr_mma.partition_fragment_A(sKG); + Tensor tCrDHB = thr_mma.partition_fragment_B(sDHBb); + auto s2r_kg = make_tiled_copy_A(LdsmN{}, mma); + auto thr_kg = s2r_kg.get_thread_slice(tid); + copy(s2r_kg, thr_kg.partition_S(sKG), thr_kg.retile_D(tCrKG)); + auto s2r_dhb = make_tiled_copy_B(LdsmT2{}, mma); + auto thr_dhb = s2r_dhb.get_thread_slice(tid); + copy(s2r_dhb, thr_dhb.partition_S(sDHBb), thr_dhb.retile_D(tCrDHB)); + + // dh[i_t] = state gradient at chunk entry (bf16) + BF16* dh_t = dh + (int64_t)i_t * p.HV * p.K * p.V; + if (p.state_v_first) { + for (int idx = tid; idx < Kp * BV; idx += kThreads) { + int k = idx / BV, n = idx % BV, vg = v0 + n; + if (k < p.K && vg < p.V) dh_t[(int64_t)vg * p.K + k] = sm.s_dhb[k * VP + n]; + } + } else { + store_tile_gmem(dh_t + v0, sm.s_dhb, p.V, p.K, cv, vec_v, tid); + } + + Tensor acc = cute::partition_fragment_C(mma, Shape, Int>{}); + clear(acc); + gemm(thr_mma, tCrKG, tCrDHB, acc); + + // s_dv2 <- -(dv + kg @ dh); the sign is flipped back on the gmem store. + Tensor tCsDV = thr_mma.partition_C(sDV); + for (int i = 0; i < int(size(acc)); i += 2) { + int m = get<0>(tCcV(i)), n = get<1>(tCcV(i)); + float vlo = acc(i) + to_f32(tCsDV(i)); + float vhi = acc(i + 1) + to_f32(tCsDV(i + 1)); + *reinterpret_cast(&sm.s_dv2[m * VP + n]) = pack_bf16(-vlo, -vhi); + } + __syncthreads(); // s_kg / s_dv / s_dhb consumed; s_dv2 complete + + stage_kgdv(i_t - 1); + store_tile_gmem(dv2 + t0 * (int64_t)p.HV * p.V + v0, + sm.s_dv2, (int64_t)p.HV * p.V, t_valid, cv, vec_v, tid); + + cp_async_wait<1>(); // qg, do and w of this chunk + __syncthreads(); + + // dh = dh * exp2(gk_last) + scale * qg^T @ do + Tensor tmp = cute::partition_fragment_C(mma, Shape, Int>{}); + clear(tmp); + Tensor tCrQG = thr_mma.partition_fragment_A(sQGa); + Tensor tCrDO = thr_mma.partition_fragment_B(sDOb); + auto s2r_qg = make_tiled_copy_A(LdsmT{}, mma); + auto thr_qg = s2r_qg.get_thread_slice(tid); + copy(s2r_qg, thr_qg.partition_S(sQGa), thr_qg.retile_D(tCrQG)); + auto s2r_do = make_tiled_copy_B(LdsmT2{}, mma); + auto thr_do = s2r_do.get_thread_slice(tid); + copy(s2r_do, thr_do.partition_S(sDOb), thr_do.retile_D(tCrDO)); + gemm(thr_mma, tCrQG, tCrDO, tmp); + for (int i = 0; i < int(size(state)); ++i) { + int j = ((i >> 1) & 1) + 2 * ((i >> 2) % kMM); + state(i) = state(i) * gval[j] + p.scale * tmp(i); + } + __syncthreads(); // s_qg / s_do consumed + stage_qgdo(i_t - 1); + + // dh -= w^T @ dv2, via the negated s_dv2 accumulated into the state + Tensor tCrW = thr_mma.partition_fragment_A(sWa); + Tensor tCrDV2 = thr_mma.partition_fragment_B(sDV2b); + auto s2r_w = make_tiled_copy_A(LdsmT{}, mma); + auto thr_w = s2r_w.get_thread_slice(tid); + copy(s2r_w, thr_w.partition_S(sWa), thr_w.retile_D(tCrW)); + auto s2r_dv2 = make_tiled_copy_B(LdsmT2{}, mma); + auto thr_dv2 = s2r_dv2.get_thread_slice(tid); + copy(s2r_dv2, thr_dv2.partition_S(sDV2b), thr_dv2.retile_D(tCrDV2)); + gemm(thr_mma, tCrW, tCrDV2, state); + __syncthreads(); // s_w / s_dv2 consumed + stage_w(i_t - 1); + } + + if (p.dh0 != nullptr) { + float* dh0 = p.dh0 + i_nh * (int64_t)p.K * p.V; + for (int i = 0; i < int(size(state)); ++i) { + int k = get<0>(tCcS(i)), n = get<1>(tCcS(i)), vg = v0 + n; + if (k < p.K && vg < p.V) { + dh0[p.state_v_first ? (int64_t)vg * p.K + k : (int64_t)k * p.V + vg] = state(i); + } + } + } +} + +// --------------------------------------------------------------------------- +// legacy kernels (K > 128): fp32 state in shared memory, scalar fragment loads. + +template +struct FwdSmem { + // persistent across the chunk loop + alignas(128) float s_h[Kp * BV]; // [Kp][BV] fp32 state + alignas(128) BF16 s_vnT[BV * kBT]; // [BV][BT] v_new, transposed (B operand of the update GEMM) + alignas(128) float s_g[Kp]; // exp2(gk_last) per K channel + union { + struct { // v_new phase + alignas(128) BF16 s_w[kBT * Kp]; // [BT][Kp] natural + alignas(128) BF16 s_hb[BV * Kp]; // [BV][Kp] state in bf16, transposed + } a; + struct { // state-update phase + alignas(128) BF16 s_kgT[Kp * kBT]; // [Kp][BT] kg, transposed + } b; + }; +}; + +template +struct BwdSmem { + alignas(128) float s_dh[Kp * BV]; // [Kp][BV] fp32 state gradient + alignas(128) BF16 s_dv2T[BV * kBT]; // [BV][BT] dv2, transposed + alignas(128) float s_g[Kp]; + union { + struct { // dv2 phase + alignas(128) BF16 s_kg[kBT * Kp]; // [BT][Kp] natural + alignas(128) BF16 s_dhb[BV * Kp]; // [BV][Kp] dh in bf16, transposed + } a; + struct { // dh-update phase (s_qwT is staged with qg first, then w) + alignas(128) BF16 s_qwT[Kp * kBT]; // [Kp][BT] qg / w, transposed + alignas(128) BF16 s_doT[BV * kBT]; // [BV][BT] do, transposed + } b; + }; +}; + +// --------------------------------------------------------------------------- +// staging helpers (all threads of the block cooperate) + +// Load a [R][C] tile from gmem (row stride `row_stride`) into smem, zero-filling +// rows >= r_valid and columns >= c_valid. +template +__device__ void stage_tile(BF16* dst, BF16 const* src, int64_t row_stride, + int r_valid, int c_valid, int tid) { + for (int idx = tid; idx < R * C; idx += kThreads) { + int r = idx / C, c = idx % C; + dst[idx] = (r < r_valid && c < c_valid) + ? src[(int64_t)r * row_stride + c] : BF16(0.f); + } +} + +// Load a gmem [R][C] tile into smem transposed as [C][R], zero-filling like above. +// Reads are 16B-vectorized along C when alignment and validity allow. +template +__device__ void stage_tile_T(BF16* dst, BF16 const* src, int64_t row_stride, + int r_valid, int c_valid, int tid) { + constexpr int kVec = 8; + constexpr int kCG = C / kVec; + bool vec_ok = (row_stride % kVec) == 0; + for (int idx = tid; idx < R * kCG; idx += kThreads) { + int r = idx / kCG, c0 = (idx % kCG) * kVec; + BF16 vals[kVec]; + if (r < r_valid && vec_ok && c0 + kVec <= c_valid) { + *reinterpret_cast(vals) = + *reinterpret_cast(src + (int64_t)r * row_stride + c0); + } else { + for (int j = 0; j < kVec; ++j) { + int c = c0 + j; + vals[j] = (r < r_valid && c < c_valid) + ? src[(int64_t)r * row_stride + c] : BF16(0.f); + } + } + for (int j = 0; j < kVec; ++j) dst[(c0 + j) * R + r] = vals[j]; + } +} + +// bf16 transposed copy of the fp32 state: dst[n][k] = BF16(s_state[k][n]). +template +__device__ void build_state_T(BF16* dst, float const* s_state, int tid) { + for (int idx = tid; idx < Kp * BV; idx += kThreads) { + int k = idx / BV, n = idx % BV; + dst[n * Kp + k] = BF16(s_state[idx]); + } +} + +// s_g[k] = exp2(gk_last[k]) for k < K, 1 otherwise (matches the Triton masked load +// with other=0 followed by exp2). +template +__device__ void load_decay_g(float* s_g, float const* gk_last, int K, int tid) { + for (int k = tid; k < Kp; k += kThreads) { + s_g[k] = (k < K) ? exp2f(gk_last[k]) : 1.0f; + } +} + +template +__device__ void decay_state(float* s_state, float const* s_g, int tid) { + for (int idx = tid; idx < Kp * BV; idx += kThreads) { + s_state[idx] *= s_g[idx / BV]; + } +} + +// Store the [Kp][BV] state tile to gmem, masked to k < K and v < V. +// state_v_first: gmem [V, K] layout (v * K + k), else [K, V] (k * V + v). +template +__device__ void store_state_gmem(OutT* dst, float const* s_state, + int K, int V, int v0, bool state_v_first, int tid) { + for (int idx = tid; idx < Kp * BV; idx += kThreads) { + int k = idx / BV, n = idx % BV; + int vg = v0 + n; + if (k < K && vg < V) { + int64_t off = state_v_first ? (int64_t)vg * K + k : (int64_t)k * V + vg; + dst[off] = OutT(s_state[idx]); + } + } +} + +// Initialize the fp32 state tile from gmem (h0/dht), zero-filling out-of-range. +template +__device__ void load_state_gmem(float* s_state, float const* src, + int K, int V, int v0, bool state_v_first, int tid) { + for (int idx = tid; idx < Kp * BV; idx += kThreads) { + int k = idx / BV, n = idx % BV; + int vg = v0 + n; + float val = 0.f; + if (src != nullptr && k < K && vg < V) { + int64_t off = state_v_first ? (int64_t)vg * K + k : (int64_t)k * V + vg; + val = src[off]; + } + s_state[idx] = val; + } +} + +// --------------------------------------------------------------------------- +// CuTe GEMM: C[M, N] += A[M, Kc] * B[N, Kc] with both operands smem-resident, +// row-major with the contraction dim contiguous. Returns the per-thread fp32 +// accumulator fragment. Scalar smem loads (correctness first; ldmatrix is a +// performance follow-up). +template +__device__ auto mma_tn(TiledMMA const& tiled_mma, BF16 const* sA, BF16 const* sB, int tid) { + auto thr_mma = tiled_mma.get_thread_slice(tid); + Tensor sAt = make_tensor(make_smem_ptr(sA), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + Tensor sBt = make_tensor(make_smem_ptr(sB), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + Tensor acc = cute::partition_fragment_C(tiled_mma, Shape, Int>{}); + clear(acc); + Tensor tCrA = thr_mma.partition_fragment_A(sAt); + Tensor tCrB = thr_mma.partition_fragment_B(sBt); + Tensor tCsA = thr_mma.partition_A(sAt); + Tensor tCsB = thr_mma.partition_B(sBt); + for (int i = 0; i < int(size(tCrA)); ++i) tCrA(i) = tCsA(i); + for (int i = 0; i < int(size(tCrB)); ++i) tCrB(i) = tCsB(i); + gemm(thr_mma, tCrA, tCrB, acc); + return acc; +} + +// Per-thread (m, n) coordinates of an accumulator fragment. +template +__device__ auto mma_coords(TiledMMA const& tiled_mma, int tid) { + auto thr_mma = tiled_mma.get_thread_slice(tid); + Tensor cC = make_identity_tensor(make_shape(Int{}, Int{})); + return thr_mma.partition_C(cC); +} + +// --------------------------------------------------------------------------- +// forward: per chunk, h[i_t] = S; v_new = u - w @ S; S = diag(exp2(gk_last)) S + kg^T @ v_new + +template +__global__ void __launch_bounds__(kThreads) kda_fwd_h_kernel(FwdParams const p) { + extern __shared__ __align__(128) unsigned char smem_raw[]; + auto& sm = *reinterpret_cast*>(smem_raw); + const int tid = threadIdx.x; + + const int i_v = blockIdx.x % p.NV; + const int64_t i_nh = blockIdx.x / p.NV; + const int64_t i_n = i_nh / p.HV; + const int i_hv = int(i_nh % p.HV); + const int i_hq = i_hv / (p.HV / p.H); + + int64_t bos; + int T, NT; + int64_t boh; + if (p.varlen) { + bos = p.cu_seqlens[i_n]; + T = int(p.cu_seqlens[i_n + 1] - bos); + NT = (T + kBT - 1) / kBT; + boh = p.chunk_offsets[i_n]; + } else { + bos = i_n * p.T; + T = p.T; + NT = (T + kBT - 1) / kBT; + boh = i_n * NT; + } + + BF16 const* kg = p.kg + (bos * p.H + i_hq) * (int64_t)p.K; + BF16 const* w = p.w + (bos * p.HV + i_hv) * (int64_t)p.K; + BF16 const* u = p.u + (bos * p.HV + i_hv) * (int64_t)p.V; + float const* gk = p.gk + (bos * p.HV + i_hv) * (int64_t)p.K; + BF16* h = p.h + (boh * p.HV + i_hv) * (int64_t)p.K * p.V; + BF16* v_new = p.v_new + (bos * p.HV + i_hv) * (int64_t)p.V; + const int v0 = i_v * BV; + + load_state_gmem(sm.s_h, + p.h0 == nullptr ? nullptr : p.h0 + i_nh * (int64_t)p.K * p.V, + p.K, p.V, v0, p.state_v_first, tid); + + auto mma_vn = make_tiled_mma(MmaAtom{}, Layout>{}); + auto mma_up = make_tiled_mma(MmaAtom{}, Layout>{}); + + __syncthreads(); + for (int i_t = 0; i_t < NT; ++i_t) { + const int t_valid = min(kBT, T - i_t * kBT); + const int64_t t0 = (int64_t)i_t * kBT; + + // h[i_t] = state at chunk entry (bf16) + store_state_gmem(h + (int64_t)i_t * p.HV * p.K * p.V, + sm.s_h, p.K, p.V, v0, p.state_v_first, tid); + // stage w, the bf16 state copy, and the decay factors + stage_tile(sm.a.s_w, w + t0 * (int64_t)p.HV * p.K, (int64_t)p.HV * p.K, + t_valid, p.K, tid); + build_state_T(sm.a.s_hb, sm.s_h, tid); + const int last = min((i_t + 1) * kBT, T) - 1; + load_decay_g(sm.s_g, gk + (int64_t)last * p.HV * p.K, p.K, tid); + __syncthreads(); + + // v_new = u - w @ h + { + auto acc = mma_tn(mma_vn, sm.a.s_w, sm.a.s_hb, tid); + auto tCcC = mma_coords(mma_vn, tid); + BF16 const* u_t = u + t0 * (int64_t)p.HV * p.V; + BF16* vn_t = v_new + t0 * (int64_t)p.HV * p.V; + for (int i = 0; i < int(size(acc)); ++i) { + int m = get<0>(tCcC(i)); + int n = get<1>(tCcC(i)); + int vg = v0 + n; + bool valid = (m < t_valid) && (vg < p.V); + float uu = valid ? to_f32(u_t[(int64_t)m * p.HV * p.V + vg]) : 0.f; + float val = uu - acc(i); + if (valid) vn_t[(int64_t)m * p.HV * p.V + vg] = BF16(val); + sm.s_vnT[n * kBT + m] = BF16(val); + } + } + __syncthreads(); // s_kgT overlaps s_w / s_hb + stage_tile_T(sm.b.s_kgT, kg + t0 * (int64_t)p.H * p.K, (int64_t)p.H * p.K, + t_valid, p.K, tid); + decay_state(sm.s_h, sm.s_g, tid); + __syncthreads(); + + // h += kg^T @ v_new + { + auto acc = mma_tn(mma_up, sm.b.s_kgT, sm.s_vnT, tid); + auto tCcC = mma_coords(mma_up, tid); + for (int i = 0; i < int(size(acc)); ++i) { + int k = get<0>(tCcC(i)); + int n = get<1>(tCcC(i)); + sm.s_h[k * BV + n] += acc(i); + } + } + __syncthreads(); + } + + if (p.ht != nullptr) { + store_state_gmem(p.ht + i_nh * (int64_t)p.K * p.V, + sm.s_h, p.K, p.V, v0, p.state_v_first, tid); + } +} + +// --------------------------------------------------------------------------- +// backward: per chunk (in reverse), dh[i_t] = S; +// dv2 = dv + kg @ S; S = S * exp2(gk_last) + scale * qg^T @ do - w^T @ dv2 + +template +__global__ void __launch_bounds__(kThreads) kda_bwd_dhu_kernel(BwdParams const p) { + extern __shared__ __align__(128) unsigned char smem_raw[]; + auto& sm = *reinterpret_cast*>(smem_raw); + const int tid = threadIdx.x; + + const int i_v = blockIdx.x % p.NV; + const int64_t i_nh = blockIdx.x / p.NV; + const int64_t i_n = i_nh / p.HV; + const int i_hv = int(i_nh % p.HV); + const int i_hq = i_hv / (p.HV / p.H); + + int64_t bos; + int T, NT; + int64_t boh; + if (p.varlen) { + bos = p.cu_seqlens[i_n]; + T = int(p.cu_seqlens[i_n + 1] - bos); + NT = (T + kBT - 1) / kBT; + boh = p.chunk_offsets[i_n]; + } else { + bos = i_n * p.T; + T = p.T; + NT = (T + kBT - 1) / kBT; + boh = i_n * NT; + } + + BF16 const* qg = p.qg + (bos * p.H + i_hq) * (int64_t)p.K; + BF16 const* kg = p.kg + (bos * p.H + i_hq) * (int64_t)p.K; + BF16 const* w = p.w + (bos * p.HV + i_hv) * (int64_t)p.K; + float const* gk = p.gk + (bos * p.HV + i_hv) * (int64_t)p.K; + BF16 const* do_ = p.do_ + (bos * p.HV + i_hv) * (int64_t)p.V; + BF16 const* dv = p.dv + (bos * p.HV + i_hv) * (int64_t)p.V; + BF16* dv2 = p.dv2 + (bos * p.HV + i_hv) * (int64_t)p.V; + BF16* dh = p.dh + (boh * p.HV + i_hv) * (int64_t)p.K * p.V; + const int v0 = i_v * BV; + + load_state_gmem(sm.s_dh, + p.dht == nullptr ? nullptr : p.dht + i_nh * (int64_t)p.K * p.V, + p.K, p.V, v0, p.state_v_first, tid); + + auto mma_dv = make_tiled_mma(MmaAtom{}, Layout>{}); + auto mma_up = make_tiled_mma(MmaAtom{}, Layout>{}); + + __syncthreads(); + for (int i_t = NT - 1; i_t >= 0; --i_t) { + const int t_valid = min(kBT, T - i_t * kBT); + const int64_t t0 = (int64_t)i_t * kBT; + + // dh[i_t] = state gradient at chunk entry (bf16) + store_state_gmem(dh + (int64_t)i_t * p.HV * p.K * p.V, + sm.s_dh, p.K, p.V, v0, p.state_v_first, tid); + stage_tile(sm.a.s_kg, kg + t0 * (int64_t)p.H * p.K, (int64_t)p.H * p.K, + t_valid, p.K, tid); + build_state_T(sm.a.s_dhb, sm.s_dh, tid); + const int last = min((i_t + 1) * kBT, T) - 1; + load_decay_g(sm.s_g, gk + (int64_t)last * p.HV * p.K, p.K, tid); + __syncthreads(); + + // dv2 = dv + kg @ dh + { + auto acc = mma_tn(mma_dv, sm.a.s_kg, sm.a.s_dhb, tid); + auto tCcC = mma_coords(mma_dv, tid); + BF16 const* dv_t = dv + t0 * (int64_t)p.HV * p.V; + BF16* dv2_t = dv2 + t0 * (int64_t)p.HV * p.V; + for (int i = 0; i < int(size(acc)); ++i) { + int m = get<0>(tCcC(i)); + int n = get<1>(tCcC(i)); + int vg = v0 + n; + bool valid = (m < t_valid) && (vg < p.V); + float dvv = valid ? to_f32(dv_t[(int64_t)m * p.HV * p.V + vg]) : 0.f; + float val = acc(i) + dvv; + if (valid) dv2_t[(int64_t)m * p.HV * p.V + vg] = BF16(val); + sm.s_dv2T[n * kBT + m] = BF16(val); + } + } + __syncthreads(); // phase-b buffers overlap s_kg / s_dhb + stage_tile_T(sm.b.s_qwT, qg + t0 * (int64_t)p.H * p.K, (int64_t)p.H * p.K, + t_valid, p.K, tid); + stage_tile_T(sm.b.s_doT, do_ + t0 * (int64_t)p.HV * p.V + v0, + (int64_t)p.HV * p.V, t_valid, min(BV, p.V - v0), tid); + decay_state(sm.s_dh, sm.s_g, tid); + __syncthreads(); + + // dh += scale * qg^T @ do + { + auto acc = mma_tn(mma_up, sm.b.s_qwT, sm.b.s_doT, tid); + auto tCcC = mma_coords(mma_up, tid); + for (int i = 0; i < int(size(acc)); ++i) { + int k = get<0>(tCcC(i)); + int n = get<1>(tCcC(i)); + sm.s_dh[k * BV + n] += p.scale * acc(i); + } + } + __syncthreads(); // restage s_qwT with w + stage_tile_T(sm.b.s_qwT, w + t0 * (int64_t)p.HV * p.K, (int64_t)p.HV * p.K, + t_valid, p.K, tid); + __syncthreads(); + // dh -= w^T @ dv2 + { + auto acc = mma_tn(mma_up, sm.b.s_qwT, sm.s_dv2T, tid); + auto tCcC = mma_coords(mma_up, tid); + for (int i = 0; i < int(size(acc)); ++i) { + int k = get<0>(tCcC(i)); + int n = get<1>(tCcC(i)); + sm.s_dh[k * BV + n] -= acc(i); + } + } + __syncthreads(); + } + + if (p.dh0 != nullptr) { + store_state_gmem(p.dh0 + i_nh * (int64_t)p.K * p.V, + sm.s_dh, p.K, p.V, v0, p.state_v_first, tid); + } +} + +// --------------------------------------------------------------------------- +// launchers + +template +void launch_fwd_pipe(FwdParams& p, int64_t grid_blocks, cudaStream_t stream) { + constexpr int kSmemBytes = int(sizeof(FwdPipeSmem)); + static bool configured = [] { + cudaFuncSetAttribute(kda_fwd_h_pipe_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes); + return true; + }(); + (void)configured; + kda_fwd_h_pipe_kernel<<>>(p); +} + +template +void launch_bwd_pipe(BwdParams& p, int64_t grid_blocks, cudaStream_t stream) { + constexpr int kSmemBytes = int(sizeof(BwdPipeSmem)); + static bool configured = [] { + cudaFuncSetAttribute(kda_bwd_dhu_pipe_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes); + return true; + }(); + (void)configured; + kda_bwd_dhu_pipe_kernel<<>>(p); +} + +template +void launch_fwd(FwdParams& p, int64_t grid_blocks, cudaStream_t stream) { + constexpr int kSmemBytes = int(sizeof(FwdSmem)); + static bool configured = [] { + cudaFuncSetAttribute(kda_fwd_h_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes); + return true; + }(); + (void)configured; + kda_fwd_h_kernel<<>>(p); +} + +template +void launch_bwd(BwdParams& p, int64_t grid_blocks, cudaStream_t stream) { + constexpr int kSmemBytes = int(sizeof(BwdSmem)); + static bool configured = [] { + cudaFuncSetAttribute(kda_bwd_dhu_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes); + return true; + }(); + (void)configured; + kda_bwd_dhu_kernel<<>>(p); +} + +// Padded K (multiple of 64) and the V tile used for it. BV=32 above K=128 keeps +// dynamic smem under the 99KB opt-in limit of sm_120. +int padded_k(int64_t K) { + TORCH_CHECK(K > 0 && K <= 256, "K must be in (0, 256]"); + return int((K + 63) / 64 * 64); +} + +int v_tile(int Kp) { + return Kp > 64 ? 32 : 64; +} + +} // anonymous namespace (internal linkage: the test harness JIT-loads this file as a +// second module alongside the pip extension; named-namespace kernels interpose +// across modules and the smem opt-in attribute then lands on the wrong kernel) + +// --------------------------------------------------------------------------- +// host entry points (fla chunk_gated_delta_rule_fwd_h / _bwd_dhu equivalents, +// KDA call shape only: USE_G=False, USE_GK=True, SAVE_NEW_VALUE=True) + +std::tuple> +chunk_gated_delta_rule_fwd_h( + torch::Tensor kg, + torch::Tensor w, + torch::Tensor u, + torch::Tensor gk, + std::optional initial_state, + bool output_final_state, + int64_t chunk_size, + bool state_v_first, + std::optional cu_seqlens, + std::optional chunk_offsets, + int64_t nt_total +) { + TORCH_CHECK(chunk_size == kBT, "only chunk_size=64 is supported"); + TORCH_CHECK(kg.is_cuda() && kg.is_contiguous() && kg.dim() == 4, "kg must be 4D contiguous CUDA"); + TORCH_CHECK(kg.scalar_type() == at::kBFloat16, "only bf16 operands are supported"); + int64_t B = kg.size(0), T = kg.size(1), H = kg.size(2), K = kg.size(3); + TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.dim() == 4 && w.scalar_type() == at::kBFloat16); + TORCH_CHECK(u.is_cuda() && u.is_contiguous() && u.dim() == 4 && u.scalar_type() == at::kBFloat16); + int64_t HV = w.size(2), V = u.size(3); + TORCH_CHECK(w.size(0) == B && w.size(1) == T && w.size(3) == K); + TORCH_CHECK(u.size(0) == B && u.size(1) == T && u.size(2) == HV); + TORCH_CHECK(gk.is_cuda() && gk.is_contiguous() && gk.scalar_type() == at::kFloat); + TORCH_CHECK(gk.dim() == 4 && gk.size(0) == B && gk.size(1) == T && gk.size(2) == HV && gk.size(3) == K); + TORCH_CHECK(HV % H == 0, "HV must be a multiple of H"); + + bool varlen = cu_seqlens.has_value(); + int64_t N, NTt; + if (varlen) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_offsets.has_value(), "chunk_offsets must be provided when cu_seqlens is provided"); + TORCH_CHECK(cu_seqlens->is_cuda() && cu_seqlens->is_contiguous() && cu_seqlens->dtype() == torch::kLong); + TORCH_CHECK(chunk_offsets->is_cuda() && chunk_offsets->is_contiguous() && chunk_offsets->dtype() == torch::kLong); + N = cu_seqlens->size(0) - 1; + NTt = nt_total; + TORCH_CHECK(NTt > 0, "nt_total must be provided for varlen"); + } else { + N = B; + NTt = (T + kBT - 1) / kBT; + } + + if (initial_state.has_value()) { + auto const& h0 = initial_state.value(); + TORCH_CHECK(h0.is_cuda() && h0.is_contiguous() && h0.scalar_type() == at::kFloat); + TORCH_CHECK(h0.dim() == 4 && h0.size(0) == N && h0.size(1) == HV); + } + + torch::Tensor h = state_v_first + ? torch::empty({B, NTt, HV, V, K}, kg.options()) + : torch::empty({B, NTt, HV, K, V}, kg.options()); + torch::Tensor v_new = torch::empty_like(u); + std::optional final_state; + if (output_final_state) { + final_state = state_v_first + ? torch::zeros({N, HV, V, K}, kg.options().dtype(at::kFloat)) + : torch::zeros({N, HV, K, V}, kg.options().dtype(at::kFloat)); + } + + FwdParams p; + p.kg = reinterpret_cast(kg.data_ptr()); + p.w = reinterpret_cast(w.data_ptr()); + p.u = reinterpret_cast(u.data_ptr()); + p.gk = gk.data_ptr(); + p.h = reinterpret_cast(h.data_ptr()); + p.v_new = reinterpret_cast(v_new.data_ptr()); + p.h0 = initial_state.has_value() ? initial_state->data_ptr() : nullptr; + p.ht = final_state.has_value() ? final_state->data_ptr() : nullptr; + p.cu_seqlens = varlen ? cu_seqlens->data_ptr() : nullptr; + p.chunk_offsets = varlen ? chunk_offsets->data_ptr() : nullptr; + p.T = int(T); p.H = int(H); p.HV = int(HV); p.K = int(K); p.V = int(V); + p.varlen = varlen; + p.state_v_first = state_v_first; + + const int Kp = padded_k(K); + p.NV = int((V + v_tile(Kp) - 1) / v_tile(Kp)); + int64_t grid_blocks = (int64_t)p.NV * N * HV; + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + switch (Kp) { + case 64: launch_fwd_pipe<64, 64>(p, grid_blocks, stream); break; + case 128: launch_fwd_pipe<128, 32>(p, grid_blocks, stream); break; + case 192: launch_fwd<192, 32>(p, grid_blocks, stream); break; + default: launch_fwd<256, 32>(p, grid_blocks, stream); break; + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {h, v_new, final_state}; +} + +std::tuple, torch::Tensor> +chunk_gated_delta_rule_bwd_dhu( + torch::Tensor qg, + torch::Tensor kg, + torch::Tensor w, + torch::Tensor gk, + torch::Tensor do_, + torch::Tensor dv, + std::optional h0, + std::optional dht, + double scale, + int64_t chunk_size, + bool state_v_first, + std::optional cu_seqlens, + std::optional chunk_offsets, + int64_t nt_total +) { + TORCH_CHECK(chunk_size == kBT, "only chunk_size=64 is supported"); + TORCH_CHECK(qg.is_cuda() && qg.is_contiguous() && qg.dim() == 4 && qg.scalar_type() == at::kBFloat16); + int64_t B = qg.size(0), T = qg.size(1), H = qg.size(2), K = qg.size(3); + TORCH_CHECK(kg.is_cuda() && kg.is_contiguous() && kg.sizes() == qg.sizes() && kg.scalar_type() == at::kBFloat16); + TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.dim() == 4 && w.scalar_type() == at::kBFloat16); + int64_t HV = w.size(2); + TORCH_CHECK(do_.is_cuda() && do_.is_contiguous() && do_.dim() == 4 && do_.scalar_type() == at::kBFloat16); + int64_t V = do_.size(3); + TORCH_CHECK(dv.is_cuda() && dv.is_contiguous() && dv.sizes() == do_.sizes() && dv.scalar_type() == at::kBFloat16); + TORCH_CHECK(gk.is_cuda() && gk.is_contiguous() && gk.scalar_type() == at::kFloat); + TORCH_CHECK(gk.dim() == 4 && gk.size(0) == B && gk.size(1) == T && gk.size(2) == HV && gk.size(3) == K); + TORCH_CHECK(HV % H == 0, "HV must be a multiple of H"); + + bool varlen = cu_seqlens.has_value(); + int64_t N, NTt; + if (varlen) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_offsets.has_value(), "chunk_offsets must be provided when cu_seqlens is provided"); + TORCH_CHECK(cu_seqlens->is_cuda() && cu_seqlens->is_contiguous() && cu_seqlens->dtype() == torch::kLong); + TORCH_CHECK(chunk_offsets->is_cuda() && chunk_offsets->is_contiguous() && chunk_offsets->dtype() == torch::kLong); + N = cu_seqlens->size(0) - 1; + NTt = nt_total; + TORCH_CHECK(NTt > 0, "nt_total must be provided for varlen"); + } else { + N = B; + NTt = (T + kBT - 1) / kBT; + } + + torch::Tensor dh = state_v_first + ? torch::empty({B, NTt, HV, V, K}, qg.options()) + : torch::empty({B, NTt, HV, K, V}, qg.options()); + std::optional dh0; + if (h0.has_value()) { + TORCH_CHECK(h0->is_cuda() && h0->is_contiguous() && h0->scalar_type() == at::kFloat); + dh0 = torch::empty_like(h0.value(), h0->options().dtype(at::kFloat)); + } + if (dht.has_value()) { + TORCH_CHECK(dht->is_cuda() && dht->is_contiguous() && dht->scalar_type() == at::kFloat); + } + torch::Tensor dv2 = torch::empty_like(dv); + + BwdParams p; + p.qg = reinterpret_cast(qg.data_ptr()); + p.kg = reinterpret_cast(kg.data_ptr()); + p.w = reinterpret_cast(w.data_ptr()); + p.gk = gk.data_ptr(); + p.do_ = reinterpret_cast(do_.data_ptr()); + p.dv = reinterpret_cast(dv.data_ptr()); + p.dv2 = reinterpret_cast(dv2.data_ptr()); + p.dh = reinterpret_cast(dh.data_ptr()); + p.dh0 = dh0.has_value() ? dh0->data_ptr() : nullptr; + p.dht = dht.has_value() ? dht->data_ptr() : nullptr; + p.cu_seqlens = varlen ? cu_seqlens->data_ptr() : nullptr; + p.chunk_offsets = varlen ? chunk_offsets->data_ptr() : nullptr; + p.scale = float(scale); + p.T = int(T); p.H = int(H); p.HV = int(HV); p.K = int(K); p.V = int(V); + p.varlen = varlen; + p.state_v_first = state_v_first; + + const int Kp = padded_k(K); + p.NV = int((V + v_tile(Kp) - 1) / v_tile(Kp)); + int64_t grid_blocks = (int64_t)p.NV * N * HV; + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + switch (Kp) { + case 64: launch_bwd_pipe<64, 64>(p, grid_blocks, stream); break; + case 128: launch_bwd_pipe<128, 32>(p, grid_blocks, stream); break; + case 192: launch_bwd<192, 32>(p, grid_blocks, stream); break; + default: launch_bwd<256, 32>(p, grid_blocks, stream); break; + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dh, dh0, dv2}; +} diff --git a/csrc/train/chunk_h_binding.cpp b/csrc/train/chunk_h_binding.cpp new file mode 100644 index 0000000..b909414 --- /dev/null +++ b/csrc/train/chunk_h_binding.cpp @@ -0,0 +1,59 @@ +#include + +std::tuple> +chunk_gated_delta_rule_fwd_h( + torch::Tensor kg, + torch::Tensor w, + torch::Tensor u, + torch::Tensor gk, + std::optional initial_state, + bool output_final_state, + int64_t chunk_size, + bool state_v_first, + std::optional cu_seqlens, + std::optional chunk_offsets, + int64_t nt_total +); + +std::tuple, torch::Tensor> +chunk_gated_delta_rule_bwd_dhu( + torch::Tensor qg, + torch::Tensor kg, + torch::Tensor w, + torch::Tensor gk, + torch::Tensor do_, + torch::Tensor dv, + std::optional h0, + std::optional dht, + double scale, + int64_t chunk_size, + bool state_v_first, + std::optional cu_seqlens, + std::optional chunk_offsets, + int64_t nt_total +); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("chunk_gated_delta_rule_fwd_h", &chunk_gated_delta_rule_fwd_h, + "KDA chunked state forward h (CUDA)", + py::arg("kg"), py::arg("w"), py::arg("u"), py::arg("gk"), + py::arg("initial_state") = py::none(), + py::arg("output_final_state") = false, + py::arg("chunk_size") = 64, + py::arg("state_v_first") = false, + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_offsets") = py::none(), + py::arg("nt_total") = 0); + m.def("chunk_gated_delta_rule_bwd_dhu", &chunk_gated_delta_rule_bwd_dhu, + "KDA chunked state backward dhu (CUDA)", + py::arg("qg"), py::arg("kg"), py::arg("w"), py::arg("gk"), + py::arg("do_"), py::arg("dv"), + py::arg("h0") = py::none(), + py::arg("dht") = py::none(), + py::arg("scale") = 1.0, + py::arg("chunk_size") = 64, + py::arg("state_v_first") = false, + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_offsets") = py::none(), + py::arg("nt_total") = 0); +} diff --git a/csrc/train/chunk_o.cu b/csrc/train/chunk_o.cu new file mode 100644 index 0000000..9714497 --- /dev/null +++ b/csrc/train/chunk_o.cu @@ -0,0 +1,420 @@ +// Output kernel of the chunked GLA forward, reused by KDA. +// Replicates fla/ops/gla/chunk.py::chunk_gla_fwd_kernel_o (BT=64, BK=BV=64). +// +// Per (v-tile, chunk, hv-head) block: +// o = scale * (q * exp2(g)) @ h + tril(Aqk, 0) @ v_new +// q*exp2(g) is computed in fp32 and rounded to the input dtype before the dot, +// as the Triton kernel does. h is bf16, [B, NT, HV, K, V] or [B, NT, HV, V, K] +// when state_v_first. GEMMs run on tensor cores (SM80 16x8x16, fp32 accum). +// +// Data movement: pure-copy tiles (h, v) are staged with 16B cp.async (masked +// rows zero-filled via src-size 0), computed tiles (q*exp2(g), tril(A)) use +// 16B vector loads/stores, smem rows are padded +8 halves against bank +// conflicts, MMA fragments load with ldmatrix, and the output goes through smem +// staging for 16B coalesced stores. GEMM/accumulation order and rounding +// points are unchanged from v1. + +#include +#include +#include + +#include + +#include "common.cuh" + +// NOTE: named namespace, not anonymous — nvcc's launch stub generation +// mis-resolves anonymous-namespace kernels when cute headers are included. +namespace chunk_o_impl { + +using namespace cute; + +template struct MmaAtom; +template <> struct MmaAtom { using type = SM80_16x8x16_F32BF16BF16F32_TN; }; +template <> struct MmaAtom { using type = SM80_16x8x16_F32F16F16F32_TN; }; + +constexpr int kBT = 64; +constexpr int kThreads = 128; +constexpr int kPad = 8; // smem row padding in halves (16B) +constexpr int kCP = kBT + kPad; // padded row stride of every 64-wide tile + +__device__ __forceinline__ void cp_async16(void* dst, void const* src, bool full) { + uint32_t s = cute::cast_smem_ptr_to_uint(dst); + int sz = full ? 16 : 0; // src-size 0 zero-fills, used for masked rows/cols + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" + :: "r"(s), "l"(src), "r"(sz) : "memory"); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory"); +} + +// s[r][c] = g[r*row_stride + c] via 16B cp.async; rows past rows_valid and col +// chunks at/past cols_valid zero-filled. Requires row_stride and cols_valid to +// be multiples of 8 halves (16B). +template +__device__ __forceinline__ void stage_tile(T* s, T const* g, int64_t row_stride, + int rows_valid, int cols_valid, int tid) { + constexpr int kCG = kBT / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + bool full = (r < rows_valid) && (c < cols_valid); + cp_async16(s + r * kCP + c, full ? g + (int64_t)r * row_stride + c : g, full); + } +} + +template +__device__ __forceinline__ void stage_tile_scalar(T* s, T const* g, int64_t row_stride, + int rows_valid, int cols_valid, int tid) { + for (int idx = tid; idx < kBT * kBT; idx += kThreads) { + int r = idx / kBT, c = idx % kBT; + s[r * kCP + c] = (r < rows_valid && c < cols_valid) ? g[(int64_t)r * row_stride + c] : T(0.0f); + } +} + +// Grid: (cdiv(V, 64), NT, B*HV). Block: 128 threads. +template +__global__ void __launch_bounds__(kThreads) chunk_gla_fwd_o_gk_kernel( + T const* __restrict__ q, + T const* __restrict__ v, + float const* __restrict__ g, + T const* __restrict__ h, + T const* __restrict__ A, + T* __restrict__ o, + float scale, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int64_t T_len, + int H, + int HV, + int K, + int V +) { + int i_v = blockIdx.x; + int64_t i_t = blockIdx.y; + int64_t i_bh = blockIdx.z; + int64_t i_b = i_bh / HV, i_hv = i_bh % HV; + int i_h = int(i_hv / (HV / H)); + + int64_t bos, t0, seq_len, i_tg; + if (IS_VARLEN) { + i_tg = i_t; // the grid chunk index is already the global h chunk index + int64_t i_n = chunk_indices[i_t * 2]; + int64_t i_tl = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + seq_len = cu_seqlens[i_n + 1] - bos; + t0 = i_tl * kBT; + } else { + int64_t NT = (T_len + kBT - 1) / kBT; + i_tg = i_b * NT + i_t; + bos = i_b * T_len; + seq_len = T_len; + t0 = i_t * kBT; + } + int64_t rem = seq_len - t0; + int rows_valid = int(rem < (int64_t)kBT ? rem : (int64_t)kBT); + int v0 = i_v * kBT; + int cols_valid = V - v0 < kBT ? V - v0 : kBT; + if (rows_valid <= 0 || cols_valid <= 0) return; + + const int64_t tok = bos + t0; + T const* qp = q + (tok * H + i_h) * (int64_t)K; + float const* gp = g + (tok * HV + i_hv) * (int64_t)K; + T const* vp = v + (tok * HV + i_hv) * (int64_t)V; + T* op = o + (tok * HV + i_hv) * (int64_t)V; + T const* hp = h + (i_tg * HV + i_hv) * (int64_t)K * V; + T const* Ap = A + (tok * HV + i_hv) * (int64_t)kBT; + + __shared__ T sA[kBT * kCP]; + __shared__ T sB[kBT * kCP]; + + int const tid = threadIdx.x; + bool const vec_ok = (K % 8) == 0 && (V % 8) == 0; + + using Atom = typename MmaAtom::type; + auto mma = make_tiled_mma(Atom{}, Layout>{}, Tile<_64, _64, _16>{}); + auto thr_mma = mma.get_thread_slice(tid); + + Copy_Atom ldsm_n; + Copy_Atom ldsm_t; + auto s2r_a = make_tiled_copy_A(ldsm_n, mma); // row-major (M,K) + auto s2r_b = make_tiled_copy_B(ldsm_n, mma); // row-major (N,K) + auto s2r_bt = make_tiled_copy_B(ldsm_t, mma); // strided (N,K) view + auto thr_s2r_a = s2r_a.get_thread_slice(tid); + auto thr_s2r_b = s2r_b.get_thread_slice(tid); + auto thr_s2r_bt = s2r_bt.get_thread_slice(tid); + + Tensor sA_rm = make_tensor(make_smem_ptr(sA), Layout, Stride, _1>>{}); + Tensor sB_rm = make_tensor(make_smem_ptr(sB), Layout, Stride, _1>>{}); + Tensor sB_st = make_tensor(make_smem_ptr(sB), Layout, Stride<_1, Int>>{}); + + Tensor cC = make_identity_tensor(Shape<_64, _64>{}); + Tensor tCcC = thr_mma.partition_C(cC); + Tensor rC = thr_mma.make_fragment_C(tCcC); + clear(rC); + + // rC += sA[64(M),64(K)] @ sB[64(N),64(K)]^T, K sliced 16-wide, ascending. + auto gemm_acc = [&](auto const& sA_t, auto const& sB_t, auto b_strided) { + constexpr bool kBStrided = decltype(b_strided)::value; + Tensor tCrA = thr_mma.partition_fragment_A(sA_t); + Tensor tCrB = thr_mma.partition_fragment_B(sB_t); + Tensor tXsA = thr_s2r_a.partition_S(sA_t); + Tensor tXrA = thr_s2r_a.retile_D(tCrA); + constexpr int KB = decltype(size<2>(tXsA))::value; + if constexpr (kBStrided) { + Tensor tXsB = thr_s2r_bt.partition_S(sB_t); + Tensor tXrB = thr_s2r_bt.retile_D(tCrB); + CUTE_UNROLL + for (int kb = 0; kb < KB; ++kb) { + copy(s2r_a, tXsA(_, _, kb), tXrA(_, _, kb)); + copy(s2r_bt, tXsB(_, _, kb), tXrB(_, _, kb)); + gemm(mma, tCrA(_, _, kb), tCrB(_, _, kb), rC); + } + } else { + Tensor tXsB = thr_s2r_b.partition_S(sB_t); + Tensor tXrB = thr_s2r_b.retile_D(tCrB); + CUTE_UNROLL + for (int kb = 0; kb < KB; ++kb) { + copy(s2r_a, tXsA(_, _, kb), tXrA(_, _, kb)); + copy(s2r_b, tXsB(_, _, kb), tXrB(_, _, kb)); + gemm(mma, tCrA(_, _, kb), tCrB(_, _, kb), rC); + } + } + }; + + // inter part: sum over K tiles of (q * exp2(g)) @ h + for (int k0 = 0; k0 < K; k0 += kBT) { + int kcols = K - k0 < kBT ? K - k0 : kBT; + if (vec_ok) { + // sA[r][c] = T(q[r, k0+c] * exp2f(g[r, k0+c])), 8 cols per thread + constexpr int kCG = kBT / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + T vals[8] = {}; + if (r < rows_valid && c < kcols) { + uint4 rq = *reinterpret_cast(qp + r * (int64_t)H * K + k0 + c); + T const* hq = reinterpret_cast(&rq); + float4 const* pg = reinterpret_cast(gp + r * (int64_t)HV * K + k0 + c); + float4 g0 = pg[0], g1 = pg[1]; + vals[0] = T(to_f32(hq[0]) * exp2f(g0.x)); + vals[1] = T(to_f32(hq[1]) * exp2f(g0.y)); + vals[2] = T(to_f32(hq[2]) * exp2f(g0.z)); + vals[3] = T(to_f32(hq[3]) * exp2f(g0.w)); + vals[4] = T(to_f32(hq[4]) * exp2f(g1.x)); + vals[5] = T(to_f32(hq[5]) * exp2f(g1.y)); + vals[6] = T(to_f32(hq[6]) * exp2f(g1.z)); + vals[7] = T(to_f32(hq[7]) * exp2f(g1.w)); + } + *reinterpret_cast(sA + r * kCP + c) = *reinterpret_cast(vals); + } + // sB: h tile, pure copy + if (!STATE_V_FIRST) { + // h is [K, V]: rows k0+r, cols v0+c; row-dim validity is kcols + stage_tile(sB, hp + (int64_t)k0 * V + v0, V, kcols, cols_valid, tid); + } else { + // h is [V, K]: rows v0+r, cols k0+c + stage_tile(sB, hp + (int64_t)v0 * K + k0, K, cols_valid, kcols, tid); + } + cp_async_commit(); + cp_async_wait<0>(); + } else { + for (int idx = tid; idx < kBT * kBT; idx += kThreads) { + int r = idx >> 6, c = idx & 63; + float val = 0.0f; + if (r < rows_valid && c < kcols) { + val = to_f32(qp[r * (int64_t)H * K + k0 + c]) * exp2f(gp[r * (int64_t)HV * K + k0 + c]); + } + sA[r * kCP + c] = T(val); + } + if (!STATE_V_FIRST) { + stage_tile_scalar(sB, hp + (int64_t)k0 * V + v0, V, kcols, cols_valid, tid); + } else { + stage_tile_scalar(sB, hp + (int64_t)v0 * K + k0, K, cols_valid, kcols, tid); + } + } + __syncthreads(); + if constexpr (STATE_V_FIRST) { gemm_acc(sA_rm, sB_rm, cute::false_type{}); } + else { gemm_acc(sA_rm, sB_st, cute::true_type{}); } + __syncthreads(); + } + + for (int i = 0; i < size(rC); ++i) rC(i) *= scale; + + // intra part: tril(Aqk, 0) @ v_new (Aqk already carries the scale) + if (vec_ok) { + constexpr int kCG = kBT / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + uint4 raw = make_uint4(0, 0, 0, 0); + if (r < rows_valid && r >= c) { + // tril mask within the 8-col group: keep cols c+j <= r + uint4 full = *reinterpret_cast(Ap + r * (int64_t)HV * kBT + c); + T const* fa = reinterpret_cast(&full); + T vals[8]; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) vals[j] = (c + j <= r) ? fa[j] : T(0.0f); + raw = *reinterpret_cast(vals); + } + *reinterpret_cast(sA + r * kCP + c) = raw; + } + stage_tile(sB, vp + v0, (int64_t)HV * V, rows_valid, cols_valid, tid); + cp_async_commit(); + cp_async_wait<0>(); + } else { + for (int idx = tid; idx < kBT * kBT; idx += kThreads) { + int r = idx >> 6, c = idx & 63; + sA[r * kCP + c] = (r < rows_valid && r >= c) ? Ap[r * (int64_t)HV * kBT + c] : T(0.0f); + } + stage_tile_scalar(sB, vp + v0, (int64_t)HV * V, rows_valid, cols_valid, tid); + } + __syncthreads(); + gemm_acc(sA_rm, sB_st, cute::true_type{}); + __syncthreads(); // sA is dead; reuse it as the o staging tile + + for (int i = 0; i < size(rC); ++i) { + sA[get<0>(tCcC(i)) * kCP + get<1>(tCcC(i))] = T(rC(i)); + } + __syncthreads(); + if (vec_ok) { + constexpr int kCG = kBT / 8; + CUTE_UNROLL + for (int idx = tid; idx < kBT * kCG; idx += kThreads) { + int r = idx / kCG, c = (idx % kCG) * 8; + if (r < rows_valid && c < cols_valid) { + *reinterpret_cast(op + r * (int64_t)HV * V + v0 + c) = + *reinterpret_cast(sA + r * kCP + c); + } + } + } else { + for (int idx = tid; idx < kBT * kBT; idx += kThreads) { + int r = idx >> 6, c = idx & 63; + if (r < rows_valid && c < cols_valid) { + op[r * (int64_t)HV * V + v0 + c] = sA[r * kCP + c]; + } + } + } +} + +template +void launch_chunk_gla_fwd_o_gk( + torch::Tensor const& q, + torch::Tensor const& v, + torch::Tensor const& g, + torch::Tensor const& A, + torch::Tensor const& h, + torch::Tensor& o, + float scale, + bool state_v_first, + int64_t const* cu_seqlens, + int64_t const* chunk_indices, + int64_t NT, + int64_t B, + int64_t T_len, + int H, + int HV, + int K, + int V, + cudaStream_t stream +) { + dim3 grid((V + kBT - 1) / kBT, NT, B * HV); + dim3 block(kThreads); + + #define LAUNCH_O(SVF, IS_VARLEN) \ + chunk_gla_fwd_o_gk_kernel<<>>( \ + reinterpret_cast(q.data_ptr()), \ + reinterpret_cast(v.data_ptr()), \ + g.data_ptr(), \ + reinterpret_cast(h.data_ptr()), \ + reinterpret_cast(A.data_ptr()), \ + reinterpret_cast(o.data_ptr()), \ + scale, cu_seqlens, chunk_indices, T_len, H, HV, K, V) + + if (state_v_first) { + if (cu_seqlens) { LAUNCH_O(true, true); } else { LAUNCH_O(true, false); } + } else { + if (cu_seqlens) { LAUNCH_O(false, true); } else { LAUNCH_O(false, false); } + } + #undef LAUNCH_O +} + +} // namespace chunk_o_impl + +using chunk_o_impl::kBT; +using chunk_o_impl::launch_chunk_gla_fwd_o_gk; + +torch::Tensor chunk_gla_fwd_o_gk( + torch::Tensor q, + torch::Tensor v, + torch::Tensor g, + torch::Tensor A, + torch::Tensor h, + double scale, + bool state_v_first, + std::optional cu_seqlens, + int64_t chunk_size, + std::optional chunk_indices +) { + TORCH_CHECK(q.is_cuda() && q.is_contiguous() && q.dim() == 4, "q must be [B, T, H, K] contiguous CUDA"); + TORCH_CHECK(v.is_cuda() && v.is_contiguous() && v.dim() == 4, "v must be [B, T, HV, V] contiguous CUDA"); + TORCH_CHECK(v.scalar_type() == q.scalar_type(), "v and q must share dtype"); + TORCH_CHECK(g.is_cuda() && g.is_contiguous() && g.scalar_type() == torch::kFloat32, + "g must be fp32 [B, T, HV, K] contiguous CUDA"); + TORCH_CHECK(A.is_cuda() && A.is_contiguous() && A.dim() == 4, "A must be [B, T, HV, BT] contiguous CUDA"); + TORCH_CHECK(A.scalar_type() == q.scalar_type(), "A and q must share dtype"); + TORCH_CHECK(h.is_cuda() && h.is_contiguous() && h.dim() == 5, "h must be 5D contiguous CUDA"); + TORCH_CHECK(h.scalar_type() == q.scalar_type(), "h and q must share dtype"); + TORCH_CHECK(chunk_size == kBT, "only chunk_size 64 is supported"); + + int64_t B = q.size(0), T_len = q.size(1); + int H = int(q.size(2)), K = int(q.size(3)); + int HV = int(v.size(2)), V = int(v.size(3)); + TORCH_CHECK(HV % H == 0, "HV must be a multiple of H"); + if (state_v_first) { + TORCH_CHECK(h.size(3) == V && h.size(4) == K, "h must be [B, NT, HV, V, K]"); + } else { + TORCH_CHECK(h.size(3) == K && h.size(4) == V, "h must be [B, NT, HV, K, V]"); + } + + int64_t const* cu_ptr = nullptr; + int64_t const* ci_ptr = nullptr; + int64_t NT; + if (cu_seqlens.has_value()) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_indices.has_value(), "chunk_indices required with cu_seqlens"); + auto const& cu = cu_seqlens.value(); + auto const& ci = chunk_indices.value(); + TORCH_CHECK(cu.dtype() == torch::kLong && cu.is_cuda() && cu.is_contiguous()); + TORCH_CHECK(ci.is_cuda() && ci.is_contiguous() && ci.dtype() == torch::kLong); + TORCH_CHECK(ci.dim() == 2 && ci.size(1) == 2); + cu_ptr = cu.data_ptr(); + ci_ptr = ci.data_ptr(); + NT = ci.size(0); + } else { + NT = (T_len + kBT - 1) / kBT; + } + + auto o = torch::zeros_like(v); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + if (q.scalar_type() == at::kBFloat16) { + launch_chunk_gla_fwd_o_gk( + q, v, g, A, h, o, float(scale), state_v_first, + cu_ptr, ci_ptr, NT, B, T_len, H, HV, K, V, stream); + } else if (q.scalar_type() == at::kHalf) { + launch_chunk_gla_fwd_o_gk( + q, v, g, A, h, o, float(scale), state_v_first, + cu_ptr, ci_ptr, NT, B, T_len, H, HV, K, V, stream); + } else { + TORCH_CHECK(false, "chunk_gla_fwd_o_gk supports bf16/fp16 only"); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return o; +} diff --git a/csrc/train/chunk_o_binding.cpp b/csrc/train/chunk_o_binding.cpp new file mode 100644 index 0000000..586c133 --- /dev/null +++ b/csrc/train/chunk_o_binding.cpp @@ -0,0 +1,24 @@ +#include + +torch::Tensor chunk_gla_fwd_o_gk( + torch::Tensor q, + torch::Tensor v, + torch::Tensor g, + torch::Tensor A, + torch::Tensor h, + double scale, + bool state_v_first, + std::optional cu_seqlens, + int64_t chunk_size, + std::optional chunk_indices +); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("chunk_gla_fwd_o_gk", &chunk_gla_fwd_o_gk, "Chunked GLA/KDA forward output (CUDA)", + py::arg("q"), py::arg("v"), py::arg("g"), py::arg("A"), py::arg("h"), + py::arg("scale"), + py::arg("state_v_first") = false, + py::arg("cu_seqlens") = py::none(), + py::arg("chunk_size") = 64, + py::arg("chunk_indices") = py::none()); +} diff --git a/csrc/train/common.cuh b/csrc/train/common.cuh new file mode 100644 index 0000000..f0ba6e8 --- /dev/null +++ b/csrc/train/common.cuh @@ -0,0 +1,48 @@ +#pragma once + +#include +#include +#include + +#include +#include + +// Numeric helpers matching fla's Triton semantics (fla/ops/utils/softplus.py, +// fla/ops/utils/op.py). All gate math is done in fp32. + +__device__ __forceinline__ float to_f32(cutlass::bfloat16_t x) { + float result; + asm("cvt.f32.bf16 %0, %1;\n" : "=f"(result) : "h"(x.storage)); + return result; +} + +__device__ __forceinline__ float to_f32(cutlass::half_t x) { + return float(x); +} + +__device__ __forceinline__ float to_f32(float x) { + return x; +} + +// softplus with threshold 20, identical PTX to fla's softplus_nv +__device__ __forceinline__ float softplus_f32(float x) { + float out; + asm( + "{\n" + ".reg .pred p;\n" + "setp.gt.f32 p, %1, 20.;\n" + "@p mov.f32 %0, %1;\n" + "@!p mul.f32 %0, %1, 1.4426950408889634;\n" + "@!p ex2.approx.ftz.f32 %0, %0;\n" + "@!p add.f32 %0, %0, 1.0;\n" + "@!p lg2.approx.ftz.f32 %0, %0;\n" + "@!p mul.f32 %0, %0, 0.6931471805599453;\n" + "}\n" + : "=f"(out) + : "f"(x)); + return out; +} + +__device__ __forceinline__ float sigmoid_f32(float x) { + return 1.0f / (1.0f + expf(-x)); +} diff --git a/csrc/train/gate.cu b/csrc/train/gate.cu new file mode 100644 index 0000000..4b26836 --- /dev/null +++ b/csrc/train/gate.cu @@ -0,0 +1,465 @@ +// KDA gate activation + chunk-local cumsum, and its backward. +// Replicates fla/ops/kda/gate.py and fla/ops/utils/cumsum.py (vector kernels). + +#include +#include + +#include "common.cuh" +#include "train_ops.h" + +namespace { + +constexpr int kBS = 64; // threads per block, one channel column each + +// Grid: (cdiv(S, BS), NT, B*H). Each thread owns one channel s of one +// (sequence, chunk, head) and walks the chunk's BT rows serially. +template +__global__ void kda_gate_chunk_cumsum_kernel( + T const* __restrict__ g, + float const* __restrict__ A_log, + float const* __restrict__ dt_bias, + float* __restrict__ o, + float scale, + float lower_bound, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int T_len, + int H, + int S, + int BT +) { + int i_s = blockIdx.x * blockDim.x + threadIdx.x; + int64_t i_t = blockIdx.y; + int64_t i_bh = blockIdx.z; + int i_h = int(i_bh % H); + if (i_s >= S) return; + + int64_t bos, i_t0; + int64_t seq_len = T_len; + if (IS_VARLEN) { + int64_t i_n = chunk_indices[i_t * 2]; + int64_t i_tl = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + seq_len = cu_seqlens[i_n + 1] - bos; + i_t0 = i_tl * BT; + } else { + bos = (i_bh / H) * (int64_t)T_len; + i_t0 = i_t * BT; + } + + float b_A = expf(A_log[i_h]); + float bias = 0.0f; + if (HAS_BIAS) bias = dt_bias[i_h * S + i_s]; + + int64_t base = (bos * H + i_h) * (int64_t)S + i_s; + int64_t stride = (int64_t)H * S; + + float run = 0.0f; + for (int t = 0; t < BT; ++t) { + int64_t tt = i_t0 + t; + if (tt >= seq_len) break; + float x = to_f32(g[base + tt * stride]); + if (HAS_BIAS) x += bias; + float gate; + if (USE_LOWER_BOUND) { + gate = lower_bound * sigmoid_f32(b_A * x); + } else { + gate = -b_A * softplus_f32(x); + } + run += gate; + o[base + tt * stride] = HAS_SCALE ? run * scale : run; + } +} + +// Grid: (cdiv(S, BS), NT, B*H). No gate activation; optional reverse (suffix) cumsum. +template +__global__ void chunk_local_cumsum_kernel( + T const* __restrict__ g, + float* __restrict__ o, + float scale, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int T_len, + int H, + int S, + int BT +) { + int i_s = blockIdx.x * blockDim.x + threadIdx.x; + int64_t i_t = blockIdx.y; + int64_t i_bh = blockIdx.z; + int i_h = int(i_bh % H); + if (i_s >= S) return; + + int64_t bos, i_t0; + int64_t seq_len = T_len; + if (IS_VARLEN) { + int64_t i_n = chunk_indices[i_t * 2]; + int64_t i_tl = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + seq_len = cu_seqlens[i_n + 1] - bos; + i_t0 = i_tl * BT; + } else { + bos = (i_bh / H) * (int64_t)T_len; + i_t0 = i_t * BT; + } + + int64_t base = (bos * H + i_h) * (int64_t)S + i_s; + int64_t stride = (int64_t)H * S; + + int64_t rem = seq_len - i_t0; + int t_valid = int(rem < (int64_t)BT ? rem : (int64_t)BT); + float run = 0.0f; + if (REVERSE) { + for (int t = t_valid - 1; t >= 0; --t) { + int64_t off = base + (i_t0 + t) * stride; + run += to_f32(g[off]); + o[off] = HAS_SCALE ? run * scale : run; + } + } else { + for (int t = 0; t < t_valid; ++t) { + int64_t off = base + (i_t0 + t) * stride; + run += to_f32(g[off]); + o[off] = HAS_SCALE ? run * scale : run; + } + } +} + +constexpr int kGateBwdBT = 32; // matches fla's fixed BT=32 in kda_gate_bwd + +// Grid: (cdiv(B*T, 32), H). Block: 128 threads over the D dimension. +// Writes dg and the per-block partial dA; the caller sums dA_partial over dim 0. +template +__global__ void kda_gate_bwd_kernel( + T const* __restrict__ g, + float const* __restrict__ A_log, + float const* __restrict__ dt_bias, + float const* __restrict__ dyg, + float* __restrict__ dg, + float* __restrict__ dA_partial, + float lower_bound, + int64_t T_total, + int H, + int D +) { + int64_t i_t = blockIdx.x; + int i_h = blockIdx.y; + int64_t t0 = i_t * kGateBwdBT; + + float b_A = expf(A_log[i_h]); + float partial = 0.0f; + + int64_t rem = T_total - t0; + int t_valid = int(rem < (int64_t)kGateBwdBT ? rem : (int64_t)kGateBwdBT); + for (int d = threadIdx.x; d < D; d += blockDim.x) { + float bias = 0.0f; + if (HAS_BIAS) bias = dt_bias[i_h * D + d]; + for (int t = 0; t < t_valid; ++t) { + int64_t off = (t0 + t) * (int64_t)H * D + i_h * D + d; + float x = to_f32(g[off]) + bias; + float dy = dyg[off]; + float dgv; + if (USE_LOWER_BOUND) { + float sig = sigmoid_f32(b_A * x); + dgv = dy * lower_bound * sig * (1.0f - sig) * b_A; + partial += dgv * x; + } else { + float yg = -b_A * softplus_f32(x); + dgv = -b_A * dy * sigmoid_f32(x); + partial += dy * yg; + } + dg[off] = dgv; + } + } + + // block reduce partial + __shared__ float smem[32]; + for (int offset = 16; offset > 0; offset >>= 1) + partial += __shfl_down_sync(0xffffffff, partial, offset); + if ((threadIdx.x & 31) == 0) smem[threadIdx.x >> 5] = partial; + __syncthreads(); + if (threadIdx.x < 32) { + int n_warps = (int(blockDim.x) + 31) / 32; + float v = threadIdx.x < n_warps ? smem[threadIdx.x] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) + v += __shfl_down_sync(0xffffffff, v, offset); + if (threadIdx.x == 0) dA_partial[i_t * H + i_h] = v; + } +} + +struct VarlenArgs { + int64_t const* cu_seqlens = nullptr; + int64_t const* chunk_indices = nullptr; + int64_t NT = 0; +}; + +VarlenArgs resolve_varlen( + std::optional const& cu_seqlens, + std::optional const& chunk_indices, + int64_t B, + int64_t T, + int64_t chunk_size +) { + VarlenArgs args; + if (cu_seqlens.has_value()) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_indices.has_value(), "chunk_indices must be provided when cu_seqlens is provided"); + auto const& ci = chunk_indices.value(); + TORCH_CHECK(ci.dtype() == torch::kLong && ci.is_cuda() && ci.is_contiguous()); + TORCH_CHECK(ci.dim() == 2 && ci.size(1) == 2); + auto const& cu = cu_seqlens.value(); + TORCH_CHECK(cu.dtype() == torch::kLong && cu.is_cuda() && cu.is_contiguous()); + args.cu_seqlens = cu.data_ptr(); + args.chunk_indices = ci.data_ptr(); + args.NT = ci.size(0); + } else { + args.NT = (T + chunk_size - 1) / chunk_size; + } + return args; +} + +void check_gate_inputs( + torch::Tensor const& g, + torch::Tensor const& A_log, + std::optional const& dt_bias, + torch::Tensor const& out +) { + TORCH_CHECK(g.is_cuda() && g.is_contiguous(), "g must be contiguous CUDA tensor"); + TORCH_CHECK(g.dim() == 4, "g must be [B, T, H, S]"); + TORCH_CHECK(A_log.is_cuda() && A_log.is_contiguous() && A_log.dtype() == torch::kFloat32); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == g.size(2), "A_log must be [H]"); + if (dt_bias.has_value()) { + TORCH_CHECK(dt_bias->is_cuda() && dt_bias->is_contiguous() && dt_bias->dtype() == torch::kFloat32); + TORCH_CHECK(dt_bias->numel() == g.size(2) * g.size(3), "dt_bias must have H*S elements"); + } + TORCH_CHECK(out.is_cuda() && out.is_contiguous() && out.dtype() == torch::kFloat32); + TORCH_CHECK(out.sizes() == g.sizes(), "out must match g shape"); +} + +template +void launch_gate_cumsum( + T const* g_ptr, + float const* A_log_ptr, + float const* bias_ptr, + float* o_ptr, + float scale, + float lower_bound, + bool has_bias, + bool use_lower_bound, + bool has_scale, + VarlenArgs const& varlen, + int T_len, + int H, + int S, + int chunk_size, + dim3 grid, + cudaStream_t stream +) { + dim3 block(kBS); + + #define LAUNCH_GATE(HAS_BIAS, USE_LB, HAS_SCALE, IS_VARLEN) \ + kda_gate_chunk_cumsum_kernel<<>>( \ + g_ptr, A_log_ptr, bias_ptr, o_ptr, scale, lower_bound, \ + varlen.cu_seqlens, varlen.chunk_indices, T_len, H, S, chunk_size) + + #define DISPATCH_VARLEN(HAS_BIAS, USE_LB, HAS_SCALE) \ + if (varlen.cu_seqlens) { LAUNCH_GATE(HAS_BIAS, USE_LB, HAS_SCALE, true); } \ + else { LAUNCH_GATE(HAS_BIAS, USE_LB, HAS_SCALE, false); } + #define DISPATCH_SCALE(HAS_BIAS, USE_LB) \ + if (has_scale) { DISPATCH_VARLEN(HAS_BIAS, USE_LB, true); } \ + else { DISPATCH_VARLEN(HAS_BIAS, USE_LB, false); } + #define DISPATCH_LB(HAS_BIAS) \ + if (use_lower_bound) { DISPATCH_SCALE(HAS_BIAS, true); } \ + else { DISPATCH_SCALE(HAS_BIAS, false); } + + if (has_bias) { DISPATCH_LB(true); } + else { DISPATCH_LB(false); } + + #undef DISPATCH_LB + #undef DISPATCH_SCALE + #undef DISPATCH_VARLEN + #undef LAUNCH_GATE +} + +template +void launch_local_cumsum( + T const* g_ptr, + float* o_ptr, + float scale, + bool has_scale, + bool reverse, + VarlenArgs const& varlen, + int T_len, + int H, + int S, + int chunk_size, + dim3 grid, + cudaStream_t stream +) { + dim3 block(kBS); + + #define LAUNCH_CUMSUM(HAS_SCALE, REVERSE, IS_VARLEN) \ + chunk_local_cumsum_kernel<<>>( \ + g_ptr, o_ptr, scale, varlen.cu_seqlens, varlen.chunk_indices, T_len, H, S, chunk_size) + + #define DISPATCH_VARLEN_C(HAS_SCALE, REVERSE) \ + if (varlen.cu_seqlens) { LAUNCH_CUMSUM(HAS_SCALE, REVERSE, true); } \ + else { LAUNCH_CUMSUM(HAS_SCALE, REVERSE, false); } + #define DISPATCH_REVERSE(HAS_SCALE) \ + if (reverse) { DISPATCH_VARLEN_C(HAS_SCALE, true); } \ + else { DISPATCH_VARLEN_C(HAS_SCALE, false); } + + if (has_scale) { DISPATCH_REVERSE(true); } + else { DISPATCH_REVERSE(false); } + + #undef DISPATCH_REVERSE + #undef DISPATCH_VARLEN_C + #undef LAUNCH_CUMSUM +} + +template +void launch_gate_bwd( + T const* g_ptr, + float const* A_log_ptr, + float const* bias_ptr, + float const* dyg_ptr, + float* dg_ptr, + float* dA_ptr, + float lower_bound, + bool use_lower_bound, + int64_t T_total, + int H, + int D, + int64_t NT, + cudaStream_t stream +) { + dim3 grid(NT, H); + dim3 block(128); + + #define LAUNCH_GATE_BWD(HAS_BIAS, USE_LB) \ + kda_gate_bwd_kernel<<>>( \ + g_ptr, A_log_ptr, bias_ptr, dyg_ptr, dg_ptr, dA_ptr, lower_bound, T_total, H, D) + + if (bias_ptr) { + if (use_lower_bound) { LAUNCH_GATE_BWD(true, true); } else { LAUNCH_GATE_BWD(true, false); } + } else { + if (use_lower_bound) { LAUNCH_GATE_BWD(false, true); } else { LAUNCH_GATE_BWD(false, false); } + } + #undef LAUNCH_GATE_BWD +} + +// Map torch scalar type to the CUDA value type used by the kernels. +template +struct KernelType { + using type = std::conditional_t, cutlass::bfloat16_t, + std::conditional_t, cutlass::half_t, float>>; +}; + +} // namespace + +void kda_gate_chunk_cumsum( + torch::Tensor g, + torch::Tensor A_log, + std::optional dt_bias, + torch::Tensor out, + double scale, + bool has_scale, + double lower_bound, + bool use_lower_bound, + int64_t chunk_size, + std::optional cu_seqlens, + std::optional chunk_indices +) { + check_gate_inputs(g, A_log, dt_bias, out); + int64_t B = g.size(0), T_len = g.size(1), H = g.size(2), S = g.size(3); + auto varlen = resolve_varlen(cu_seqlens, chunk_indices, B, T_len, chunk_size); + + dim3 grid((S + kBS - 1) / kBS, varlen.NT, B * H); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kBFloat16, at::kHalf, g.scalar_type(), "kda_gate_chunk_cumsum", [&] { + using T = typename KernelType::type; + launch_gate_cumsum( + reinterpret_cast(g.data_ptr()), A_log.data_ptr(), + dt_bias.has_value() ? dt_bias->data_ptr() : nullptr, + out.data_ptr(), + float(scale), float(lower_bound), + dt_bias.has_value(), use_lower_bound, has_scale, + varlen, int(T_len), int(H), int(S), int(chunk_size), grid, stream + ); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void chunk_local_cumsum( + torch::Tensor g, + torch::Tensor out, + double scale, + bool has_scale, + bool reverse, + int64_t chunk_size, + std::optional cu_seqlens, + std::optional chunk_indices +) { + TORCH_CHECK(g.is_cuda() && g.is_contiguous() && g.dim() == 4, "g must be 4D contiguous CUDA tensor"); + TORCH_CHECK(out.is_cuda() && out.is_contiguous() && out.dtype() == torch::kFloat32); + TORCH_CHECK(out.sizes() == g.sizes(), "out must match g shape"); + int64_t B = g.size(0), T_len = g.size(1), H = g.size(2), S = g.size(3); + auto varlen = resolve_varlen(cu_seqlens, chunk_indices, B, T_len, chunk_size); + + dim3 grid((S + kBS - 1) / kBS, varlen.NT, B * H); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kBFloat16, at::kHalf, g.scalar_type(), "chunk_local_cumsum", [&] { + using T = typename KernelType::type; + launch_local_cumsum( + reinterpret_cast(g.data_ptr()), out.data_ptr(), + float(scale), has_scale, reverse, + varlen, int(T_len), int(H), int(S), int(chunk_size), grid, stream + ); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void kda_gate_bwd( + torch::Tensor g, + torch::Tensor A_log, + std::optional dt_bias, + torch::Tensor dyg, + torch::Tensor dg, + torch::Tensor dA_partial, + double lower_bound, + bool use_lower_bound +) { + TORCH_CHECK(g.is_cuda() && g.is_contiguous() && g.dim() == 4, "g must be 4D contiguous CUDA tensor"); + TORCH_CHECK(A_log.is_cuda() && A_log.is_contiguous() && A_log.dtype() == torch::kFloat32); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == g.size(2), "A_log must be [H]"); + if (dt_bias.has_value()) { + TORCH_CHECK(dt_bias->is_cuda() && dt_bias->is_contiguous() && dt_bias->dtype() == torch::kFloat32); + TORCH_CHECK(dt_bias->numel() == g.size(2) * g.size(3), "dt_bias must have H*D elements"); + } + TORCH_CHECK(dyg.is_cuda() && dyg.is_contiguous() && dyg.dtype() == torch::kFloat32); + TORCH_CHECK(dg.is_cuda() && dg.is_contiguous() && dg.dtype() == torch::kFloat32); + TORCH_CHECK(dyg.sizes() == g.sizes() && dg.sizes() == g.sizes()); + + int64_t B = g.size(0), T_len = g.size(1), H = g.size(2), D = g.size(3); + int64_t T_total = B * T_len; + int64_t NT = (T_total + kGateBwdBT - 1) / kGateBwdBT; + TORCH_CHECK(dA_partial.dim() == 2 && dA_partial.size(0) == NT && dA_partial.size(1) == H, + "dA_partial must be [cdiv(B*T, 32), H]"); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kBFloat16, at::kHalf, g.scalar_type(), "kda_gate_bwd", [&] { + using T = typename KernelType::type; + launch_gate_bwd( + reinterpret_cast(g.data_ptr()), A_log.data_ptr(), + dt_bias.has_value() ? dt_bias->data_ptr() : nullptr, + dyg.data_ptr(), dg.data_ptr(), dA_partial.data_ptr(), + float(lower_bound), use_lower_bound, T_total, int(H), int(D), NT, stream + ); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/csrc/train/intra.cu b/csrc/train/intra.cu new file mode 100644 index 0000000..c25400b --- /dev/null +++ b/csrc/train/intra.cu @@ -0,0 +1,1193 @@ +// KDA forward intra-chunk kernels, replicating fla/ops/kda/chunk_intra.py +// (chunk_kda_fwd_kernel_intra_sub_chunk, chunk_kda_fwd_kernel_inter_solve_fused) +// and fla/ops/kda/chunk_intra_token_parallel.py. +// +// All gate math is fp32 in the log2 domain (g is the chunk-local inclusive +// cumsum scaled by RCP_LN2). 16x16 GEMMs run on tensor cores through the CuTe +// SM80 tf32 atom (SM80_16x8x8_F32TF32TF32F32_TN) to match Triton's tf32 dots; +// the token_parallel kernel replicates Triton's fp32 elementwise + tl.sum +// semantics and therefore intentionally uses no MMA. + +#include +#include + +#include +#include + +#include "common.cuh" + +// Named namespace (not anonymous): nvcc's registration stub cannot spell +// template instantiations from anonymous namespaces (ambiguous mangled-name +// references in the generated stub). +namespace kda_train_intra { + +using namespace cute; + +using BF16 = cutlass::bfloat16_t; +using FP16 = cutlass::half_t; +using TF32 = cutlass::tfloat32_t; + +constexpr int kBC = 16; // sub-chunk size (fla BC) +constexpr int kKC2 = 32; // K staging chunk; keeps the ascending-K mma order + +// tl.math.exp2 lowers to ex2.approx; keep the same instruction. +__device__ __forceinline__ float exp2_ftz(float x) { + float result; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x)); + return result; +} + +// fp32 -> tf32 with the same rounding Triton applies before tf32 mma. +__device__ __forceinline__ TF32 f32_to_tf32(float x) { + uint32_t r; + asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(r) : "f"(x)); + return TF32::bitcast(r); +} + +__device__ __forceinline__ uint4 pack4_tf32(float const* v) { + return make_uint4(f32_to_tf32(v[0]).storage, f32_to_tf32(v[1]).storage, + f32_to_tf32(v[2]).storage, f32_to_tf32(v[3]).storage); +} + +// One warp computing C[16,16] (fp32) += A[16,Kc] @ B[16,Kc]^T with tf32 +// operands in shared memory (row-major, K contiguous). The tiled MMA is the +// 16x8x8 tf32 atom value-tiled to 16x16x8. +using MmaTF32_16 = decltype(make_tiled_mma( + SM80_16x8x8_F32TF32TF32F32_TN{}, + Layout>{}, + Tile<_16, _16, _8>{} +)); + +template +struct Mma16Ctx { + Mma mma; + ThrMma thr; + Acc acc; + Coord coord; // identity-tensor partition mapping acc(v) -> (i, j) +}; + +template +__device__ __forceinline__ auto make_mma16_ctx(int lane) { + Mma mma{}; + auto thr = mma.get_slice(lane); + auto coord = thr.partition_C(make_identity_tensor(Shape<_16, _16>{})); + auto acc = thr.make_fragment_C(coord); + clear(acc); + return Mma16Ctx{mma, thr, acc, coord}; +} + +// Row strides are padded (+4 floats) so the tf32 fragment loads of lanes +// grouped 4-apart (rows lane/4) land on distinct bank quads: stride%32 == 4 +// keeps rows r and r+1 eight banks apart, making the scalar ld.shared in +// mma16_accum / merge_gemm_16 conflict-free. +constexpr int kPadHalf = 32 + 4; +constexpr int kPad16 = 16 + 4; +using SmemLayoutHalf = Layout, Stride, _1>>; // [16, 32] staging tiles +using SmemLayout16 = Layout, Stride, _1>>; // [16, 16] merge tiles + +template +__device__ __forceinline__ void mma16_accum( + Ctx& ctx, + TF32 const* sA, // [16, K] row-major + TF32 const* sB, // [16, K] row-major; computes A @ B^T + SmemLayout const& lay +) { + Tensor sAt = make_tensor(make_smem_ptr(sA), lay); + Tensor sBt = make_tensor(make_smem_ptr(sB), lay); + Tensor tCsA = ctx.thr.partition_A(sAt); + Tensor tCsB = ctx.thr.partition_B(sBt); + Tensor tCrA = ctx.thr.partition_fragment_A(sAt); + Tensor tCrB = ctx.thr.partition_fragment_B(sBt); + copy(tCsA, tCrA); + copy(tCsB, tCrB); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrA); ++kb) { + gemm(ctx.mma, tCrA(_, _, kb), tCrB(_, _, kb), ctx.acc); + } +} + +// In-place forward substitution on a 16x16 fp32 smem block holding -tril(C,-1): +// turns it into (I + tril(C,-1))^-1 minus the identity (the caller adds I). +// Equivalent to the gmem read-back loop in the Triton kernels: the row read +// there is exactly the current (negated) row of the block. lanes 0..15 only; +// each lane owns one column. +__device__ __forceinline__ void fwd_subst_16(float* sAi, int lim, int lane) { + // only lanes 0..15 participate; sync with an explicit partial mask + float a[kBC]; + for (int i = 2; i < lim; ++i) { + CUTE_UNROLL + for (int r = 0; r < kBC; ++r) a[r] = (r < i) ? sAi[i * kBC + r] : 0.f; + float acc = a[lane]; + CUTE_UNROLL + for (int r = 0; r < kBC; ++r) acc += a[r] * sAi[r * kBC + lane]; + __syncwarp(0xffffu); + sAi[i * kBC + lane] = acc; + __syncwarp(0xffffu); + } +} + +struct VarlenArgs { + int64_t const* cu_seqlens = nullptr; + int64_t const* chunk_indices = nullptr; + int64_t NT = 0; +}; + +VarlenArgs resolve_varlen_intra( + std::optional const& cu_seqlens, + std::optional const& chunk_indices, + int64_t B, + int64_t T, + int64_t chunk_size +) { + VarlenArgs args; + if (cu_seqlens.has_value()) { + TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); + TORCH_CHECK(chunk_indices.has_value(), "chunk_indices must be provided when cu_seqlens is provided"); + auto const& ci = chunk_indices.value(); + TORCH_CHECK(ci.dtype() == torch::kLong && ci.is_cuda() && ci.is_contiguous()); + TORCH_CHECK(ci.dim() == 2 && ci.size(1) == 2); + auto const& cu = cu_seqlens.value(); + TORCH_CHECK(cu.dtype() == torch::kLong && cu.is_cuda() && cu.is_contiguous()); + args.cu_seqlens = cu.data_ptr(); + args.chunk_indices = ci.data_ptr(); + args.NT = ci.size(0); + } else { + args.NT = (T + chunk_size - 1) / chunk_size; + } + return args; +} + +void check_intra_inputs( + torch::Tensor const& q, + torch::Tensor const& k, + torch::Tensor const& g, + torch::Tensor const& beta, + torch::Tensor const& Aqk, + torch::Tensor const& Akkd, + int64_t chunk_size +) { + TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA tensor"); + TORCH_CHECK(k.is_cuda() && k.is_contiguous(), "k must be contiguous CUDA tensor"); + TORCH_CHECK(g.is_cuda() && g.is_contiguous(), "g must be contiguous CUDA tensor"); + TORCH_CHECK(beta.is_cuda() && beta.is_contiguous(), "beta must be contiguous CUDA tensor"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "q/k must be [B, T, H, K]"); + TORCH_CHECK(q.sizes() == k.sizes() && q.scalar_type() == k.scalar_type(), "q/k must match"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16 || q.scalar_type() == at::kHalf, "q/k must be bf16/fp16"); + int64_t B = k.size(0), T = k.size(1), H = k.size(2), K = k.size(3); + TORCH_CHECK(K <= 256, "K must be <= 256"); + TORCH_CHECK(chunk_size == 32 || chunk_size == 64, "chunk_size must be 32 or 64"); + TORCH_CHECK(g.dim() == 4 && g.scalar_type() == at::kFloat, "g must be fp32 [B, T, HV, K]"); + TORCH_CHECK(g.size(0) == B && g.size(1) == T && g.size(3) == K, "g shape mismatch"); + int64_t HV = g.size(2); + TORCH_CHECK(HV % H == 0, "HV must be a multiple of H"); + TORCH_CHECK(beta.dim() == 3 && beta.size(0) == B && beta.size(1) == T && beta.size(2) == HV, + "beta must be [B, T, HV]"); + TORCH_CHECK(beta.scalar_type() == at::kBFloat16 || beta.scalar_type() == at::kHalf || + beta.scalar_type() == at::kFloat, "beta must be bf16/fp16/fp32"); + TORCH_CHECK(Aqk.is_cuda() && Aqk.is_contiguous() && Aqk.scalar_type() == k.scalar_type()); + TORCH_CHECK(Aqk.dim() == 4 && Aqk.size(0) == B && Aqk.size(1) == T && Aqk.size(2) == HV && + Aqk.size(3) == chunk_size, "Aqk must be [B, T, HV, BT]"); + TORCH_CHECK(Akkd.is_cuda() && Akkd.is_contiguous() && Akkd.scalar_type() == at::kFloat); + TORCH_CHECK(Akkd.dim() == 4 && Akkd.size(0) == B && Akkd.size(1) == T && Akkd.size(2) == HV && + Akkd.size(3) == kBC, "Akkd must be [B, T, HV, 16] fp32"); + TORCH_CHECK(B * HV <= 65535, "B*HV exceeds gridDim.z limit"); +} + +// --------------------------------------------------------------------------- +// Kernel 1: chunk_kda_fwd_kernel_intra_sub_chunk (safe_gate path) +// grid (NT, B*HV), block 128. Each warp computes one 16x16 diagonal block of +// the chunk (warp w <-> sub-chunk w): the Aqk diag block (scaled, bf16) and +// the inverted Akk diag block (fp32 into Akkd). Warps are fully independent +// (private smem scratch, warp-local mma and forward substitution, no +// block-wide barriers), so the four diag blocks of a chunk are computed +// concurrently instead of serialized on warp 0 as in v1. Per-element math +// and the ascending-K mma order are unchanged. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(128) chunk_kda_fwd_kernel_intra_sub_chunk_cuda( + QKT const* __restrict__ q, + QKT const* __restrict__ k, + float const* __restrict__ g, + BetaT const* __restrict__ beta, + QKT* __restrict__ Aqk, + float* __restrict__ Akkd, + float scale, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int T, + int H, + int HV, + int K +) { + constexpr int NC = BT / kBC; + int64_t i_t = blockIdx.x; + int64_t i_bh = blockIdx.y; + int i_hv = int(i_bh % HV); + int i_h = i_hv / (HV / H); + + int64_t bos; + int Tseq = T; + if (IS_VARLEN) { + int64_t i_n = chunk_indices[i_t * 2]; + i_t = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + Tseq = int(cu_seqlens[i_n + 1] - bos); + } else { + bos = (i_bh / HV) * (int64_t)T; + } + if (i_t * BT >= Tseq) return; + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const i_i = warp; // sub-chunk handled by this warp + if (i_i >= NC) return; // no block-wide barriers in this kernel + int const i_ti = int(i_t) * BT + i_i * kBC; + if (i_ti >= Tseq) return; + + __shared__ TF32 sW[4][3][kBC * kPadHalf]; // per-warp Aq/Ak/B staging tiles + __shared__ float sAi[4][kBC * kBC]; + __shared__ float sBeta[NC][kBC]; + + TF32* sAq = &sW[warp][0][0]; + TF32* sAk = &sW[warp][1][0]; + TF32* sB = &sW[warp][2][0]; + float* sAiW = &sAi[warp][0]; + + if (lane < kBC) { + int t = i_ti + lane; + sBeta[warp][lane] = (t < Tseq) ? to_f32(beta[(bos + t) * HV + i_hv]) : 0.f; + } + // midpoint reference row of the sub-chunk (numerical stability) + int const gn_row = i_ti + min(kBC / 2, Tseq - i_ti - 1); + + auto ctx_qk = make_mma16_ctx(lane); + auto ctx_kk = make_mma16_ctx(lane); + + int const n_kchunks = (K + kKC2 - 1) / kKC2; + bool const kvec = (K % 8) == 0; + // lane covers col group cg = (lane%4)*8 and rows lane/4, lane/4+8 + int const cg = (lane & 3) * 8; + int const r0 = lane >> 2; + float const* gn_base = g + ((bos + gn_row) * HV + i_hv) * (int64_t)K; + if (kvec) { + // Pass-level pipeline: the next 8-col pass's raw gmem payload is + // issued before the current pass is gated/packed/stored, hiding the + // gmem latency behind the exp2/pack math and the fragment MMAs. + // Raw payload of one pass: reference-row gates + row gates + q/k. + struct RawPass { + float4 gn0, gn1, gv0, gv1; + uint4 rq, rk; + unsigned flags; // bit0: row valid, bit1: full 8-col group + }; + auto load_pass = [&](int kc, int pass, RawPass& R) { + int const kk = kc * kKC2 + cg; + int const r = r0 + pass * 8; + int const t = i_ti + r; + bool const ok = t < Tseq; + bool const v8 = (K - kk) >= 8; // K%8==0: group is full or empty + R.flags = (ok ? 1u : 0u) | (v8 ? 2u : 0u); + if (v8) { + float4 const* pg = reinterpret_cast(gn_base + kk); + R.gn0 = pg[0]; R.gn1 = pg[1]; + if (ok) { + float4 const* pr = reinterpret_cast(g + ((bos + t) * HV + i_hv) * (int64_t)K + kk); + R.gv0 = pr[0]; R.gv1 = pr[1]; + R.rq = *reinterpret_cast(q + ((bos + t) * H + i_h) * (int64_t)K + kk); + R.rk = *reinterpret_cast(k + ((bos + t) * H + i_h) * (int64_t)K + kk); + } + } + }; + auto store_pass = [&](RawPass const& R, int pass) { + int const r = r0 + pass * 8; + bool const ok = (R.flags & 1u) != 0; + bool const v8 = (R.flags & 2u) != 0; + float gnv[8] = {}, gv[8] = {}, qv[8] = {}, kv[8] = {}; + if (v8) { + gnv[0] = R.gn0.x; gnv[1] = R.gn0.y; gnv[2] = R.gn0.z; gnv[3] = R.gn0.w; + gnv[4] = R.gn1.x; gnv[5] = R.gn1.y; gnv[6] = R.gn1.z; gnv[7] = R.gn1.w; + if (ok) { + gv[0] = R.gv0.x; gv[1] = R.gv0.y; gv[2] = R.gv0.z; gv[3] = R.gv0.w; + gv[4] = R.gv1.x; gv[5] = R.gv1.y; gv[6] = R.gv1.z; gv[7] = R.gv1.w; + QKT const* hq = reinterpret_cast(&R.rq); + QKT const* hk = reinterpret_cast(&R.rk); + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { qv[j] = to_f32(hq[j]); kv[j] = to_f32(hk[j]); } + } + } + float oq[8], okk2[8], ob[8]; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { + float gq = ok ? exp2_ftz(gv[j] - gnv[j]) : 0.f; + float gk = ok ? exp2_ftz(gnv[j] - gv[j]) : 0.f; + oq[j] = qv[j] * gq; + okk2[j] = kv[j] * gq; + ob[j] = kv[j] * gk; + } + int const soff = r * kPadHalf + cg; + *reinterpret_cast(sAq + soff) = pack4_tf32(oq); + *reinterpret_cast(sAq + soff + 4) = pack4_tf32(oq + 4); + *reinterpret_cast(sAk + soff) = pack4_tf32(okk2); + *reinterpret_cast(sAk + soff + 4) = pack4_tf32(okk2 + 4); + *reinterpret_cast(sB + soff) = pack4_tf32(ob); + *reinterpret_cast(sB + soff + 4) = pack4_tf32(ob + 4); + }; + RawPass cur, nxt; + load_pass(0, 0, cur); + for (int kc = 0; kc < n_kchunks; ++kc) { + CUTE_UNROLL + for (int pass = 0; pass < 2; ++pass) { + int const nkc = kc + pass; + if (nkc < n_kchunks) load_pass(nkc, (pass + 1) & 1, nxt); + store_pass(cur, pass); + if (pass == 1) { + __syncwarp(); + mma16_accum(ctx_qk, sAq, sB, SmemLayoutHalf{}); + mma16_accum(ctx_kk, sAk, sB, SmemLayoutHalf{}); + __syncwarp(); + } + cur = nxt; + } + } + } else { + for (int kc = 0; kc < n_kchunks; ++kc) { + int const kk = kc * kKC2 + cg; + int const rem = K - kk; + CUTE_UNROLL + for (int pass = 0; pass < 2; ++pass) { + int const r = r0 + pass * 8; + int const t = i_ti + r; + bool const ok = t < Tseq; + float gnv[8] = {}, gv[8] = {}, qv[8] = {}, kv[8] = {}; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { + bool okk = j < rem; + gnv[j] = okk ? gn_base[kk + j] : 0.f; + if (ok && okk) { + gv[j] = g[((bos + t) * HV + i_hv) * (int64_t)K + kk + j]; + qv[j] = to_f32(q[((bos + t) * H + i_h) * (int64_t)K + kk + j]); + kv[j] = to_f32(k[((bos + t) * H + i_h) * (int64_t)K + kk + j]); + } + } + float oq[8], okk2[8], ob[8]; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { + float gq = (ok && j < rem) ? exp2_ftz(gv[j] - gnv[j]) : 0.f; + float gk = (ok && j < rem) ? exp2_ftz(gnv[j] - gv[j]) : 0.f; + oq[j] = qv[j] * gq; + okk2[j] = kv[j] * gq; + ob[j] = kv[j] * gk; + } + int const soff = r * kPadHalf + cg; + *reinterpret_cast(sAq + soff) = pack4_tf32(oq); + *reinterpret_cast(sAq + soff + 4) = pack4_tf32(oq + 4); + *reinterpret_cast(sAk + soff) = pack4_tf32(okk2); + *reinterpret_cast(sAk + soff + 4) = pack4_tf32(okk2 + 4); + *reinterpret_cast(sB + soff) = pack4_tf32(ob); + *reinterpret_cast(sB + soff + 4) = pack4_tf32(ob + 4); + } + __syncwarp(); + mma16_accum(ctx_qk, sAq, sB, SmemLayoutHalf{}); + mma16_accum(ctx_kk, sAk, sB, SmemLayoutHalf{}); + __syncwarp(); + } + } + + // Aqk diagonal block: lower triangle incl. diagonal, scale applied + CUTE_UNROLL + for (int v = 0; v < size(ctx_qk.acc); ++v) { + auto c = ctx_qk.coord(v); + int i = int(get<0>(c)), j = int(get<1>(c)); + if (i_ti + i < Tseq) { + float val = (i >= j) ? ctx_qk.acc(v) * scale : 0.f; + Aqk[((bos + i_ti + i) * HV + i_hv) * (int64_t)BT + i_i * kBC + j] = QKT(val); + } + } + // sAi = -tril(beta * Akk, -1) + CUTE_UNROLL + for (int v = 0; v < size(ctx_kk.acc); ++v) { + auto c = ctx_kk.coord(v); + int i = int(get<0>(c)), j = int(get<1>(c)); + sAiW[i * kBC + j] = (i > j) ? -ctx_kk.acc(v) * sBeta[warp][i] : 0.f; + } + __syncwarp(); + + if (lane < kBC) { + fwd_subst_16(sAiW, min(kBC, Tseq - i_ti), lane); + sAiW[lane * kBC + lane] += 1.f; + } + __syncwarp(); + + for (int idx = lane; idx < kBC * kBC; idx += 32) { + int i = idx / kBC; + if (i_ti + i < Tseq) { + Akkd[((bos + i_ti + i) * HV + i_hv) * (int64_t)kBC + idx % kBC] = sAiW[idx]; + } + } +} + +// --------------------------------------------------------------------------- +// Kernel 2: chunk_kda_fwd_kernel_intra_token_parallel (non safe_gate path) +// grid (B*T, cdiv(HV, 4)), block 128 (one warp per value head). Replicates +// Triton's fp32 elementwise + tl.sum semantics (no MMA by design). +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(128) chunk_kda_fwd_kernel_intra_token_parallel_cuda( + QKT const* __restrict__ q, + QKT const* __restrict__ k, + float const* __restrict__ g, + BetaT const* __restrict__ beta, + QKT* __restrict__ Aqk, + float* __restrict__ Akkd, + float scale, + int64_t const* __restrict__ cu_seqlens, + int64_t N, + int T, + int H, + int HV, + int K +) { + int64_t i_tg = blockIdx.x; + + int64_t bos; + int64_t i_t; + int Tseq; + if (IS_VARLEN) { + // unrolled binary search for the sequence containing token i_tg + int left = 0, right = int(N); + CUTE_UNROLL + for (int it = 0; it < 20; ++it) { + if (left < right) { + int mid = (left + right) >> 1; + if (i_tg < cu_seqlens[mid + 1]) right = mid; + else left = mid + 1; + } + } + bos = cu_seqlens[left]; + Tseq = int(cu_seqlens[left + 1] - bos); + i_t = i_tg - bos; + } else { + bos = (i_tg / T) * T; + i_t = i_tg % T; + Tseq = T; + } + if (i_t >= Tseq) return; + + int warp = threadIdx.x >> 5; + int lane = threadIdx.x & 31; + int i_hv = blockIdx.y * 4 + warp; + if (i_hv >= HV) return; + int i_h = i_hv / (HV / H); + + int64_t i_c = i_t / BT; + int i_s = int((i_t % BT) / kBC); + int64_t i_ts = i_c * BT + i_s * kBC; + + constexpr int MAXE = 8; // K <= 256, 32 lanes * 8 + float q_r[MAXE], kb_r[MAXE], g_r[MAXE]; + int64_t qk_base = ((bos + i_t) * H + i_h) * (int64_t)K; + int64_t g_base = ((bos + i_t) * HV + i_hv) * (int64_t)K; + float beta_i = to_f32(beta[(bos + i_t) * HV + i_hv]); + CUTE_UNROLL + for (int e = 0; e < MAXE; ++e) { + int kk = lane + e * 32; + bool ok = kk < K; + q_r[e] = ok ? to_f32(q[qk_base + kk]) : 0.f; + kb_r[e] = ok ? to_f32(k[qk_base + kk]) * beta_i : 0.f; + g_r[e] = ok ? g[g_base + kk] : 0.f; + } + + int64_t j_end = i_t + 1; + if ((int64_t)Tseq < j_end) j_end = Tseq; + if (i_ts + kBC < j_end) j_end = i_ts + kBC; + for (int64_t j = i_ts; j < j_end; ++j) { + int64_t kjb = ((bos + j) * H + i_h) * (int64_t)K; + int64_t gjb = ((bos + j) * HV + i_hv) * (int64_t)K; + float aq = 0.f, ak = 0.f; + CUTE_UNROLL + for (int e = 0; e < MAXE; ++e) { + int kk = lane + e * 32; + if (kk < K) { + float kgj = to_f32(k[kjb + kk]) * exp2_ftz(g_r[e] - g[gjb + kk]); + aq += q_r[e] * kgj; + ak += kb_r[e] * kgj; + } + } + CUTE_UNROLL + for (int off = 16; off > 0; off >>= 1) { + aq += __shfl_down_sync(0xffffffffu, aq, off); + ak += __shfl_down_sync(0xffffffffu, ak, off); + } + if (lane == 0) { + Aqk[((bos + i_t) * HV + i_hv) * (int64_t)BT + (j % BT)] = QKT(aq * scale); + Akkd[((bos + i_t) * HV + i_hv) * (int64_t)kBC + (j - i_ts)] = (j < i_t) ? ak : 0.f; + } + } +} + +// --------------------------------------------------------------------------- +// Kernel 3: chunk_kda_fwd_kernel_inter_solve_fused +// grid (NT, B*HV), block 128. Computes off-diagonal Aqk/Akk blocks, inverts +// the diagonal blocks (non safe_gate only; safe_gate blocks arrive inverted +// from kernel 1), merges the block-triangular inverse, and writes Akk (bf16). +// --------------------------------------------------------------------------- + +// Warp-local 16x16 GEMM: D = alpha * (A @ B) with fp32 smem operands via tf32 +// MMA. Same staging math, atom, and K order as v1's block-cooperative version; +// sMA/sMB are this warp's private tf32 scratch. Contains __syncwarp only. +__device__ __forceinline__ void merge_gemm_16_warp( + float const* A, + float const* B, + float* D, + float alpha, + TF32* sMA, + TF32* sMB, + int lane +) { + CUTE_UNROLL + for (int idx = lane; idx < kBC * kBC; idx += 32) { + int n = idx / kBC, kk = idx % kBC; + sMA[n * kPad16 + kk] = f32_to_tf32(A[idx]); + sMB[n * kPad16 + kk] = f32_to_tf32(B[kk * kBC + n]); // sMB[n][kk] = B[kk][n]: mma16_accum computes A @ sMB^T = A @ B + } + __syncwarp(); + auto ctx = make_mma16_ctx(lane); + mma16_accum(ctx, sMA, sMB, SmemLayout16{}); + CUTE_UNROLL + for (int v = 0; v < size(ctx.acc); ++v) { + auto c = ctx.coord(v); + D[int(get<0>(c)) * kBC + int(get<1>(c))] = alpha * ctx.acc(v); + } + __syncwarp(); +} + +template +__global__ void __launch_bounds__(128) chunk_kda_fwd_kernel_inter_solve_fused_cuda( + QKT const* __restrict__ q, + QKT const* __restrict__ k, + float const* __restrict__ g, + BetaT const* __restrict__ beta, + QKT* __restrict__ Aqk, + float const* __restrict__ Akkd, + QKT* __restrict__ Akk, + float scale, + int64_t const* __restrict__ cu_seqlens, + int64_t const* __restrict__ chunk_indices, + int T, + int H, + int HV, + int K +) { + constexpr int NC = BT / kBC; + constexpr int NP = (NC == 4) ? 6 : 1; // off-diagonal block pairs + int64_t i_t = blockIdx.x; + int64_t i_bh = blockIdx.y; + int i_hv = int(i_bh % HV); + int i_h = i_hv / (HV / H); + + int64_t bos; + int Tseq = T; + if (IS_VARLEN) { + int64_t i_n = chunk_indices[i_t * 2]; + i_t = chunk_indices[i_t * 2 + 1]; + bos = cu_seqlens[i_n]; + Tseq = int(cu_seqlens[i_n + 1] - bos); + } else { + bos = (i_bh / HV) * (int64_t)T; + } + if (i_t * BT >= Tseq) return; + + // Phase 1+2 warp scratch: per-warp Aq/Ak/B staging tiles [kBC][kKC2]. + // (Replaces the old block-wide sAq/sAk/sB staging; see below.) + __shared__ TF32 sW[4][3][kBC * kPadHalf]; + __shared__ float sAkkOff[6][kBC * kBC]; // off-diagonal Akk blocks (beta applied) + __shared__ float sAi[NC][kBC * kBC]; // diagonal inverse blocks + __shared__ float sAiX[6][kBC * kBC]; // merged off-diagonal inverse blocks + __shared__ float sBeta[NC][kBC]; + + // Phase-5 scratch overlays sW (dead once phase 1+2 finishes): per-warp + // tf32 operand staging plus the fp32 partial products of the parallel + // block-triangular solve. + struct P5Scratch { + TF32 sMA[4][kBC * kPad16]; + TF32 sMB[4][kBC * kPad16]; + float T[3][kBC * kBC]; // stage-A products, later t20/t31/t30 partial sums + float P[6][kBC * kBC]; // independent products of the off-diagonal merge + }; + static_assert(sizeof(P5Scratch) <= sizeof(TF32) * 4 * 3 * kBC * kPadHalf); + + int tid = threadIdx.x; + int lane = tid & 31; + int warp = tid >> 5; + + int i_tc[NC]; + CUTE_UNROLL + for (int c = 0; c < NC; ++c) i_tc[c] = int(i_t) * BT + c * kBC; + + for (int idx = tid; idx < NP * kBC * kBC; idx += blockDim.x) { + sAkkOff[idx / (kBC * kBC)][idx % (kBC * kBC)] = 0.f; + } + CUTE_UNROLL + for (int c = 0; c < NC; ++c) { + if (tid < kBC) { + int t = i_tc[c] + tid; + sBeta[c][tid] = (t < Tseq) ? to_f32(beta[(bos + t) * HV + i_hv]) : 0.f; + } + } + // Diagonal blocks from Akkd (fp32), hoisted ahead of phase 1+2 so the gmem + // latency overlaps the off-diagonal staging/MMA work. + for (int idx = tid; idx < NC * kBC * kBC; idx += blockDim.x) { + int c = idx / (kBC * kBC); + int i = (idx / kBC) % kBC; + int j = idx % kBC; + sAi[c][i * kBC + j] = + (i_tc[c] + i < Tseq) ? Akkd[((bos + i_tc[c] + i) * HV + i_hv) * (int64_t)kBC + j] : 0.f; + } + __syncthreads(); + + // Phase 1+2: off-diagonal blocks. Pair p covers block row c, block col cp. + // Each warp owns whole pairs (p = warp, warp+4) and works independently + // through its private smem scratch, so stagings/MMAs of different pairs + // overlap and no block-wide barrier sits inside the loop. The K loop of a + // pair runs as a pass-level pipeline: the next 8-col pass's raw gmem + // payload is issued before the current pass is gated/packed/stored, so + // the gmem latency hides behind the exp2/pack math and the fragment MMAs. + // Elementwise fallback when K % 8 != 0. + // Raw gmem payload of one staging pass (one 8-col group of one row). + struct RawPass { + float4 gn0, gn1; // reference-row gates + float4 gv0, gv1; // row-block gates + float4 gp0, gp1; // col-block gates + uint4 rq, rk; // row-block q/k + uint4 rp; // col-block k + unsigned flags; // bit0: row valid, bit1: col valid, bit2: full 8-col group + }; + int const n_kchunks = (K + kKC2 - 1) / kKC2; + bool const kvec = (K % 8) == 0; + // lane covers col group cg = (lane%4)*8 and rows lane/4, lane/4+8 + int const cg = (lane & 3) * 8; + int const r0 = lane >> 2; + for (int p = warp; p < NP; p += 4) { + int const c = (p == 0) ? 1 : (p < 3 ? 2 : 3); + int const cp = p - c * (c - 1) / 2; + if (i_tc[c] >= Tseq) continue; + int gn_row = i_tc[c]; // reference: first row of the row block + TF32* sAq = &sW[warp][0][0]; + TF32* sAk = &sW[warp][1][0]; + TF32* sB = &sW[warp][2][0]; + auto ctx_qk = make_mma16_ctx(lane); + auto ctx_kk = make_mma16_ctx(lane); + float const* gn_base = g + ((bos + gn_row) * HV + i_hv) * (int64_t)K; + if (kvec) { + auto load_pass = [&](int kc, int pass, RawPass& R) { + int const kk = kc * kKC2 + cg; + int const r = r0 + pass * 8; + int const t = i_tc[c] + r; // row block token + int const tp = i_tc[cp] + r; // col block token + bool const ok_r = t < Tseq; + bool const ok_c = tp < Tseq; + bool const v8 = (K - kk) >= 8; // K%8==0: group is full or empty + R.flags = (ok_r ? 1u : 0u) | (ok_c ? 2u : 0u) | (v8 ? 4u : 0u); + if (v8) { + float4 const* pg = reinterpret_cast(gn_base + kk); + R.gn0 = pg[0]; R.gn1 = pg[1]; + if (ok_r) { + float4 const* pr = reinterpret_cast(g + ((bos + t) * HV + i_hv) * (int64_t)K + kk); + R.gv0 = pr[0]; R.gv1 = pr[1]; + R.rq = *reinterpret_cast(q + ((bos + t) * H + i_h) * (int64_t)K + kk); + R.rk = *reinterpret_cast(k + ((bos + t) * H + i_h) * (int64_t)K + kk); + } + if (ok_c) { + float4 const* pp = reinterpret_cast(g + ((bos + tp) * HV + i_hv) * (int64_t)K + kk); + R.gp0 = pp[0]; R.gp1 = pp[1]; + R.rp = *reinterpret_cast(k + ((bos + tp) * H + i_h) * (int64_t)K + kk); + } + } + }; + auto store_pass = [&](RawPass const& R, int pass) { + int const r = r0 + pass * 8; + bool const ok_r = (R.flags & 1u) != 0; + bool const ok_c = (R.flags & 2u) != 0; + bool const v8 = (R.flags & 4u) != 0; + float gnv[8] = {}, gv[8] = {}, qv[8] = {}, kv[8] = {}, gvp[8] = {}, kvp[8] = {}; + if (v8) { + gnv[0] = R.gn0.x; gnv[1] = R.gn0.y; gnv[2] = R.gn0.z; gnv[3] = R.gn0.w; + gnv[4] = R.gn1.x; gnv[5] = R.gn1.y; gnv[6] = R.gn1.z; gnv[7] = R.gn1.w; + if (ok_r) { + gv[0] = R.gv0.x; gv[1] = R.gv0.y; gv[2] = R.gv0.z; gv[3] = R.gv0.w; + gv[4] = R.gv1.x; gv[5] = R.gv1.y; gv[6] = R.gv1.z; gv[7] = R.gv1.w; + QKT const* hq = reinterpret_cast(&R.rq); + QKT const* hk = reinterpret_cast(&R.rk); + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { qv[j] = to_f32(hq[j]); kv[j] = to_f32(hk[j]); } + } + if (ok_c) { + gvp[0] = R.gp0.x; gvp[1] = R.gp0.y; gvp[2] = R.gp0.z; gvp[3] = R.gp0.w; + gvp[4] = R.gp1.x; gvp[5] = R.gp1.y; gvp[6] = R.gp1.z; gvp[7] = R.gp1.w; + QKT const* hp = reinterpret_cast(&R.rp); + CUTE_UNROLL + for (int j = 0; j < 8; ++j) kvp[j] = to_f32(hp[j]); + } + } + float oq[8], ok2[8], ob[8]; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { + float gqn = ok_r ? exp2_ftz(gv[j] - gnv[j]) : 0.f; + float gkn = ok_c ? exp2_ftz(gnv[j] - gvp[j]) : 0.f; + oq[j] = qv[j] * gqn; + ok2[j] = kv[j] * gqn; + ob[j] = kvp[j] * gkn; + } + int soff = r * kPadHalf + cg; + *reinterpret_cast(sAq + soff) = pack4_tf32(oq); + *reinterpret_cast(sAq + soff + 4) = pack4_tf32(oq + 4); + *reinterpret_cast(sAk + soff) = pack4_tf32(ok2); + *reinterpret_cast(sAk + soff + 4) = pack4_tf32(ok2 + 4); + *reinterpret_cast(sB + soff) = pack4_tf32(ob); + *reinterpret_cast(sB + soff + 4) = pack4_tf32(ob + 4); + }; + RawPass cur, nxt; + load_pass(0, 0, cur); + for (int kc = 0; kc < n_kchunks; ++kc) { + CUTE_UNROLL + for (int pass = 0; pass < 2; ++pass) { + int const nkc = kc + pass; + if (nkc < n_kchunks) load_pass(nkc, (pass + 1) & 1, nxt); + store_pass(cur, pass); + if (pass == 1) { + __syncwarp(); + mma16_accum(ctx_qk, sAq, sB, SmemLayoutHalf{}); + mma16_accum(ctx_kk, sAk, sB, SmemLayoutHalf{}); + __syncwarp(); + } + cur = nxt; + } + } + } else { + for (int kc = 0; kc < n_kchunks; ++kc) { + int kk = kc * kKC2 + cg; + int rem = K - kk; // valid cols in this 8-col group (may be <= 0) + CUTE_UNROLL + for (int pass = 0; pass < 2; ++pass) { + int r = r0 + pass * 8; + int t = i_tc[c] + r; // row block token + int tp = i_tc[cp] + r; // col block token + bool ok_r = t < Tseq; + bool ok_c = tp < Tseq; + float gnv[8] = {}, gv[8] = {}, qv[8] = {}, kv[8] = {}, gvp[8] = {}, kvp[8] = {}; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { + bool okk = j < rem; + gnv[j] = okk ? gn_base[kk + j] : 0.f; + if (ok_r && okk) { + gv[j] = g[((bos + t) * HV + i_hv) * (int64_t)K + kk + j]; + qv[j] = to_f32(q[((bos + t) * H + i_h) * (int64_t)K + kk + j]); + kv[j] = to_f32(k[((bos + t) * H + i_h) * (int64_t)K + kk + j]); + } + if (ok_c && okk) { + gvp[j] = g[((bos + tp) * HV + i_hv) * (int64_t)K + kk + j]; + kvp[j] = to_f32(k[((bos + tp) * H + i_h) * (int64_t)K + kk + j]); + } + } + float oq[8], ok2[8], ob[8]; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) { + float gqn = (ok_r && j < rem) ? exp2_ftz(gv[j] - gnv[j]) : 0.f; + float gkn = (ok_c && j < rem) ? exp2_ftz(gnv[j] - gvp[j]) : 0.f; + oq[j] = qv[j] * gqn; + ok2[j] = kv[j] * gqn; + ob[j] = kvp[j] * gkn; + } + int soff = r * kPadHalf + cg; + *reinterpret_cast(sAq + soff) = pack4_tf32(oq); + *reinterpret_cast(sAq + soff + 4) = pack4_tf32(oq + 4); + *reinterpret_cast(sAk + soff) = pack4_tf32(ok2); + *reinterpret_cast(sAk + soff + 4) = pack4_tf32(ok2 + 4); + *reinterpret_cast(sB + soff) = pack4_tf32(ob); + *reinterpret_cast(sB + soff + 4) = pack4_tf32(ob + 4); + } + __syncwarp(); + mma16_accum(ctx_qk, sAq, sB, SmemLayoutHalf{}); + mma16_accum(ctx_kk, sAk, sB, SmemLayoutHalf{}); + __syncwarp(); + } + } + // Aqk off-diagonal block: scale applied at store + CUTE_UNROLL + for (int v = 0; v < size(ctx_qk.acc); ++v) { + auto crd = ctx_qk.coord(v); + int i = int(get<0>(crd)), j = int(get<1>(crd)); + if (i_tc[c] + i < Tseq) { + Aqk[((bos + i_tc[c] + i) * HV + i_hv) * (int64_t)BT + cp * kBC + j] = + QKT(ctx_qk.acc(v) * scale); + } + } + // Akk off-diagonal block: beta on rows, kept fp32 in smem + CUTE_UNROLL + for (int v = 0; v < size(ctx_kk.acc); ++v) { + auto crd = ctx_kk.coord(v); + int i = int(get<0>(crd)), j = int(get<1>(crd)); + sAkkOff[p][i * kBC + j] = ctx_kk.acc(v) * sBeta[c][i]; + } + } + __syncthreads(); + + // Phase 4: invert diagonal blocks (non safe_gate only; warp c handles block c) + if (!SAFE_GATE) { + for (int idx = tid; idx < NC * kBC * kBC; idx += blockDim.x) { + int i = (idx / kBC) % kBC; + int j = idx % kBC; + float* p = &sAi[0][0] + idx; + *p = (i > j) ? -*p : 0.f; + } + __syncthreads(); + if (warp < NC && lane < kBC) { + int c = warp; + fwd_subst_16(&sAi[c][0], min(kBC, Tseq - i_tc[c]), lane); + sAi[c][lane * kBC + lane] += 1.f; + } + __syncthreads(); + } + + // Phase 5: merged inverse (block lower-triangular solve), tf32 dots. + // sAiX index for block pair (c, cp) is c*(c-1)/2 + cp. + // The solve is a chain of tiny 16x16 GEMMs; independent links run on + // separate warps through per-warp scratch (sW is dead after phase 1+2), + // cutting the serial depth from 15 chained GEMMs to 6 stages. Every + // product and fp32 partial-sum add keeps the v1 operand and accumulation + // order bit-for-bit (t30 = ((Akk30@Ai00 + Akk31@Ai10) + Akk32@Ai20)). + auto* p5 = reinterpret_cast(&sW[0][0][0]); + TF32* sMAw = p5->sMA[warp]; + TF32* sMBw = p5->sMB[warp]; + if constexpr (NC == 4) { + float (*T)[kBC * kBC] = p5->T; + float (*P)[kBC * kBC] = p5->P; + // stage A: diagonal-times-Akk products, one per warp + if (warp == 0) merge_gemm_16_warp(&sAi[1][0], &sAkkOff[0][0], T[0], 1.f, sMAw, sMBw, lane); + if (warp == 1) merge_gemm_16_warp(&sAi[2][0], &sAkkOff[2][0], T[1], 1.f, sMAw, sMBw, lane); + if (warp == 2) merge_gemm_16_warp(&sAi[3][0], &sAkkOff[5][0], T[2], 1.f, sMAw, sMBw, lane); + if (warp == 3) merge_gemm_16_warp(&sAkkOff[3][0], &sAi[0][0], P[0], 1.f, sMAw, sMBw, lane); + __syncthreads(); + // stage B: two-link merges Ai10/Ai21/Ai32 plus Akk31@Ai11 + if (warp == 0) merge_gemm_16_warp(T[0], &sAi[0][0], &sAiX[0][0], -1.f, sMAw, sMBw, lane); + if (warp == 1) merge_gemm_16_warp(T[1], &sAi[1][0], &sAiX[2][0], -1.f, sMAw, sMBw, lane); + if (warp == 2) merge_gemm_16_warp(T[2], &sAi[2][0], &sAiX[5][0], -1.f, sMAw, sMBw, lane); + if (warp == 3) merge_gemm_16_warp(&sAkkOff[4][0], &sAi[1][0], P[1], 1.f, sMAw, sMBw, lane); + __syncthreads(); + // stage C: remaining products of the three-link sums + if (warp == 0) merge_gemm_16_warp(&sAkkOff[2][0], &sAiX[0][0], P[3], 1.f, sMAw, sMBw, lane); + if (warp == 1) merge_gemm_16_warp(&sAkkOff[1][0], &sAi[0][0], P[2], 1.f, sMAw, sMBw, lane); + if (warp == 2) merge_gemm_16_warp(&sAkkOff[5][0], &sAiX[2][0], P[4], 1.f, sMAw, sMBw, lane); + if (warp == 3) merge_gemm_16_warp(&sAkkOff[4][0], &sAiX[0][0], P[5], 1.f, sMAw, sMBw, lane); + __syncthreads(); + for (int idx = tid; idx < kBC * kBC; idx += blockDim.x) { + T[0][idx] = P[2][idx] + P[3][idx]; // t20 = Akk20@Ai00 + Akk21@Ai10 + T[1][idx] = P[1][idx] + P[4][idx]; // t31 = Akk31@Ai11 + Akk32@Ai21 + T[2][idx] = P[0][idx] + P[5][idx]; // t30 partial = Akk30@Ai00 + Akk31@Ai10 + } + __syncthreads(); + // stage D: Ai20 and Ai31 + if (warp == 0) merge_gemm_16_warp(&sAi[2][0], T[0], &sAiX[1][0], -1.f, sMAw, sMBw, lane); + if (warp == 2) merge_gemm_16_warp(&sAi[3][0], T[1], &sAiX[4][0], -1.f, sMAw, sMBw, lane); + __syncthreads(); + // stage E: last product of the t30 sum + if (warp == 1) merge_gemm_16_warp(&sAkkOff[5][0], &sAiX[1][0], P[2], 1.f, sMAw, sMBw, lane); + __syncthreads(); + for (int idx = tid; idx < kBC * kBC; idx += blockDim.x) T[2][idx] += P[2][idx]; + __syncthreads(); + // stage F: Ai30 + if (warp == 0) merge_gemm_16_warp(&sAi[3][0], T[2], &sAiX[3][0], -1.f, sMAw, sMBw, lane); + __syncthreads(); + } else { + // NC == 2: single off-diagonal block Ai10 = -(Ai11 @ Akk10) @ Ai00 + if (warp == 0) { + merge_gemm_16_warp(&sAi[1][0], &sAkkOff[0][0], p5->T[0], 1.f, sMAw, sMBw, lane); + merge_gemm_16_warp(p5->T[0], &sAi[0][0], &sAiX[0][0], -1.f, sMAw, sMBw, lane); + } + __syncthreads(); + } + + // Phase 6: store the full block-lower-triangular inverse to Akk (bf16). + // Each 16-col row segment is contiguous, so a warp-group of 32 threads + // stores a 16x16 block as 32 16B vectors instead of 256 scalar stores. + // Pair index p enumerates (c, cp) with p = c*(c+1)/2 + cp. + { + int const slot = tid >> 5; // pair slot per pass + int const r = (tid >> 1) & (kBC - 1); // row within the block + int const h8 = tid & 1; // which 8-col half of the row + constexpr int NPAIR = NC * (NC + 1) / 2; + for (int p = slot; p < NPAIR; p += int(blockDim.x >> 5)) { + int c = 0; + while ((c + 1) * (c + 2) / 2 <= p) ++c; + int cp = p - c * (c + 1) / 2; + if (i_tc[c] >= Tseq || i_tc[c] + r >= Tseq) continue; + float const* src = (cp == c) ? &sAi[c][0] : &sAiX[c * (c - 1) / 2 + cp][0]; + QKT vals[8]; + CUTE_UNROLL + for (int j = 0; j < 8; ++j) vals[j] = QKT(src[r * kBC + h8 * 8 + j]); + *reinterpret_cast( + Akk + ((bos + i_tc[c] + r) * HV + i_hv) * (int64_t)BT + cp * kBC + h8 * 8) = + *reinterpret_cast(vals); + } + } +} + +// --------------------------------------------------------------------------- +// Host launchers +// --------------------------------------------------------------------------- + +template +void launch_sub_chunk( + torch::Tensor const& q, + torch::Tensor const& k, + torch::Tensor const& g, + torch::Tensor const& beta, + torch::Tensor& Aqk, + torch::Tensor& Akkd, + float scale, + int64_t chunk_size, + VarlenArgs const& varlen, + int64_t B, + int64_t T, + int64_t H, + int64_t HV, + int64_t K, + cudaStream_t stream +) { + dim3 grid(varlen.NT, B * HV); + dim3 block(128); + + #define LAUNCH_SUB_CHUNK(BT, IS_VARLEN) \ + chunk_kda_fwd_kernel_intra_sub_chunk_cuda<<>>( \ + reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), \ + g.data_ptr(), reinterpret_cast(beta.data_ptr()), \ + reinterpret_cast(Aqk.data_ptr()), Akkd.data_ptr(), \ + scale, varlen.cu_seqlens, varlen.chunk_indices, int(T), int(H), int(HV), int(K)) + + if (chunk_size == 64) { + if (varlen.cu_seqlens) { LAUNCH_SUB_CHUNK(64, true); } else { LAUNCH_SUB_CHUNK(64, false); } + } else { + if (varlen.cu_seqlens) { LAUNCH_SUB_CHUNK(32, true); } else { LAUNCH_SUB_CHUNK(32, false); } + } + #undef LAUNCH_SUB_CHUNK +} + +template +void launch_token_parallel( + torch::Tensor const& q, + torch::Tensor const& k, + torch::Tensor const& g, + torch::Tensor const& beta, + torch::Tensor& Aqk, + torch::Tensor& Akkd, + float scale, + int64_t chunk_size, + std::optional const& cu_seqlens, + int64_t B, + int64_t T, + int64_t H, + int64_t HV, + int64_t K, + cudaStream_t stream +) { + int64_t N = cu_seqlens.has_value() ? cu_seqlens->numel() - 1 : B; + dim3 grid(B * T, (HV + 3) / 4); + dim3 block(128); + int64_t const* cu_ptr = + cu_seqlens.has_value() ? cu_seqlens->data_ptr() : nullptr; + + #define LAUNCH_TOKEN_PARALLEL(BT, IS_VARLEN) \ + chunk_kda_fwd_kernel_intra_token_parallel_cuda<<>>( \ + reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), \ + g.data_ptr(), reinterpret_cast(beta.data_ptr()), \ + reinterpret_cast(Aqk.data_ptr()), Akkd.data_ptr(), \ + scale, cu_ptr, N, int(T), int(H), int(HV), int(K)) + + if (chunk_size == 64) { + if (cu_ptr) { LAUNCH_TOKEN_PARALLEL(64, true); } else { LAUNCH_TOKEN_PARALLEL(64, false); } + } else { + if (cu_ptr) { LAUNCH_TOKEN_PARALLEL(32, true); } else { LAUNCH_TOKEN_PARALLEL(32, false); } + } + #undef LAUNCH_TOKEN_PARALLEL +} + +template +void launch_inter_solve_fused( + torch::Tensor const& q, + torch::Tensor const& k, + torch::Tensor const& g, + torch::Tensor const& beta, + torch::Tensor& Aqk, + torch::Tensor const& Akkd, + torch::Tensor& Akk, + float scale, + int64_t chunk_size, + bool safe_gate, + VarlenArgs const& varlen, + int64_t B, + int64_t T, + int64_t H, + int64_t HV, + int64_t K, + cudaStream_t stream +) { + dim3 grid(varlen.NT, B * HV); + dim3 block(128); + + #define LAUNCH_FUSED(BT, SAFE_GATE, IS_VARLEN) \ + chunk_kda_fwd_kernel_inter_solve_fused_cuda<<>>( \ + reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), \ + g.data_ptr(), reinterpret_cast(beta.data_ptr()), \ + reinterpret_cast(Aqk.data_ptr()), Akkd.data_ptr(), \ + reinterpret_cast(Akk.data_ptr()), \ + scale, varlen.cu_seqlens, varlen.chunk_indices, int(T), int(H), int(HV), int(K)) + + #define DISPATCH_FUSED_VARLEN(BT, SAFE_GATE) \ + if (varlen.cu_seqlens) { LAUNCH_FUSED(BT, SAFE_GATE, true); } \ + else { LAUNCH_FUSED(BT, SAFE_GATE, false); } + #define DISPATCH_FUSED_SAFE(BT) \ + if (safe_gate) { DISPATCH_FUSED_VARLEN(BT, true); } \ + else { DISPATCH_FUSED_VARLEN(BT, false); } + + if (chunk_size == 64) { DISPATCH_FUSED_SAFE(64); } + else { DISPATCH_FUSED_SAFE(32); } + + #undef DISPATCH_FUSED_SAFE + #undef DISPATCH_FUSED_VARLEN + #undef LAUNCH_FUSED +} + +template