From d0cb63797af3b1ec87b1618e2f08a8da6172e0dd Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Thu, 16 Apr 2026 15:09:08 +0800 Subject: [PATCH 01/34] doc: update README with bigger logo (#53) --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 25dc13dd..81385c83 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@
-cuLA Logo +cuLA Logo # cuLA — CUDA Linear Attention From e8a70986e53894706e14749823678048c3f525c7 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Thu, 16 Apr 2026 15:51:38 +0800 Subject: [PATCH 02/34] [KDA] Adapt recompute_wu and delta_h for backward (#54) * adapt reomp_wu and delta_h into backward * add deter check * lint * Update cula/kda/chunk_bwd.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * update roadmap --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- README.md | 2 +- benchmarks/bench_kda_fwd_bwd_e2e.py | 623 ++++++++++++++++++++++++++++ cula/kda/chunk.py | 2 +- cula/kda/chunk_bwd.py | 604 +++++++++++++++++++++++++++ 4 files changed, 1229 insertions(+), 2 deletions(-) create mode 100644 benchmarks/bench_kda_fwd_bwd_e2e.py create mode 100644 cula/kda/chunk_bwd.py diff --git a/README.md b/README.md index 81385c83..d5764052 100644 --- a/README.md +++ b/README.md @@ -184,7 +184,7 @@ See [REPO_LAYOUT.md](REPO_LAYOUT.md) for the full directory structure and a summ * [x] Modular KDA Forward (SM10X, compatible with [Kimi CP](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/cp/README.md)) * [x] kda chunk intra * [x] chunk gated delta h - * [ ] recompute wu + * [x] recompute wu * [x] chunk fwd o * [ ] Modular GDN Forward / Backward Kernels (compatible with [Kimi CP](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/cp/README.md)) diff --git a/benchmarks/bench_kda_fwd_bwd_e2e.py b/benchmarks/bench_kda_fwd_bwd_e2e.py new file mode 100644 index 00000000..c6b4117b --- /dev/null +++ b/benchmarks/bench_kda_fwd_bwd_e2e.py @@ -0,0 +1,623 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +bench_kda_fwd_bwd_e2e.py — Benchmark: cuLA CuTe DSL vs FLA Triton baseline + for chunk_kda forward + backward (end-to-end) + +Compares: + - Accuracy: err_ratio, relative max diff between cuLA and FLA outputs & gradients + - Performance: kernel execution time (ms) with CUDA events + +Modes: + - Fixed-length: B=1, B=2 with various T + - Varlen: ~20 seqs with 2-3x length variation + +Phases: + - forward: forward pass only + - e2e: forward + backward (end-to-end) + +Usage: + python bench_kda_fwd_bwd_e2e.py [--mode fixed|varlen|both] [--phase forward|e2e] [--ncu] + +With --ncu, warmup=1 and iters=1 for ncu profiling: + ncu --set full -o report python bench_kda_fwd_bwd_e2e.py --mode varlen --ncu +""" + +import argparse +import os +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) + +import torch +from fla.ops.kda import chunk_kda as fla_chunk_kda + +from benchmarks.utils import ( + SEED, + build_varlen_configs, + exclusive_cumsum, + generate_random_seq_lens, + prepare_safe_gate_inputs, + set_seed, +) +from cula.kda import chunk_kda as cula_chunk_kda + +# ============================================================ +# Constants +# ============================================================ +H, D = 64, 128 +WARMUP = 25 +N_ITERS = 100 +NCU_MODE = False +SANITIZER_MODE = False +DISABLE_RECOMPUTE = False +PHASE = "e2e" # "forward" or "e2e" + + +# ============================================================ +# Helpers +# ============================================================ +def time_kernel(fn, warmup=None, n_iters=None): + if warmup is None: + warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + if n_iters is None: + n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() + for _ in range(n_iters): + fn() + end_evt.record() + torch.cuda.synchronize() + return start_evt.elapsed_time(end_evt) / n_iters + + +def accuracy_stats(ref, out): + """Compute err_ratio, relative max diff, and mean absolute difference.""" + ref_f = ref.float() + out_f = out.float() + diff = (ref_f - out_f).abs() + err = diff.flatten().pow(2).mean().sqrt().item() + base = ref_f.flatten().pow(2).mean().sqrt().item() + err_ratio = err / (base + 1e-8) + max_diff = diff.max().item() + denom = ref_f.abs().max().item() + rel_max = max_diff / denom if denom > 0 else 0.0 + mean_diff = diff.mean().item() + return err_ratio, rel_max, mean_diff + + +def run_kda_e2e(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound, do, dht, fn): + """Run KDA forward (+ backward if PHASE == 'e2e'). + + Clears gradients, runs forward, optionally backward. + """ + q.grad = None + k.grad = None + v.grad = None + g.grad = None + beta.grad = None + init_state.grad = None + + out, ht = fn( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + initial_state=init_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=lower_bound, + disable_recompute=DISABLE_RECOMPUTE, + ) + if PHASE == "e2e": + out.backward(do) + return out, ht + + +def run_kda_e2e_with_grads(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound, do, dht, fn): + """Run KDA forward + backward and return outputs + gradients for accuracy check.""" + q_c = q.detach().clone().requires_grad_(True) + k_c = k.detach().clone().requires_grad_(True) + v_c = v.detach().clone().requires_grad_(True) + g_c = g.detach().clone().requires_grad_(True) + b_c = beta.detach().clone().requires_grad_(True) + h_c = init_state.detach().clone().requires_grad_(True) + + out, ht = fn( + q=q_c, + k=k_c, + v=v_c, + g=g_c, + beta=b_c, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h_c, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=lower_bound, + disable_recompute=DISABLE_RECOMPUTE, + ) + loss = (out * do).sum() + (ht * dht).sum() + loss.backward() + + return dict( + o=out, + ht=ht, + dq=q_c.grad, + dk=k_c.grad, + dv=v_c.grad, + dg=g_c.grad, + dbeta=b_c.grad, + dh0=h_c.grad, + ) + + +# ============================================================ +# Determinism check +# ============================================================ +def check_determinism(num_seqs=5, T=512, iters=20): + """Verify that cuLA chunk_kda produces identical outputs across repeated runs.""" + device = torch.device("cuda") + set_seed(SEED) + + seq_lens = generate_random_seq_lens(num_seqs, T, 63, seed=SEED) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + + inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] + + set_seed(SEED + 1) + do = torch.randn_like(v) + dht = torch.randn_like(init_state) + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + init_state=init_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + do=do, + dht=dht, + ) + + ref = run_kda_e2e_with_grads(**common, fn=cula_chunk_kda) + for i in range(iters): + out = run_kda_e2e_with_grads(**common, fn=cula_chunk_kda) + for name in ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"): + assert torch.equal(out[name], ref[name]), f"[determinism] cuLA {name} mismatch at iter {i}" + return True + + +# ============================================================ +# Fixed-length benchmark +# ============================================================ +def bench_fixed(configs): + print("\n" + "=" * 120) + print(f" Fixed-Length E2E Benchmark: cuLA vs FLA phase={PHASE} disable_recompute={DISABLE_RECOMPUTE}") + print("=" * 120) + results = [] + + for B, T in configs: + set_seed(SEED) + device = torch.device("cuda") + torch.cuda.empty_cache() + + seq_lens = [T] * B + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + + inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] + + # Generate do, dht for backward + set_seed(SEED + 1) + do = torch.randn_like(v) + dht = torch.randn_like(init_state) + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + init_state=init_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + do=do, + dht=dht, + ) + + # Accuracy: compare outputs and gradients + acc = {} + if PHASE == "e2e": + fla_results = run_kda_e2e_with_grads(**common, fn=fla_chunk_kda) + cula_results = run_kda_e2e_with_grads(**common, fn=cula_chunk_kda) + torch.cuda.synchronize() + + for name in ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"): + err_ratio, rel_max, mean_diff = accuracy_stats(fla_results[name], cula_results[name]) + acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + else: + # forward-only accuracy + o_fla, ht_fla = run_kda_e2e(**common, fn=fla_chunk_kda) + o_cula, ht_cula = run_kda_e2e(**common, fn=cula_chunk_kda) + torch.cuda.synchronize() + for name, ref, out in [("o", o_fla, o_cula), ("ht", ht_fla, ht_cula)]: + err_ratio, rel_max, mean_diff = accuracy_stats(ref, out) + acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + + # For timing, use leaf tensors with requires_grad + q_t = q.detach().clone().requires_grad_(True) + k_t = k.detach().clone().requires_grad_(True) + v_t = v.detach().clone().requires_grad_(True) + g_t = g.detach().clone().requires_grad_(True) + beta_t = beta.detach().clone().requires_grad_(True) + h0_t = init_state.detach().clone().requires_grad_(True) + + timing_common = dict( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + init_state=h0_t, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + do=do, + dht=dht, + ) + + def fn_fla(**kw): + return lambda: run_kda_e2e(**kw, fn=fla_chunk_kda) + + def fn_cula(**kw): + return lambda: run_kda_e2e(**kw, fn=cula_chunk_kda) + + ms_fla = time_kernel(fn_fla(**timing_common)) + ms_cula = time_kernel(fn_cula(**timing_common)) + speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") + + r = { + "B": B, + "T": T, + "accuracy": acc, + "ms_fla": ms_fla, + "ms_cula": ms_cula, + "speedup": speedup, + } + results.append(r) + + del q, k, v, g, beta, A_log, dt_bias, inputs, do, dht + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Varlen benchmark +# ============================================================ +def bench_varlen(configs): + print("\n" + "=" * 120) + print(f" Varlen E2E Benchmark: cuLA vs FLA phase={PHASE} disable_recompute={DISABLE_RECOMPUTE}") + print("=" * 120) + results = [] + + for seq_lens, total_len, dist in configs: + set_seed(SEED) + device = torch.device("cuda") + torch.cuda.empty_cache() + + T = total_len + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + + inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] + + # Generate do, dht for backward + set_seed(SEED + 1) + do = torch.randn_like(v) + dht = torch.randn_like(init_state) + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + init_state=init_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + do=do, + dht=dht, + ) + + # Accuracy: compare outputs and gradients + acc = {} + if PHASE == "e2e": + fla_results = run_kda_e2e_with_grads(**common, fn=fla_chunk_kda) + cula_results = run_kda_e2e_with_grads(**common, fn=cula_chunk_kda) + torch.cuda.synchronize() + + for name in ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"): + err_ratio, rel_max, mean_diff = accuracy_stats(fla_results[name], cula_results[name]) + acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + else: + o_fla, ht_fla = run_kda_e2e(**common, fn=fla_chunk_kda) + o_cula, ht_cula = run_kda_e2e(**common, fn=cula_chunk_kda) + torch.cuda.synchronize() + for name, ref, out in [("o", o_fla, o_cula), ("ht", ht_fla, ht_cula)]: + err_ratio, rel_max, mean_diff = accuracy_stats(ref, out) + acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + + # For timing, use leaf tensors with requires_grad + q_t = q.detach().clone().requires_grad_(True) + k_t = k.detach().clone().requires_grad_(True) + v_t = v.detach().clone().requires_grad_(True) + g_t = g.detach().clone().requires_grad_(True) + beta_t = beta.detach().clone().requires_grad_(True) + h0_t = init_state.detach().clone().requires_grad_(True) + + timing_common = dict( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + init_state=h0_t, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + do=do, + dht=dht, + ) + + def fn_fla(**kw): + return lambda: run_kda_e2e(**kw, fn=fla_chunk_kda) + + def fn_cula(**kw): + return lambda: run_kda_e2e(**kw, fn=cula_chunk_kda) + + ms_fla = time_kernel(fn_fla(**timing_common)) + ms_cula = time_kernel(fn_cula(**timing_common)) + speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") + + n_seqs = len(seq_lens) + min_l, max_l = min(seq_lens), max(seq_lens) + avg_l = T // n_seqs + tag = f"{dist:>7s} {n_seqs:>2d}seqs T={T} [{min_l}..{max_l}] avg={avg_l}" + + r = { + "tag": tag, + "dist": dist, + "T_total": T, + "n_seqs": n_seqs, + "accuracy": acc, + "ms_fla": ms_fla, + "ms_cula": ms_cula, + "speedup": speedup, + } + results.append(r) + + del q, k, v, g, beta, A_log, dt_bias, inputs, do, dht + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Report +# ============================================================ +def print_report(fixed_results, varlen_results): + sep = "=" * 130 + print(f"\n\n{sep}") + print(" BENCHMARK REPORT: chunk_kda forward+backward (E2E)") + print(" cuLA CuTe DSL vs FLA Triton") + print( + f" H={H} D={D} dtype=bf16 safe_gate=True phase={PHASE} disable_recompute={DISABLE_RECOMPUTE}" + ) + wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") + print(f" Warmup={wu} Iters={ni}{mode_tag}") + print(sep) + + # Determine which accuracy keys to show + if PHASE == "e2e": + acc_keys = ["o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"] + else: + acc_keys = ["o", "ht"] + + acc_header = " ".join(f"{k:>10s}" for k in acc_keys) + + if fixed_results: + print("\n [Fixed-Length]") + print(f" {'─' * 125}") + + # Header + print(f" {'B':>3s} {'T':>5s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s} │ {'':>10s}{acc_header}") + print(f" {'─' * 125}") + + for r in fixed_results: + rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in acc_keys) + err_ratio_vals = " ".join(f"{r['accuracy'].get(k, {}).get('err_ratio', 0.0):10.6f}" for k in acc_keys) + # Line 1: timing + rel_max + print( + f" {r['B']:3d} {r['T']:5d} │ " + f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x │ " + f"{'rel_max:':>10s}{rel_max_vals}" + ) + # Line 2: err_ratio (no timing columns) + print(f" {'':3s} {'':5s} │ {'':9s} {'':11s} {'':8s} │ {'err_ratio:':>10s}{err_ratio_vals}") + print(f" {'─' * 125}") + + if varlen_results: + print("\n [Varlen]") + print(f" {'─' * 140}") + + print(f" {'Config':>45s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s} │ {'':>10s}{acc_header}") + print(f" {'─' * 140}") + + for r in varlen_results: + rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in acc_keys) + err_ratio_vals = " ".join(f"{r['accuracy'].get(k, {}).get('err_ratio', 0.0):10.6f}" for k in acc_keys) + # Line 1: timing + rel_max + print( + f" {r['tag']:>45s} │ " + f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x │ " + f"{'rel_max:':>10s}{rel_max_vals}" + ) + # Line 2: err_ratio (no config/timing columns) + print(f" {'':>45s} │ {'':9s} {'':11s} {'':8s} │ {'err_ratio:':>10s}{err_ratio_vals}") + print(f" {'─' * 140}") + + print(f"\n{sep}\n") + + +# ============================================================ +# Main +# ============================================================ +def main(): + parser = argparse.ArgumentParser(description="bench_kda_fwd_bwd_e2e: cuLA vs FLA (forward + backward)") + parser.add_argument( + "--mode", + type=str, + default="both", + choices=["fixed", "varlen", "both"], + help="Which benchmark mode to run (default: both)", + ) + parser.add_argument( + "--phase", + type=str, + default="e2e", + choices=["forward", "e2e"], + help="Benchmark phase: forward only or end-to-end (default: e2e)", + ) + parser.add_argument( + "--ncu", + action="store_true", + help="NCU profiling mode: warmup=1, iters=1", + ) + parser.add_argument( + "--sanitizer", + action="store_true", + help="Sanitizer mode: warmup=1, iters=1 (avoid Triton memory leak under compute-sanitizer)", + ) + parser.add_argument( + "--disable_recompute", + action="store_true", + help="Disable recompute in both FLA and cuLA (pre-compute QG)", + ) + parser.add_argument( + "--check_determinism", + action="store_true", + help="Run determinism check: verify cuLA produces identical outputs across repeated runs", + ) + args = parser.parse_args() + + global NCU_MODE, SANITIZER_MODE, DISABLE_RECOMPUTE, PHASE + if args.ncu: + NCU_MODE = True + print("[NCU mode] warmup=1, iters=1") + if args.sanitizer: + SANITIZER_MODE = True + print("[Sanitizer mode] warmup=1, iters=1") + if args.disable_recompute: + DISABLE_RECOMPUTE = True + print("[Disable recompute] pre-compute QG in forward") + PHASE = args.phase + + if args.check_determinism: + det_configs = [(5, 1024), (10, 4096), (10, 8192), (10, 16384)] + print("\n[Determinism Check] cuLA chunk_kda E2E ...") + for num_seqs, T in det_configs: + result = check_determinism(num_seqs=num_seqs, T=T, iters=20) + print(f" num_seqs={num_seqs} T={T:5d} iters=20 {'PASS' if result else 'FAIL'}") + print("[Determinism Check] All passed.\n") + return + + fixed_configs = [ + # (B, T) + (1, 512), + (1, 1024), + (1, 4096), + (1, 8192), + (1, 16384), + (2, 512), + (2, 1024), + (2, 4096), + (2, 8192), + (2, 16384), + ] + + varlen_configs = build_varlen_configs( + num_seqs_list=(10, 20), + total_lens=(4096, 8192, 16384), + dists=("uniform", "random", "skewed"), + ) + + fixed_res, varlen_res = [], [] + + if args.mode in ("fixed", "both"): + fixed_res = bench_fixed(fixed_configs) + + if args.mode in ("varlen", "both"): + varlen_res = bench_varlen(varlen_configs) + + print_report(fixed_res, varlen_res) + return fixed_res, varlen_res + + +if __name__ == "__main__": + main() diff --git a/cula/kda/chunk.py b/cula/kda/chunk.py index 3e4a7ffd..b89c7806 100644 --- a/cula/kda/chunk.py +++ b/cula/kda/chunk.py @@ -18,10 +18,10 @@ import torch from fla.modules.l2norm import l2norm_bwd, l2norm_fwd from fla.ops.cp import FLACPContext -from fla.ops.kda.chunk_bwd import chunk_kda_bwd from fla.ops.utils.index import prepare_chunk_indices from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +from cula.kda.chunk_bwd import chunk_kda_bwd from cula.kda.chunk_fwd import chunk_kda_fwd diff --git a/cula/kda/chunk_bwd.py b/cula/kda/chunk_bwd.py new file mode 100644 index 00000000..859b6be9 --- /dev/null +++ b/cula/kda/chunk_bwd.py @@ -0,0 +1,604 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import importlib + +import torch +import triton +import triton.language as tl +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + expand_h0, +) +from fla.ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum +from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +import cula.cudac as cula_cuda +from cula.kda.chunk_intra import chunk_kda_bwd_intra +from cula.utils import prepare_uniform_cu_seqlens + +_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h") +chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [64, 128] if check_shared_mem("ampere") else [16, 32] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_bwd_kernel_dAv( + q, + k, + v, + A, + do, + dv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + dA += (bos * H + i_h) * BT + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (V, T), (1, H * V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BV, BT] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BT, BT] + b_dA += tl.dot(b_do, b_v) + # [BT, BV] + b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.0) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=["BT", "TRANSPOSE_STATE"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_bwd_kernel_wy_dqkg_fused( + q, + k, + v, + v_new, + g, + beta, + A, + h, + do, + dh, + dq, + dk, + dv, + dv2, + dg, + db, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t.to(tl.int64) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int64) + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_last = o_t == min(T, i_t * BT + BT) - 1 + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + v_new += (bos * H + i_h) * V + g += (bos * H + i_h) * K + beta += bos * H + i_h + A += (bos * H + i_h) * BT + h += (i_tg * H + i_h) * K * V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K * V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + dv2 += (bos * H + i_h) * V + dg += (bos * H + i_h) * K + db += bos * H + i_h + dA += (bos * H + i_h) * BT + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + b_db = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + + p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H * K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk = tl.zeros([BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v_new = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + if TRANSPOSE_STATE: + p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT, BV] + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BT, BV] + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + + b_dgk += tl.sum(b_h * b_dh, axis=0) + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) + b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) + tl.debug_barrier() # DO NOT REMOVE THIS LINE! + if i_k == 0: + p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_v = tl.load(p_v, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dv, tl.trans(b_v)) + + b_dvb = tl.dot(b_A, b_dv) + b_dv2 = b_dvb * b_beta[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + + tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) + + b_gk_exp = exp2(b_g) + b_gb = b_gk_exp * b_beta[:, None] + b_dgk *= exp2(b_gn) + b_dq = b_dq * b_gk_exp * scale + b_dk = b_dk * tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0) + + b_kg = b_k * b_gk_exp + + b_dw = -b_dw.to(b_A.dtype) + b_dA += tl.dot(b_dw, tl.trans(b_kg.to(b_A.dtype))) + + b_dkgb = tl.dot(b_A, b_dw) + b_db += tl.sum(b_dkgb * b_kg, 1) + + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kdk = b_k * b_dk + b_dgk += tl.sum(b_kdk, axis=0) + b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_beta[:, None] + b_dk = b_dk + b_dkgb * b_gb + + p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA * b_beta[None, :], 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + b_dA = tl.where(m_A, -b_dA, 0) + + p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def chunk_kda_bwd_dAv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + A: torch.Tensor | None = None, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # H100 can have larger block size + if check_shared_mem("hopper", k.device.index): + CONST_TILING = 128 + elif check_shared_mem(): + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dA = v.new_empty(B, T, H, BT, dtype=torch.float) + dv = torch.empty_like(do) + grid = (NT, B * H) + chunk_kda_bwd_kernel_dAv[grid]( + q=q, + k=k, + v=v, + A=A, + do=do, + dv=dv, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dA, dv + + +def chunk_kda_bwd_wy_dqkg_fused( + 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: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + transpose_state_layout: bool = False, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + dv2 = torch.empty_like(v) + dg = torch.empty_like(g, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + + grid = (NT, B * H) + chunk_kda_bwd_kernel_wy_dqkg_fused[grid]( + q=q, + k=k, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=A, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dv=dv, + dv2=dv2, + dg=dg, + db=db, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + TRANSPOSE_STATE=transpose_state_layout, + ) + dv = dv2 + return dq, dk, dv, db, dg, dA + + +def chunk_kda_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + g: torch.Tensor | None = None, + g_org: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + disable_recompute: bool = False, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + **kwargs, +): + assert transpose_state_layout is False, "transpose_state_layout=True is not supported for training." + if disable_recompute is False: + B, T, _, _ = k.shape + if use_gate_in_kernel: + g = kda_gate_chunk_cumsum( + g=g_org, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + reset_cu_seqlens = False + if cu_seqlens is None: + reset_cu_seqlens = True + cu_seqlens = prepare_uniform_cu_seqlens(B, T, q.device, torch.int32) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + w = torch.empty_like(k) + u = torch.empty_like(v) + qg = torch.empty_like(q) if q is not None else None + kg = torch.empty_like(k) if g is not None else None + cula_cuda.recompute_w_u_cuda(k, v, beta, Akk, g, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q, qg) + if cp_context is not None: + # Restore the full initial_state tensor from the compressed version. + # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. + initial_state = expand_h0(initial_state, context=cp_context) + if reset_cu_seqlens: + cu_seqlens = None + chunk_indices = None + # TODO: update to support only varlen (1,T,H,D) format + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + w, u, qg, kg, v_new, h = kwargs["w"], kwargs["u"], kwargs["qg"], kwargs["kg"], kwargs["v_new"], kwargs["h"] + if cp_context is not None: + # Restore the full initial_state tensor from the compressed version. + # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. + initial_state = expand_h0(initial_state, context=cp_context) + + # dAqk = do @ v.T + # dv = A @ do + dAqk, dv = chunk_kda_bwd_dAv( + q=q, + k=k, + v=v_new, + do=do, + A=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=qg, + k=kg, + w=w, + do=do, + dv=dv, + gk=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + use_exp2=True, + context=cp_context, + transpose_state_layout=transpose_state_layout, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=qg, + k=kg, + w=w, + gk=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=transpose_state_layout, + ) + + dq, dk, dv, db, dg, dAkk = 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=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + transpose_state_layout=transpose_state_layout, + ) + + dq, dk, db, dg = chunk_kda_bwd_intra( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dk=dk, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + ) + + dA, dbias = None, None + dg = chunk_local_cumsum( + dg, + chunk_size=chunk_size, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + if use_gate_in_kernel: + dg, dA, dbias = kda_gate_bwd(g=g_org, A_log=A_log, dt_bias=dt_bias, dyg=dg, lower_bound=lower_bound) + + return dq, dk, dv, db, dg, dh0, dA, dbias From a78819fbdfec51955bc4f2e53408358c927414c8 Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Sun, 19 Apr 2026 00:15:09 +0800 Subject: [PATCH 03/34] perf(kda_decode): optimize GMEM coalescing, V-tile parallelism (#52) * perf(kda_decode): optimize GMEM coalescing, V-tile parallelism, use small-batch kernel * bench(kda_decode): add auto doc generation * bench(kda_decode): update GB200 results * chore: drop unused naming * chore: move the cutedsl kda_decode to cula.ops * chore: move the cutedsl la_decode to cula.ops --------- Co-authored-by: boyu.zbw --- BENCHMARK_KDA_DECODE_GB200.md | 145 ++ BENCHMARK_KDA_DECODE_H203E.md | 145 ++ benchmarks/bench_kda_decode.py | 462 ++++++ benchmarks/bench_la_decode_vs_fla.py | 2 +- cula/kda/__init__.py | 3 +- cula/kda/kda_decode.py | 1525 ------------------- cula/lightning/__init__.py | 2 +- cula/ops/__init__.py | 9 + cula/ops/kda_decode.py | 2091 ++++++++++++++++++++++++++ cula/ops/kda_decode_fla.py | 247 +++ cula/{lightning => ops}/la_decode.py | 0 tests/test_kda_decode.py | 268 +++- tests/test_la_decode.py | 2 +- 13 files changed, 3335 insertions(+), 1566 deletions(-) create mode 100644 BENCHMARK_KDA_DECODE_GB200.md create mode 100644 BENCHMARK_KDA_DECODE_H203E.md create mode 100644 benchmarks/bench_kda_decode.py delete mode 100644 cula/kda/kda_decode.py create mode 100644 cula/ops/kda_decode.py create mode 100644 cula/ops/kda_decode_fla.py rename cula/{lightning => ops}/la_decode.py (100%) diff --git a/BENCHMARK_KDA_DECODE_GB200.md b/BENCHMARK_KDA_DECODE_GB200.md new file mode 100644 index 00000000..b1ebd4c2 --- /dev/null +++ b/BENCHMARK_KDA_DECODE_GB200.md @@ -0,0 +1,145 @@ +# Benchmark Results - KDA Decode + +> Auto-generated by `benchmarks/bench_kda_decode.py` on 2026-04-18 23:38:47. + +> **GPU:** NVIDIA GB200 | **CUDA:** 13.0 | **PyTorch:** 2.11.0+cu130 | **Python:** 3.12.3 + +> Decode setting: single-token (T=1), K=128, V=128. + +## Summary + +- v-last speedup (FLA/cuLA): avg=1.02x, min=0.94x, max=1.16x +- k-last speedup (FLA/cuLA): avg=1.15x, min=1.01x, max=1.31x +- Batch sizes: [1, 4, 8, 16, 32, 64, 128, 256] +- Q/K heads (H): [8, 32, 64] +- V heads (HV): 128 +- Timing params: warmup=30, rep=200, ncu_mode=False + +## KDA Decode (H=8, HV=128, K=128, V=128) + +### Performance + +| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | +|---|------------------:|------------------:|----------------:|---------------:|---------------:| +| 1 | 0.0423 | 0.0408 | 0.0490 | **1.16x** | **1.20x** | +| 4 | 0.0533 | 0.0495 | 0.0503 | **0.94x** | **1.02x** | +| 8 | 0.0409 | 0.0379 | 0.0415 | **1.01x** | **1.09x** | +| 16 | 0.0705 | 0.0623 | 0.0693 | **0.98x** | **1.11x** | +| 32 | 0.1382 | 0.1194 | 0.1313 | **0.95x** | **1.10x** | +| 64 | 0.2630 | 0.2322 | 0.2505 | **0.95x** | **1.08x** | +| 128 | 0.5008 | 0.4275 | 0.4892 | **0.98x** | **1.14x** | +| 256 | 0.9801 | 0.8387 | 0.9666 | **0.99x** | **1.15x** | + +### Accuracy (Output) + +| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel | +|---|----------------:|---------------:|----------------:|---------------:| +| 1 | 2.980e-08 | 1.359e-04 | 2.980e-08 | 1.359e-04 | +| 4 | 1.490e-08 | 9.084e-05 | 1.490e-08 | 9.084e-05 | +| 8 | 8.942e-08 | 7.267e-04 | 8.942e-08 | 7.267e-04 | +| 16 | 1.350e-07 | 1.453e-03 | 1.350e-07 | 1.453e-03 | +| 32 | 1.315e-07 | 1.073e-03 | 1.315e-07 | 1.073e-03 | +| 64 | 1.212e-07 | 1.073e-03 | 1.212e-07 | 1.073e-03 | +| 128 | 1.035e-07 | 7.440e-04 | 1.035e-07 | 7.440e-04 | +| 256 | 1.143e-07 | 1.603e-03 | 1.143e-07 | 1.603e-03 | + +### Accuracy (State) + +| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel | +|---|------------------:|-----------------:|------------------:|-----------------:| +| 1 | 4.369e-09 | 2.442e-07 | 4.369e-09 | 2.442e-07 | +| 4 | 3.985e-09 | 1.784e-07 | 3.985e-09 | 1.784e-07 | +| 8 | 4.310e-09 | 2.378e-07 | 4.310e-09 | 2.378e-07 | +| 16 | 4.269e-09 | 2.378e-07 | 4.269e-09 | 2.378e-07 | +| 32 | 4.244e-09 | 2.378e-07 | 4.244e-09 | 2.378e-07 | +| 64 | 4.254e-09 | 2.071e-07 | 4.254e-09 | 2.071e-07 | +| 128 | 4.349e-09 | 2.966e-07 | 4.349e-09 | 2.966e-07 | +| 256 | 4.390e-09 | 2.382e-07 | 4.390e-09 | 2.382e-07 | + +## KDA Decode (H=32, HV=128, K=128, V=128) + +### Performance + +| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | +|---|------------------:|------------------:|----------------:|---------------:|---------------:| +| 1 | 0.0455 | 0.0441 | 0.0499 | **1.09x** | **1.13x** | +| 4 | 0.0532 | 0.0499 | 0.0505 | **0.95x** | **1.01x** | +| 8 | 0.0407 | 0.0363 | 0.0431 | **1.06x** | **1.19x** | +| 16 | 0.0746 | 0.0664 | 0.0705 | **0.95x** | **1.06x** | +| 32 | 0.1314 | 0.1135 | 0.1312 | **1.00x** | **1.16x** | +| 64 | 0.2538 | 0.2227 | 0.2516 | **0.99x** | **1.13x** | +| 128 | 0.4659 | 0.3944 | 0.4962 | **1.06x** | **1.26x** | +| 256 | 0.9040 | 0.7630 | 0.9710 | **1.07x** | **1.27x** | + +### Accuracy (Output) + +| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel | +|---|----------------:|---------------:|----------------:|---------------:| +| 1 | 3.255e-11 | 9.719e-08 | 3.255e-11 | 9.719e-08 | +| 4 | 1.195e-07 | 7.962e-04 | 1.195e-07 | 7.962e-04 | +| 8 | 8.492e-08 | 6.793e-04 | 8.492e-08 | 6.793e-04 | +| 16 | 6.297e-08 | 6.614e-04 | 6.297e-08 | 6.614e-04 | +| 32 | 1.012e-07 | 5.274e-04 | 1.012e-07 | 5.274e-04 | +| 64 | 1.039e-07 | 1.055e-03 | 1.039e-07 | 1.055e-03 | +| 128 | 1.409e-07 | 1.374e-03 | 1.409e-07 | 1.374e-03 | +| 256 | 1.113e-07 | 1.603e-03 | 1.113e-07 | 1.603e-03 | + +### Accuracy (State) + +| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel | +|---|------------------:|-----------------:|------------------:|-----------------:| +| 1 | 3.774e-09 | 1.763e-07 | 3.774e-09 | 1.763e-07 | +| 4 | 4.163e-09 | 1.661e-07 | 4.163e-09 | 1.661e-07 | +| 8 | 4.248e-09 | 1.993e-07 | 4.248e-09 | 1.993e-07 | +| 16 | 4.296e-09 | 2.451e-07 | 4.296e-09 | 2.451e-07 | +| 32 | 4.345e-09 | 2.178e-07 | 4.345e-09 | 2.178e-07 | +| 64 | 4.384e-09 | 2.178e-07 | 4.384e-09 | 2.178e-07 | +| 128 | 4.400e-09 | 2.194e-07 | 4.400e-09 | 2.194e-07 | +| 256 | 4.430e-09 | 2.091e-07 | 4.430e-09 | 2.091e-07 | + +## KDA Decode (H=64, HV=128, K=128, V=128) + +### Performance + +| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | +|---|------------------:|------------------:|----------------:|---------------:|---------------:| +| 1 | 0.0453 | 0.0444 | 0.0505 | **1.11x** | **1.14x** | +| 4 | 0.0544 | 0.0502 | 0.0518 | **0.95x** | **1.03x** | +| 8 | 0.0401 | 0.0359 | 0.0392 | **0.98x** | **1.09x** | +| 16 | 0.0685 | 0.0605 | 0.0699 | **1.02x** | **1.15x** | +| 32 | 0.1324 | 0.1156 | 0.1316 | **0.99x** | **1.14x** | +| 64 | 0.2394 | 0.2020 | 0.2520 | **1.05x** | **1.25x** | +| 128 | 0.4578 | 0.3845 | 0.4923 | **1.08x** | **1.28x** | +| 256 | 0.8939 | 0.7465 | 0.9784 | **1.09x** | **1.31x** | + +### Accuracy (Output) + +| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel | +|---|----------------:|---------------:|----------------:|---------------:| +| 1 | 1.490e-08 | 5.830e-05 | 1.490e-08 | 5.830e-05 | +| 4 | 7.465e-09 | 3.427e-05 | 7.465e-09 | 3.427e-05 | +| 8 | 5.990e-08 | 2.741e-04 | 5.990e-08 | 2.741e-04 | +| 16 | 7.778e-08 | 5.482e-04 | 7.778e-08 | 5.482e-04 | +| 32 | 8.604e-08 | 5.482e-04 | 8.604e-08 | 5.482e-04 | +| 64 | 1.411e-07 | 1.953e-03 | 1.411e-07 | 1.953e-03 | +| 128 | 1.015e-07 | 8.741e-04 | 1.015e-07 | 8.741e-04 | +| 256 | 1.038e-07 | 6.510e-04 | 1.038e-07 | 6.510e-04 | + +### Accuracy (State) + +| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel | +|---|------------------:|-----------------:|------------------:|-----------------:| +| 1 | 3.888e-09 | 2.085e-07 | 3.888e-09 | 2.085e-07 | +| 4 | 4.209e-09 | 1.955e-07 | 4.209e-09 | 1.955e-07 | +| 8 | 4.347e-09 | 1.858e-07 | 4.347e-09 | 1.858e-07 | +| 16 | 4.411e-09 | 1.858e-07 | 4.411e-09 | 1.858e-07 | +| 32 | 4.382e-09 | 1.717e-07 | 4.382e-09 | 1.717e-07 | +| 64 | 4.389e-09 | 1.625e-07 | 4.389e-09 | 1.625e-07 | +| 128 | 4.400e-09 | 2.728e-07 | 4.400e-09 | 2.728e-07 | +| 256 | 4.430e-09 | 2.334e-07 | 4.430e-09 | 2.334e-07 | + +## Reproduce + +```bash +python benchmarks/bench_kda_decode.py +``` diff --git a/BENCHMARK_KDA_DECODE_H203E.md b/BENCHMARK_KDA_DECODE_H203E.md new file mode 100644 index 00000000..c9e56201 --- /dev/null +++ b/BENCHMARK_KDA_DECODE_H203E.md @@ -0,0 +1,145 @@ +# Benchmark Results - KDA Decode + +> Auto-generated by `benchmarks/bench_kda_decode.py` on 2026-04-18 16:59:00. + +> **GPU:** NVIDIA H20-3e | **CUDA:** 12.8 | **PyTorch:** 2.9.0+cu128 | **Python:** 3.12.3 + +> Decode setting: single-token (T=1), K=128, V=128. + +## Summary + +- v-last speedup (FLA/cuLA): avg=1.05x, min=0.90x, max=1.50x +- k-last speedup (FLA/cuLA): avg=1.26x, min=1.13x, max=1.56x +- Batch sizes: [1, 4, 8, 16, 32, 64, 128, 256] +- Q/K heads (H): [8, 32, 64] +- V heads (HV): 128 +- Timing params: warmup=30, rep=200, ncu_mode=False + +## KDA Decode (H=8, HV=128, K=128, V=128) + +### Performance + +| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | +|---|------------------:|------------------:|----------------:|---------------:|---------------:| +| 1 | 0.0322 | 0.0309 | 0.0480 | **1.49x** | **1.55x** | +| 4 | 0.0389 | 0.0333 | 0.0402 | **1.03x** | **1.20x** | +| 8 | 0.0734 | 0.0595 | 0.0703 | **0.96x** | **1.18x** | +| 16 | 0.1387 | 0.1104 | 0.1334 | **0.96x** | **1.21x** | +| 32 | 0.2864 | 0.2259 | 0.2575 | **0.90x** | **1.14x** | +| 64 | 0.5250 | 0.4255 | 0.5095 | **0.97x** | **1.20x** | +| 128 | 1.1560 | 0.8632 | 1.0765 | **0.93x** | **1.25x** | +| 256 | 2.1451 | 1.7678 | 2.0018 | **0.93x** | **1.13x** | + +### Accuracy (Output) + +| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel | +|---|----------------:|---------------:|----------------:|---------------:| +| 1 | 2.675e-09 | 8.798e-06 | 2.675e-09 | 8.798e-06 | +| 4 | 1.145e-08 | 4.542e-05 | 1.145e-08 | 4.542e-05 | +| 8 | 4.424e-08 | 3.634e-04 | 4.424e-08 | 3.634e-04 | +| 16 | 5.518e-08 | 3.634e-04 | 5.518e-08 | 3.634e-04 | +| 32 | 1.354e-07 | 1.106e-03 | 1.354e-07 | 1.106e-03 | +| 64 | 5.708e-08 | 4.980e-04 | 5.708e-08 | 4.980e-04 | +| 128 | 1.036e-07 | 7.530e-04 | 1.036e-07 | 7.530e-04 | +| 256 | 9.156e-08 | 8.446e-04 | 9.156e-08 | 8.446e-04 | + +### Accuracy (State) + +| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel | +|---|------------------:|-----------------:|------------------:|-----------------:| +| 1 | 4.346e-09 | 1.611e-07 | 4.346e-09 | 1.611e-07 | +| 4 | 3.985e-09 | 1.195e-07 | 3.985e-09 | 1.195e-07 | +| 8 | 4.313e-09 | 2.389e-07 | 4.313e-09 | 2.389e-07 | +| 16 | 4.269e-09 | 2.389e-07 | 4.269e-09 | 2.389e-07 | +| 32 | 4.229e-09 | 2.220e-07 | 4.229e-09 | 2.220e-07 | +| 64 | 4.275e-09 | 2.338e-07 | 4.275e-09 | 2.338e-07 | +| 128 | 4.385e-09 | 1.906e-07 | 4.385e-09 | 1.906e-07 | +| 256 | 4.397e-09 | 2.039e-07 | 4.397e-09 | 2.039e-07 | + +## KDA Decode (H=32, HV=128, K=128, V=128) + +### Performance + +| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | +|---|------------------:|------------------:|----------------:|---------------:|---------------:| +| 1 | 0.0326 | 0.0312 | 0.0483 | **1.48x** | **1.55x** | +| 4 | 0.0394 | 0.0334 | 0.0390 | **0.99x** | **1.17x** | +| 8 | 0.0762 | 0.0628 | 0.0708 | **0.93x** | **1.13x** | +| 16 | 0.1350 | 0.1096 | 0.1362 | **1.01x** | **1.24x** | +| 32 | 0.2870 | 0.2172 | 0.2613 | **0.91x** | **1.20x** | +| 64 | 0.5228 | 0.4124 | 0.5169 | **0.99x** | **1.25x** | +| 128 | 1.0221 | 0.8133 | 1.1053 | **1.08x** | **1.36x** | +| 256 | 2.0099 | 1.7715 | 2.0360 | **1.01x** | **1.15x** | + +### Accuracy (Output) + +| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel | +|---|----------------:|---------------:|----------------:|---------------:| +| 1 | 1.666e-08 | 5.008e-05 | 1.666e-08 | 5.008e-05 | +| 4 | 6.031e-08 | 4.006e-04 | 6.031e-08 | 4.006e-04 | +| 8 | 1.271e-07 | 6.410e-04 | 1.271e-07 | 6.410e-04 | +| 16 | 2.567e-07 | 2.475e-03 | 2.567e-07 | 2.475e-03 | +| 32 | 1.907e-07 | 1.923e-03 | 1.907e-07 | 1.923e-03 | +| 64 | 1.473e-07 | 1.603e-03 | 1.473e-07 | 1.603e-03 | +| 128 | 1.815e-07 | 1.712e-03 | 1.815e-07 | 1.712e-03 | +| 256 | 1.310e-07 | 1.462e-03 | 1.310e-07 | 1.462e-03 | + +### Accuracy (State) + +| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel | +|---|------------------:|-----------------:|------------------:|-----------------:| +| 1 | 3.809e-09 | 1.417e-07 | 3.809e-09 | 1.417e-07 | +| 4 | 4.141e-09 | 1.976e-07 | 4.141e-09 | 1.976e-07 | +| 8 | 4.246e-09 | 1.976e-07 | 4.246e-09 | 1.976e-07 | +| 16 | 4.294e-09 | 1.867e-07 | 4.294e-09 | 1.867e-07 | +| 32 | 4.329e-09 | 2.351e-07 | 4.329e-09 | 2.351e-07 | +| 64 | 4.357e-09 | 2.140e-07 | 4.357e-09 | 2.140e-07 | +| 128 | 4.400e-09 | 2.523e-07 | 4.400e-09 | 2.523e-07 | +| 256 | 4.415e-09 | 1.824e-07 | 4.415e-09 | 1.824e-07 | + +## KDA Decode (H=64, HV=128, K=128, V=128) + +### Performance + +| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | +|---|------------------:|------------------:|----------------:|---------------:|---------------:| +| 1 | 0.0327 | 0.0315 | 0.0491 | **1.50x** | **1.56x** | +| 4 | 0.0392 | 0.0332 | 0.0398 | **1.02x** | **1.20x** | +| 8 | 0.0727 | 0.0591 | 0.0713 | **0.98x** | **1.21x** | +| 16 | 0.1449 | 0.1119 | 0.1362 | **0.94x** | **1.22x** | +| 32 | 0.2591 | 0.2036 | 0.2631 | **1.02x** | **1.29x** | +| 64 | 0.5023 | 0.4048 | 0.5163 | **1.03x** | **1.28x** | +| 128 | 0.9907 | 0.8056 | 1.1256 | **1.14x** | **1.40x** | +| 256 | 1.9737 | 1.8047 | 2.0436 | **1.04x** | **1.13x** | + +### Accuracy (Output) + +| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel | +|---|----------------:|---------------:|----------------:|---------------:| +| 1 | 3.332e-08 | 1.149e-04 | 3.332e-08 | 1.149e-04 | +| 4 | 1.687e-08 | 6.975e-05 | 1.687e-08 | 6.975e-05 | +| 8 | 1.207e-08 | 6.975e-05 | 1.207e-08 | 6.975e-05 | +| 16 | 8.538e-08 | 5.580e-04 | 8.538e-08 | 5.580e-04 | +| 32 | 9.205e-08 | 5.580e-04 | 9.205e-08 | 5.580e-04 | +| 64 | 1.621e-07 | 1.866e-03 | 1.621e-07 | 1.866e-03 | +| 128 | 9.375e-08 | 6.510e-04 | 9.375e-08 | 6.510e-04 | +| 256 | 8.748e-08 | 6.250e-04 | 8.748e-08 | 6.250e-04 | + +### Accuracy (State) + +| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel | +|---|------------------:|-----------------:|------------------:|-----------------:| +| 1 | 3.870e-09 | 2.050e-07 | 3.870e-09 | 2.050e-07 | +| 4 | 4.201e-09 | 1.986e-07 | 4.201e-09 | 1.986e-07 | +| 8 | 4.346e-09 | 1.869e-07 | 4.346e-09 | 1.869e-07 | +| 16 | 4.414e-09 | 1.869e-07 | 4.414e-09 | 1.869e-07 | +| 32 | 4.350e-09 | 2.492e-07 | 4.350e-09 | 2.492e-07 | +| 64 | 4.384e-09 | 3.529e-07 | 4.384e-09 | 3.529e-07 | +| 128 | 4.416e-09 | 2.027e-07 | 4.416e-09 | 2.027e-07 | +| 256 | 4.432e-09 | 2.033e-07 | 4.432e-09 | 2.033e-07 | + +## Reproduce + +```bash +python benchmarks/bench_kda_decode.py +``` diff --git a/benchmarks/bench_kda_decode.py b/benchmarks/bench_kda_decode.py new file mode 100644 index 00000000..0658623c --- /dev/null +++ b/benchmarks/bench_kda_decode.py @@ -0,0 +1,462 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +bench_kda_decode.py — 3-way benchmark for KDA decode (single-token, T=1) + +Compares three routes: + 1. cuLA with v-last states (..., K, V) + 2. cuLA with k-last states (..., V, K) + 3. FLA with v-last states (..., K, V) + +Fairness note: + - All three routes use the fused decode entry point. + - State buffers are reset before each timed iteration. + - The reset copy is done outside the CUDA event window and is NOT counted. + +Usage: + python benchmarks/bench_kda_decode.py + python benchmarks/bench_kda_decode.py --batch-sizes 1 4 16 64 128 256 + python benchmarks/bench_kda_decode.py --Hs 8 32 64 + python benchmarks/bench_kda_decode.py --ncu + +Note: + - This benchmark is currently restricted to K=128 and V=128. + - By default it reports H=8, H=32, and H=64. +""" + +import argparse +import os +import pathlib +import platform +import re +import sys +from datetime import datetime + +os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) + +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.kda import fused_sigmoid_gating_delta_rule_update as cula_fused +from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as fla_fused + + +# ────────────────────────────────────────────────────────────────────── +# Timing utility +# ────────────────────────────────────────────────────────────────────── +def benchmark_fn(fn, *, setup_fn=None, warmup=30, rep=200): + """Benchmark using CUDA events. + + If provided, setup_fn runs before each iteration and is excluded from the + timing window. This is used to reset the mutable recurrent state fairly. + """ + for _ in range(warmup): + if setup_fn is not None: + setup_fn() + fn() + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + + for i in range(rep): + if setup_fn is not None: + setup_fn() + starts[i].record() + fn() + ends[i].record() + + torch.cuda.synchronize() + times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends)) + n = len(times) + if n <= 2: + return sum(times) / max(len(times), 1) + iqr = times[n // 4 : 3 * n // 4] + return sum(iqr) / len(iqr) + + +# ────────────────────────────────────────────────────────────────────── +# Input generation +# ────────────────────────────────────────────────────────────────────── +def make_inputs(N, H, HV, K, V, device="cuda", seed=42): + """Generate random inputs for KDA decode benchmark.""" + torch.manual_seed(seed) + q = torch.randn(N, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn(N, H, K, device=device, dtype=torch.bfloat16) + v = torch.randn(N, HV, V, device=device, dtype=torch.bfloat16) + a = (torch.randn(N, HV, K, device=device, dtype=torch.float32) * 0.1).to(torch.bfloat16) + b = torch.randn(N, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) * 2 + dt_bias = torch.randn(HV, K, device=device, dtype=torch.float32) * 0.1 + state = torch.randn(N, HV, V, K, device=device, dtype=torch.float32) * 0.01 + return q, k, v, a, b, A_log, dt_bias, state + + +# ────────────────────────────────────────────────────────────────────── +# Accuracy check +# ────────────────────────────────────────────────────────────────────── +def accuracy_stats(ref, out): + ref_f, out_f = ref.float(), out.float() + diff = (ref_f - out_f).abs() + rmse = diff.pow(2).mean().sqrt().item() + max_diff = diff.max().item() + denom = ref_f.abs().max().item() + rel_max = max_diff / denom if denom > 0 else 0.0 + return rmse, rel_max + + +def to_v_last_state(state: torch.Tensor, layout: str) -> torch.Tensor: + if layout == "kv": + return state + if layout == "vk": + return state.permute(0, 1, 3, 2).contiguous() + raise ValueError(f"Unsupported layout={layout}") + + +def normalize_gpu_type(gpu_name: str) -> str: + """Normalize GPU name to a compact token for report filename.""" + upper = gpu_name.upper() + tokens = re.findall(r"[A-Z0-9]+", upper) + ignore = {"NVIDIA", "GEFORCE", "TESLA", "GRAPHICS", "CORPORATION", "INC"} + tokens = [t for t in tokens if t not in ignore] + + # Handle SKUs like "H20-3e" -> "H203E". + for i in range(len(tokens) - 1): + left, right = tokens[i], tokens[i + 1] + if re.fullmatch(r"H\d+", left) and re.fullmatch(r"\d+[A-Z]+", right): + return f"{left}{right}" + + if "RTX" in tokens: + i = tokens.index("RTX") + if i + 1 < len(tokens) and tokens[i + 1].isdigit(): + return f"RTX_{tokens[i + 1]}" + + digit_tokens = [t for t in tokens if any(ch.isdigit() for ch in t)] + if digit_tokens: + return digit_tokens[-1] + + return "_".join(tokens) if tokens else "UNKNOWN_GPU" + + +def write_markdown_report(args, gpu_name: str, sections: list[tuple[int, int, list[dict]]], output_path: pathlib.Path): + """Write benchmark results into a markdown report.""" + now = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + py_ver = platform.python_version() + cuda_ver = torch.version.cuda or "unknown" + torch_ver = torch.__version__ + + speedups_v = [r["speedup_v_last"] for _, _, rows in sections for r in rows] + speedups_k = [r["speedup_k_last"] for _, _, rows in sections for r in rows] + + def summary(vals): + if not vals: + return "n/a" + return f"avg={sum(vals)/len(vals):.2f}x, min={min(vals):.2f}x, max={max(vals):.2f}x" + + lines = [] + lines.append("# Benchmark Results - KDA Decode") + lines.append("") + lines.append(f"> Auto-generated by `benchmarks/bench_kda_decode.py` on {now}.") + lines.append("") + lines.append(f"> **GPU:** {gpu_name} | **CUDA:** {cuda_ver} | **PyTorch:** {torch_ver} | **Python:** {py_ver}") + lines.append("") + lines.append("> Decode setting: single-token (T=1), K=128, V=128.") + lines.append("") + lines.append("## Summary") + lines.append("") + lines.append(f"- v-last speedup (FLA/cuLA): {summary(speedups_v)}") + lines.append(f"- k-last speedup (FLA/cuLA): {summary(speedups_k)}") + lines.append(f"- Batch sizes: {args.batch_sizes}") + lines.append(f"- Q/K heads (H): {args.Hs}") + lines.append(f"- V heads (HV): {args.HV}") + lines.append(f"- Timing params: warmup={args.warmup}, rep={args.rep}, ncu_mode={args.ncu}") + lines.append("") + + for h_dim, v_dim, results in sections: + lines.append(f"## KDA Decode (H={h_dim}, HV={args.HV}, K={args.K}, V={v_dim})") + lines.append("") + lines.append("### Performance") + lines.append("") + lines.append("| N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup |") + lines.append("|---|------------------:|------------------:|----------------:|---------------:|---------------:|") + for r in results: + lines.append( + f"| {r['N']} | {r['t_cula_v_last_ms']:.4f} | {r['t_cula_k_last_ms']:.4f} | {r['t_fla_v_last_ms']:.4f} | " + f"**{r['speedup_v_last']:.2f}x** | **{r['speedup_k_last']:.2f}x** |" + ) + lines.append("") + + lines.append("### Accuracy (Output)") + lines.append("") + lines.append("| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel |") + lines.append("|---|----------------:|---------------:|----------------:|---------------:|") + for r in results: + lines.append( + f"| {r['N']} | {r['out_v_last_rmse']:.3e} | {r['out_v_last_rel']:.3e} | " + f"{r['out_k_last_rmse']:.3e} | {r['out_k_last_rel']:.3e} |" + ) + lines.append("") + + lines.append("### Accuracy (State)") + lines.append("") + lines.append("| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel |") + lines.append("|---|------------------:|-----------------:|------------------:|-----------------:|") + for r in results: + lines.append( + f"| {r['N']} | {r['state_v_last_rmse']:.3e} | {r['state_v_last_rel']:.3e} | " + f"{r['state_k_last_rmse']:.3e} | {r['state_k_last_rel']:.3e} |" + ) + lines.append("") + + lines.append("## Reproduce") + lines.append("") + lines.append("```bash") + lines.append("python benchmarks/bench_kda_decode.py") + lines.append("```") + lines.append("") + + output_path.write_text("\n".join(lines), encoding="utf-8") + + +# ────────────────────────────────────────────────────────────────────── +# Run one config +# ────────────────────────────────────────────────────────────────────── +def run_config(N, H, HV, K, V, warmup, rep, ncu_mode): + device = "cuda" + scale = K**-0.5 + + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V, device) + + q_4d = q.unsqueeze(1).contiguous() + k_4d = k.unsqueeze(1).contiguous() + v_4d = v.unsqueeze(1).contiguous() + a_flat = a.reshape(N, 1, -1).contiguous() + b_3d = b.unsqueeze(1).contiguous() + indices = torch.arange(N, device=device, dtype=torch.int32) + + state_init_k_last = state.clone().contiguous() # (N, HV, V, K) + state_init_v_last = state_init_k_last.permute(0, 1, 3, 2).contiguous() # (N, HV, K, V) + + state_cula_v_last = state_init_v_last.clone() + state_cula_k_last = state_init_k_last.clone() + state_fla_v_last = state_init_v_last.clone() + + def call_cula_v_last(state_buf): + return cula_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="kv", + ) + + def call_cula_k_last(state_buf): + return cula_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="vk", + ) + + def call_fla_v_last(state_buf): + return fla_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + ) + + with torch.no_grad(): + o_cula_v_last = call_cula_v_last(state_cula_v_last) + o_cula_k_last = call_cula_k_last(state_cula_k_last) + o_fla_v_last = call_fla_v_last(state_fla_v_last) + + out_v_last_rmse, out_v_last_rel = accuracy_stats(o_fla_v_last, o_cula_v_last) + out_k_last_rmse, out_k_last_rel = accuracy_stats(o_fla_v_last, o_cula_k_last) + state_v_last_rmse, state_v_last_rel = accuracy_stats(state_fla_v_last, state_cula_v_last) + state_k_last_rmse, state_k_last_rel = accuracy_stats(state_fla_v_last, to_v_last_state(state_cula_k_last, "vk")) + + if ncu_mode: + w, r = 1, 1 + else: + w, r = warmup, rep + + state_bench_cula_v_last = state_init_v_last.clone() + state_bench_cula_k_last = state_init_k_last.clone() + state_bench_fla_v_last = state_init_v_last.clone() + + def setup_cula_v_last(): + state_bench_cula_v_last.copy_(state_init_v_last) + + def setup_cula_k_last(): + state_bench_cula_k_last.copy_(state_init_k_last) + + def setup_fla_v_last(): + state_bench_fla_v_last.copy_(state_init_v_last) + + with torch.no_grad(): + t_cula_v_last = benchmark_fn(lambda: call_cula_v_last(state_bench_cula_v_last), setup_fn=setup_cula_v_last, warmup=w, rep=r) + t_cula_k_last = benchmark_fn(lambda: call_cula_k_last(state_bench_cula_k_last), setup_fn=setup_cula_k_last, warmup=w, rep=r) + t_fla_v_last = benchmark_fn(lambda: call_fla_v_last(state_bench_fla_v_last), setup_fn=setup_fla_v_last, warmup=w, rep=r) + + return { + "N": N, + "H": H, + "HV": HV, + "K": K, + "V": V, + "t_cula_v_last_ms": t_cula_v_last, + "t_cula_k_last_ms": t_cula_k_last, + "t_fla_v_last_ms": t_fla_v_last, + "speedup_v_last": t_fla_v_last / t_cula_v_last if t_cula_v_last > 0 else float("inf"), + "speedup_k_last": t_fla_v_last / t_cula_k_last if t_cula_k_last > 0 else float("inf"), + "out_v_last_rmse": out_v_last_rmse, + "out_v_last_rel": out_v_last_rel, + "out_k_last_rmse": out_k_last_rmse, + "out_k_last_rel": out_k_last_rel, + "state_v_last_rmse": state_v_last_rmse, + "state_v_last_rel": state_v_last_rel, + "state_k_last_rmse": state_k_last_rmse, + "state_k_last_rel": state_k_last_rel, + } + + +# ────────────────────────────────────────────────────────────────────── +# Main +# ────────────────────────────────────────────────────────────────────── +def main(): + parser = argparse.ArgumentParser(description="Benchmark KDA decode: cuLA vs FLA") + parser.add_argument("--batch-sizes", nargs="+", type=int, default=[1, 4, 8, 16, 32, 64, 128, 256]) + parser.add_argument("--Hs", nargs="+", type=int, default=[8, 32, 64], help="Q/K head counts to benchmark") + parser.add_argument("--HV", type=int, default=128, help="Number of V heads (GVA)") + parser.add_argument("--K", type=int, default=128, help="Head dim K (only 128 is supported)") + parser.add_argument("--V", type=int, default=128, help="Head dim V (only 128 is supported)") + parser.add_argument("--warmup", type=int, default=30) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--ncu", action="store_true", help="NCU mode: warmup=1, rep=1") + args = parser.parse_args() + + if args.K != 128 or args.V != 128: + raise ValueError(f"bench_kda_decode.py currently only supports K=128 and V=128, got K={args.K}, V={args.V}") + + gpu_name = torch.cuda.get_device_name(0) + print(f"GPU: {gpu_name}") + print(f"Config: Hs={args.Hs}, HV={args.HV}, K={args.K}, V={args.V}") + print("Timing note: state reset uses copy_() before each timed iteration and is not counted.") + print("Accuracy reference: FLA v-last route.") + print() + + all_sections: list[tuple[int, int, list[dict]]] = [] + + def print_section(h_dim: int, v_dim: int): + print(f"Config: H={h_dim}, HV={args.HV}, K={args.K}, V={v_dim}") + hdr_perf = ( + f"{'N':>5} | {'cuLA v-last':>12} | {'cuLA k-last':>12} | {'FLA v-last':>12} | " + f"{'v-last spd':>10} | {'k-last spd':>10}" + ) + print(hdr_perf) + print("-" * len(hdr_perf)) + + results = [] + for N in args.batch_sizes: + res = run_config( + N, + h_dim, + args.HV, + args.K, + v_dim, + args.warmup, + args.rep, + args.ncu, + ) + results.append(res) + print( + f"{res['N']:5d} | {res['t_cula_v_last_ms']:12.4f} | {res['t_cula_k_last_ms']:12.4f} | " + f"{res['t_fla_v_last_ms']:12.4f} | {res['speedup_v_last']:9.2f}x | {res['speedup_k_last']:9.2f}x" + ) + + print() + hdr_out = ( + f"{'N':>5} | {'cuLA v out RMSE':>16} | {'rel':>10} | " + f"{'cuLA k out RMSE':>16} | {'rel':>10}" + ) + print(hdr_out) + print("-" * len(hdr_out)) + for res in results: + print( + f"{res['N']:5d} | {res['out_v_last_rmse']:16.3e} | {res['out_v_last_rel']:10.3e} | " + f"{res['out_k_last_rmse']:16.3e} | {res['out_k_last_rel']:10.3e}" + ) + + print() + hdr_state = ( + f"{'N':>5} | {'cuLA v state RMSE':>18} | {'rel':>10} | " + f"{'cuLA k state RMSE':>18} | {'rel':>10}" + ) + print(hdr_state) + print("-" * len(hdr_state)) + for res in results: + print( + f"{res['N']:5d} | {res['state_v_last_rmse']:18.3e} | {res['state_v_last_rel']:10.3e} | " + f"{res['state_k_last_rmse']:18.3e} | {res['state_k_last_rel']:10.3e}" + ) + + all_sections.append((h_dim, v_dim, results)) + + for h_dim in args.Hs: + print_section(h_dim, args.V) + print() + + gpu_type = normalize_gpu_type(gpu_name) + report_path = pathlib.Path(__file__).resolve().parent.parent / f"BENCHMARK_KDA_DECODE_{gpu_type}.md" + write_markdown_report(args, gpu_name, all_sections, report_path) + print(f"Markdown report written to: {report_path}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_la_decode_vs_fla.py b/benchmarks/bench_la_decode_vs_fla.py index 964555e4..471bc226 100644 --- a/benchmarks/bench_la_decode_vs_fla.py +++ b/benchmarks/bench_la_decode_vs_fla.py @@ -55,7 +55,7 @@ from fla.ops.common.fused_recurrent import fused_recurrent_fwd, fused_recurrent_fwd_kernel -from cula.lightning.la_decode import _get_compiled_kernel, linear_attention_decode +from cula.ops.la_decode import _get_compiled_kernel, linear_attention_decode from cula.utils import USE_FAST_MATH diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index f7ee65d2..01190143 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -14,10 +14,11 @@ from cula.kda.chunk import chunk_kda from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper -from cula.kda.kda_decode import kda_decode +from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode __all__ = [ "chunk_kda", "kda_decode", + "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", ] diff --git a/cula/kda/kda_decode.py b/cula/kda/kda_decode.py deleted file mode 100644 index b399f2ce..00000000 --- a/cula/kda/kda_decode.py +++ /dev/null @@ -1,1525 +0,0 @@ -"""CuTe DSL Fused Sigmoid Gating Delta Rule Kernel for KDA Decode. - -This version uses production / Triton-compatible VK state layout: - state.shape == (pool_size, HV, V, K) - -The kernel still computes on a logical (K, V) matrix in shared memory. Global -state loads/stores therefore explicitly map: - global(V, K) <-> shared(K, V) - -Notes: -- This is a correctness-first implementation for decode. -- It keeps the original small-batch / large-batch split. -- It preserves the previous PAD semantics: if pool_idx < 0 the block does not - load / update / write output or state, consistent with the earlier CuTe path. -""" - -import logging -from typing import Dict, Optional, Tuple - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute -import torch -from cutlass.cute.runtime import from_dlpack - -logger = logging.getLogger(__name__) - -_compiled_kernels: Dict[Tuple, object] = {} -_cu_seqlens_cache: Dict[Tuple, torch.Tensor] = {} - -TILE_K = 128 -TILE_V = 32 -TILE_V_PADDED = 36 -TILE_V_SMALL = 16 -TILE_V_SMALL_PADDED = 20 -NUM_STAGES = 2 -NUM_THREADS = 128 -NUM_BLOCKS_PER_STATE_SMALL = 8 -NUM_THREADS_LARGE = 256 -NUM_WARPS_LARGE = 8 -V_PER_WARP = 4 -ROWS_PER_ITER = 8 -NUM_K_ITERS = TILE_K // ROWS_PER_ITER -SMALL_BATCH_THRESHOLD = 32 - - -def _define_kernels(): - """Define CuTe DSL kernels for KDA normal and varlen decode modes.""" - - NUM_WARPS_SMALL = 4 - V_PER_WARP_SMALL = TILE_V_SMALL // NUM_WARPS_SMALL - ROWS_PER_ITER_SMALL = 32 // V_PER_WARP_SMALL - NUM_K_ITERS_SMALL = TILE_K // ROWS_PER_ITER_SMALL - - @cute.kernel - def kda_kernel_small_batch( - tiled_copy_load: cute.TiledCopy, - h0_source: cute.Tensor, - smem_layout_staged: cute.Layout, - num_v_tiles: cutlass.Constexpr[int], - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - o: cute.Tensor, - h0_indices: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - use_qk_l2norm: cutlass.Constexpr[bool], - ): - """Small batch KDA kernel for dense decode: q/k/v shapes (N, 1, ...).""" - del tiled_copy_load - tidx, _, _ = cute.arch.thread_idx() - in_warp_tid = tidx % 32 - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - block_idx, _, _ = cute.arch.block_idx() - - batch_idx = block_idx // NUM_BLOCKS_PER_STATE_SMALL - batch_inner = block_idx % NUM_BLOCKS_PER_STATE_SMALL - num_v_tiles_per_block = num_v_tiles // NUM_BLOCKS_PER_STATE_SMALL - start_v_tile = batch_inner * num_v_tiles_per_block - - i_n = batch_idx // HV - i_hv = batch_idx % HV - i_h = i_hv // (HV // H) - - pool_idx = h0_indices[i_n] - - if pool_idx >= 0: - k_local = in_warp_tid // V_PER_WARP_SMALL - v_local = in_warp_tid % V_PER_WARP_SMALL - v_base = warp_idx * V_PER_WARP_SMALL - v_idx = v_base + v_local - - smem = cutlass.utils.SmemAllocator() - sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) - smem_o_layout = cute.make_layout((TILE_V_SMALL,), stride=(1,)) - smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) - smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) - sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) - sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) - sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) - - if tidx < TILE_K: - sK[tidx] = cutlass.Float32(k[i_n, 0, i_h, tidx]) - sQ[tidx] = cutlass.Float32(q[i_n, 0, i_h, tidx]) - - r_A_log = cutlass.Float32(A_log[i_hv]) - r_exp_A = cute.exp(r_A_log) - if tidx < TILE_K: - r_a_k = cutlass.Float32(a[i_n, 0, i_hv, tidx]) - r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) - x = r_a_k + r_dt_bias_k - beta_x = softplus_beta * x - softplus_x = 0.0 - if beta_x <= softplus_threshold: - exp_beta_x = cute.exp(beta_x) - log_input = cutlass.Float32(1.0 + exp_beta_x) - log_result = cutlass.Float32(cute.log(log_input)) - softplus_x = cutlass.Float32( - (cutlass.Float32(1.0) / softplus_beta) * log_result - ) - else: - softplus_x = x - sG[tidx] = cute.exp(-r_exp_A * softplus_x) - - r_beta = 0.0 - if in_warp_tid == 0: - r_b = cutlass.Float32(b[i_n, 0, i_hv]) - r_beta = 1.0 / (1.0 + cute.exp(-r_b)) - r_beta = cute.arch.shuffle_sync(r_beta, 0) - - cute.arch.barrier() - - if use_qk_l2norm: - sum_q_partial = 0.0 - sum_k_partial = 0.0 - if tidx < TILE_K: - q_val = sQ[tidx] - k_val = sK[tidx] - sum_q_partial = q_val * q_val - sum_k_partial = k_val * k_val - - for offset in [16, 8, 4, 2, 1]: - sum_q_partial += cute.arch.shuffle_sync_bfly( - sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - sum_k_partial += cute.arch.shuffle_sync_bfly( - sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - - if in_warp_tid == 0: - smem_o[warp_idx] = sum_q_partial - smem_o[warp_idx + 4] = sum_k_partial - cute.arch.barrier() - - if warp_idx == 0: - local_sum_q = 0.0 - local_sum_k = 0.0 - if in_warp_tid < NUM_WARPS_SMALL: - local_sum_q = smem_o[in_warp_tid] - local_sum_k = smem_o[in_warp_tid + 4] - for offset in [2, 1]: - local_sum_q += cute.arch.shuffle_sync_bfly( - local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 - ) - local_sum_k += cute.arch.shuffle_sync_bfly( - local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 - ) - if in_warp_tid == 0: - smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) - smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) - cute.arch.barrier() - - inv_norm_q = smem_o[0] - inv_norm_k = smem_o[1] - - if tidx < TILE_K: - sK[tidx] = sK[tidx] * inv_norm_k - sQ[tidx] = sQ[tidx] * scale * inv_norm_q - cute.arch.barrier() - else: - if tidx < TILE_K: - sQ[tidx] = sQ[tidx] * scale - cute.arch.barrier() - - for v_tile_offset in range(num_v_tiles_per_block): - stage = v_tile_offset % NUM_STAGES - v_tile = start_v_tile + v_tile_offset - - for k_iter in range(NUM_K_ITERS_SMALL): - flat_idx = tidx + k_iter * NUM_THREADS - k_load = flat_idx // TILE_V_SMALL - v_load = flat_idx % TILE_V_SMALL - if k_load < TILE_K: - v_global_load = v_tile * TILE_V_SMALL + v_load - h_val = 0.0 - if v_global_load < v.shape[3]: - h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] - ) - sData[(k_load, v_load, stage)] = h_val - - cute.arch.barrier() - - v_global = v_tile * TILE_V_SMALL + v_idx - r_v = 0.0 - if v_global < v.shape[3]: - r_v = cutlass.Float32(v[i_n, 0, i_hv, v_global]) - - sum_hk = 0.0 - for k_iter in range(NUM_K_ITERS_SMALL): - k_base = k_iter * ROWS_PER_ITER_SMALL - k_idx = k_base + k_local - sum_hk += sData[(k_idx, v_idx, stage)] * sG[k_idx] * sK[k_idx] - - for offset in [4, 2, 1]: - sum_hk += cute.arch.shuffle_sync_bfly( - sum_hk, - offset=offset * V_PER_WARP_SMALL, - mask=-1, - mask_and_clamp=31, - ) - - v_new = (r_v - sum_hk) * r_beta - v_new = cute.arch.shuffle_sync(v_new, v_local) - - sum_hq = 0.0 - for k_iter in range(NUM_K_ITERS_SMALL): - k_base = k_iter * ROWS_PER_ITER_SMALL - k_idx = k_base + k_local - h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] - h_new = h_old + sK[k_idx] * v_new - sData[(k_idx, v_idx, stage)] = h_new - sum_hq += h_new * sQ[k_idx] - - for offset in [4, 2, 1]: - sum_hq += cute.arch.shuffle_sync_bfly( - sum_hq, - offset=offset * V_PER_WARP_SMALL, - mask=-1, - mask_and_clamp=31, - ) - - if k_local == 0 and v_global < v.shape[3]: - o[(i_n, 0, i_hv, v_global)] = cutlass.BFloat16(sum_hq) - - cute.arch.barrier() - - for k_iter in range(NUM_K_ITERS_SMALL): - flat_idx = tidx + k_iter * NUM_THREADS - k_write = flat_idx // TILE_V_SMALL - v_write = flat_idx % TILE_V_SMALL - if k_write < TILE_K: - v_global_write = v_tile * TILE_V_SMALL + v_write - if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( - sData[(k_write, v_write, stage)] - ) - - cute.arch.barrier() - - @cute.kernel - def kda_kernel_small_batch_varlen( - tiled_copy_load: cute.TiledCopy, - h0_source: cute.Tensor, - smem_layout_staged: cute.Layout, - num_v_tiles: cutlass.Constexpr[int], - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - o: cute.Tensor, - h0_indices: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - use_qk_l2norm: cutlass.Constexpr[bool], - ): - """Small batch KDA kernel for varlen decode: q/k/v shapes (1, N, ...).""" - del tiled_copy_load - tidx, _, _ = cute.arch.thread_idx() - in_warp_tid = tidx % 32 - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - block_idx, _, _ = cute.arch.block_idx() - - batch_idx = block_idx // NUM_BLOCKS_PER_STATE_SMALL - batch_inner = block_idx % NUM_BLOCKS_PER_STATE_SMALL - num_v_tiles_per_block = num_v_tiles // NUM_BLOCKS_PER_STATE_SMALL - start_v_tile = batch_inner * num_v_tiles_per_block - - i_n = batch_idx // HV - i_hv = batch_idx % HV - i_h = i_hv // (HV // H) - - pool_idx = h0_indices[i_n] - - if pool_idx >= 0: - k_local = in_warp_tid // V_PER_WARP_SMALL - v_local = in_warp_tid % V_PER_WARP_SMALL - v_base = warp_idx * V_PER_WARP_SMALL - v_idx = v_base + v_local - - smem = cutlass.utils.SmemAllocator() - sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) - smem_o_layout = cute.make_layout((TILE_V_SMALL,), stride=(1,)) - smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) - smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) - sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) - sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) - sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) - - if tidx < TILE_K: - sK[tidx] = cutlass.Float32(k[0, i_n, i_h, tidx]) - sQ[tidx] = cutlass.Float32(q[0, i_n, i_h, tidx]) - - r_A_log = cutlass.Float32(A_log[i_hv]) - r_exp_A = cute.exp(r_A_log) - if tidx < TILE_K: - r_a_k = cutlass.Float32(a[i_n, i_hv, tidx]) - r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) - x = r_a_k + r_dt_bias_k - beta_x = softplus_beta * x - softplus_x = 0.0 - if beta_x <= softplus_threshold: - exp_beta_x = cute.exp(beta_x) - log_input = cutlass.Float32(1.0 + exp_beta_x) - log_result = cutlass.Float32(cute.log(log_input)) - softplus_x = cutlass.Float32( - (cutlass.Float32(1.0) / softplus_beta) * log_result - ) - else: - softplus_x = x - sG[tidx] = cute.exp(-r_exp_A * softplus_x) - - r_beta = 0.0 - if in_warp_tid == 0: - r_b = cutlass.Float32(b[i_n, i_hv]) - r_beta = 1.0 / (1.0 + cute.exp(-r_b)) - r_beta = cute.arch.shuffle_sync(r_beta, 0) - - cute.arch.barrier() - - if use_qk_l2norm: - sum_q_partial = 0.0 - sum_k_partial = 0.0 - if tidx < TILE_K: - q_val = sQ[tidx] - k_val = sK[tidx] - sum_q_partial = q_val * q_val - sum_k_partial = k_val * k_val - - for offset in [16, 8, 4, 2, 1]: - sum_q_partial += cute.arch.shuffle_sync_bfly( - sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - sum_k_partial += cute.arch.shuffle_sync_bfly( - sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - - if in_warp_tid == 0: - smem_o[warp_idx] = sum_q_partial - smem_o[warp_idx + 4] = sum_k_partial - cute.arch.barrier() - - if warp_idx == 0: - local_sum_q = 0.0 - local_sum_k = 0.0 - if in_warp_tid < NUM_WARPS_SMALL: - local_sum_q = smem_o[in_warp_tid] - local_sum_k = smem_o[in_warp_tid + 4] - for offset in [2, 1]: - local_sum_q += cute.arch.shuffle_sync_bfly( - local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 - ) - local_sum_k += cute.arch.shuffle_sync_bfly( - local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 - ) - if in_warp_tid == 0: - smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) - smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) - cute.arch.barrier() - - inv_norm_q = smem_o[0] - inv_norm_k = smem_o[1] - - if tidx < TILE_K: - sK[tidx] = sK[tidx] * inv_norm_k - sQ[tidx] = sQ[tidx] * scale * inv_norm_q - cute.arch.barrier() - else: - if tidx < TILE_K: - sQ[tidx] = sQ[tidx] * scale - cute.arch.barrier() - - for v_tile_offset in range(num_v_tiles_per_block): - stage = v_tile_offset % NUM_STAGES - v_tile = start_v_tile + v_tile_offset - - for k_iter in range(NUM_K_ITERS_SMALL): - flat_idx = tidx + k_iter * NUM_THREADS - k_load = flat_idx // TILE_V_SMALL - v_load = flat_idx % TILE_V_SMALL - if k_load < TILE_K: - v_global_load = v_tile * TILE_V_SMALL + v_load - h_val = 0.0 - if v_global_load < v.shape[3]: - h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] - ) - sData[(k_load, v_load, stage)] = h_val - - cute.arch.barrier() - - v_global = v_tile * TILE_V_SMALL + v_idx - r_v = 0.0 - if v_global < v.shape[3]: - r_v = cutlass.Float32(v[0, i_n, i_hv, v_global]) - - sum_hk = 0.0 - for k_iter in range(NUM_K_ITERS_SMALL): - k_base = k_iter * ROWS_PER_ITER_SMALL - k_idx = k_base + k_local - sum_hk += sData[(k_idx, v_idx, stage)] * sG[k_idx] * sK[k_idx] - - for offset in [4, 2, 1]: - sum_hk += cute.arch.shuffle_sync_bfly( - sum_hk, - offset=offset * V_PER_WARP_SMALL, - mask=-1, - mask_and_clamp=31, - ) - - v_new = (r_v - sum_hk) * r_beta - v_new = cute.arch.shuffle_sync(v_new, v_local) - - sum_hq = 0.0 - for k_iter in range(NUM_K_ITERS_SMALL): - k_base = k_iter * ROWS_PER_ITER_SMALL - k_idx = k_base + k_local - h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] - h_new = h_old + sK[k_idx] * v_new - sData[(k_idx, v_idx, stage)] = h_new - sum_hq += h_new * sQ[k_idx] - - for offset in [4, 2, 1]: - sum_hq += cute.arch.shuffle_sync_bfly( - sum_hq, - offset=offset * V_PER_WARP_SMALL, - mask=-1, - mask_and_clamp=31, - ) - - if k_local == 0 and v_global < v.shape[3]: - o[(0, i_n, i_hv, v_global)] = cutlass.BFloat16(sum_hq) - - cute.arch.barrier() - - for k_iter in range(NUM_K_ITERS_SMALL): - flat_idx = tidx + k_iter * NUM_THREADS - k_write = flat_idx // TILE_V_SMALL - v_write = flat_idx % TILE_V_SMALL - if k_write < TILE_K: - v_global_write = v_tile * TILE_V_SMALL + v_write - if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( - sData[(k_write, v_write, stage)] - ) - - cute.arch.barrier() - - @cute.kernel - def kda_kernel_large_batch( - tiled_copy_load: cute.TiledCopy, - h0_source: cute.Tensor, - smem_layout_staged: cute.Layout, - num_v_tiles: cutlass.Constexpr[int], - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - o: cute.Tensor, - h0_indices: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - use_qk_l2norm: cutlass.Constexpr[bool], - ): - """Large batch KDA kernel for dense decode: q/k/v shapes (N, 1, ...).""" - del tiled_copy_load - tidx, _, _ = cute.arch.thread_idx() - in_warp_tid = tidx % 32 - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - batch_idx, _, _ = cute.arch.block_idx() - - i_n = batch_idx // HV - i_hv = batch_idx % HV - i_h = i_hv // (HV // H) - - pool_idx = h0_indices[i_n] - - if pool_idx >= 0: - k_local = in_warp_tid // V_PER_WARP - v_local = in_warp_tid % V_PER_WARP - v_base = warp_idx * V_PER_WARP - v_idx = v_base + v_local - - smem = cutlass.utils.SmemAllocator() - sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) - smem_o_layout = cute.make_layout((TILE_V,), stride=(1,)) - smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) - smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) - sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) - sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) - sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) - - if tidx < TILE_K: - sK[tidx] = cutlass.Float32(k[i_n, 0, i_h, tidx]) - sQ[tidx] = cutlass.Float32(q[i_n, 0, i_h, tidx]) - - r_A_log = cutlass.Float32(A_log[i_hv]) - r_exp_A = cute.exp(r_A_log) - if tidx < TILE_K: - r_a_k = cutlass.Float32(a[i_n, 0, i_hv, tidx]) - r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) - x = r_a_k + r_dt_bias_k - beta_x = softplus_beta * x - softplus_x = 0.0 - if beta_x <= softplus_threshold: - exp_beta_x = cute.exp(beta_x) - log_input = cutlass.Float32(1.0 + exp_beta_x) - log_result = cutlass.Float32(cute.log(log_input)) - softplus_x = cutlass.Float32( - (cutlass.Float32(1.0) / softplus_beta) * log_result - ) - else: - softplus_x = x - sG[tidx] = cute.exp(-r_exp_A * softplus_x) - - r_beta = 0.0 - if in_warp_tid == 0: - r_b = cutlass.Float32(b[i_n, 0, i_hv]) - r_beta = 1.0 / (1.0 + cute.exp(-r_b)) - r_beta = cute.arch.shuffle_sync(r_beta, 0) - - cute.arch.barrier() - - if use_qk_l2norm: - sum_q_partial = 0.0 - sum_k_partial = 0.0 - if tidx < TILE_K: - q_val = sQ[tidx] - k_val = sK[tidx] - sum_q_partial = q_val * q_val - sum_k_partial = k_val * k_val - - for offset in [16, 8, 4, 2, 1]: - sum_q_partial += cute.arch.shuffle_sync_bfly( - sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - sum_k_partial += cute.arch.shuffle_sync_bfly( - sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - - if in_warp_tid == 0: - smem_o[warp_idx] = sum_q_partial - smem_o[warp_idx + 8] = sum_k_partial - cute.arch.barrier() - - if warp_idx == 0: - local_sum_q = 0.0 - local_sum_k = 0.0 - if in_warp_tid < NUM_WARPS_LARGE: - local_sum_q = smem_o[in_warp_tid] - local_sum_k = smem_o[in_warp_tid + 8] - for offset in [4, 2, 1]: - local_sum_q += cute.arch.shuffle_sync_bfly( - local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 - ) - local_sum_k += cute.arch.shuffle_sync_bfly( - local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 - ) - if in_warp_tid == 0: - smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) - smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) - cute.arch.barrier() - - inv_norm_q = smem_o[0] - inv_norm_k = smem_o[1] - - if tidx < TILE_K: - sK[tidx] = sK[tidx] * inv_norm_k - sQ[tidx] = sQ[tidx] * scale * inv_norm_q - cute.arch.barrier() - else: - if tidx < TILE_K: - sQ[tidx] = sQ[tidx] * scale - cute.arch.barrier() - - for v_tile in range(num_v_tiles): - stage = v_tile % NUM_STAGES - - for k_iter in range(NUM_K_ITERS): - flat_idx = tidx + k_iter * NUM_THREADS_LARGE - k_load = flat_idx // TILE_V - v_load = flat_idx % TILE_V - if k_load < TILE_K: - v_global_load = v_tile * TILE_V + v_load - h_val = 0.0 - if v_global_load < v.shape[3]: - h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] - ) - sData[(k_load, v_load, stage)] = h_val - - cute.arch.barrier() - - v_global = v_tile * TILE_V + v_idx - r_v = 0.0 - if v_global < v.shape[3]: - r_v = cutlass.Float32(v[i_n, 0, i_hv, v_global]) - - sum_hk = 0.0 - for k_iter in range(NUM_K_ITERS): - k_base = k_iter * ROWS_PER_ITER - k_idx = k_base + k_local - sum_hk += sData[(k_idx, v_idx, stage)] * sG[k_idx] * sK[k_idx] - - for offset in [4, 2, 1]: - sum_hk += cute.arch.shuffle_sync_bfly( - sum_hk, - offset=offset * V_PER_WARP, - mask=-1, - mask_and_clamp=31, - ) - - v_new = (r_v - sum_hk) * r_beta - v_new = cute.arch.shuffle_sync(v_new, v_local) - - sum_hq = 0.0 - for k_iter in range(NUM_K_ITERS): - k_base = k_iter * ROWS_PER_ITER - k_idx = k_base + k_local - h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] - h_new = h_old + sK[k_idx] * v_new - sData[(k_idx, v_idx, stage)] = h_new - sum_hq += h_new * sQ[k_idx] - - for offset in [4, 2, 1]: - sum_hq += cute.arch.shuffle_sync_bfly( - sum_hq, - offset=offset * V_PER_WARP, - mask=-1, - mask_and_clamp=31, - ) - - if k_local == 0 and v_global < v.shape[3]: - o[(i_n, 0, i_hv, v_global)] = cutlass.BFloat16(sum_hq) - - cute.arch.barrier() - - for k_iter in range(NUM_K_ITERS): - flat_idx = tidx + k_iter * NUM_THREADS_LARGE - k_write = flat_idx // TILE_V - v_write = flat_idx % TILE_V - if k_write < TILE_K: - v_global_write = v_tile * TILE_V + v_write - if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( - sData[(k_write, v_write, stage)] - ) - - cute.arch.barrier() - - @cute.kernel - def kda_kernel_large_batch_varlen( - tiled_copy_load: cute.TiledCopy, - h0_source: cute.Tensor, - smem_layout_staged: cute.Layout, - num_v_tiles: cutlass.Constexpr[int], - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - o: cute.Tensor, - h0_indices: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - use_qk_l2norm: cutlass.Constexpr[bool], - ): - """Large batch KDA kernel for varlen decode: q/k/v shapes (1, N, ...).""" - del tiled_copy_load - tidx, _, _ = cute.arch.thread_idx() - in_warp_tid = tidx % 32 - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - batch_idx, _, _ = cute.arch.block_idx() - - i_n = batch_idx // HV - i_hv = batch_idx % HV - i_h = i_hv // (HV // H) - - pool_idx = h0_indices[i_n] - - if pool_idx >= 0: - k_local = in_warp_tid // V_PER_WARP - v_local = in_warp_tid % V_PER_WARP - v_base = warp_idx * V_PER_WARP - v_idx = v_base + v_local - - smem = cutlass.utils.SmemAllocator() - sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) - smem_o_layout = cute.make_layout((TILE_V,), stride=(1,)) - smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) - smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) - smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) - sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) - sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) - sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) - - if tidx < TILE_K: - sK[tidx] = cutlass.Float32(k[0, i_n, i_h, tidx]) - sQ[tidx] = cutlass.Float32(q[0, i_n, i_h, tidx]) - - r_A_log = cutlass.Float32(A_log[i_hv]) - r_exp_A = cute.exp(r_A_log) - if tidx < TILE_K: - r_a_k = cutlass.Float32(a[i_n, i_hv, tidx]) - r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) - x = r_a_k + r_dt_bias_k - beta_x = softplus_beta * x - softplus_x = 0.0 - if beta_x <= softplus_threshold: - exp_beta_x = cute.exp(beta_x) - log_input = cutlass.Float32(1.0 + exp_beta_x) - log_result = cutlass.Float32(cute.log(log_input)) - softplus_x = cutlass.Float32( - (cutlass.Float32(1.0) / softplus_beta) * log_result - ) - else: - softplus_x = x - sG[tidx] = cute.exp(-r_exp_A * softplus_x) - - r_beta = 0.0 - if in_warp_tid == 0: - r_b = cutlass.Float32(b[i_n, i_hv]) - r_beta = 1.0 / (1.0 + cute.exp(-r_b)) - r_beta = cute.arch.shuffle_sync(r_beta, 0) - - cute.arch.barrier() - - if use_qk_l2norm: - sum_q_partial = 0.0 - sum_k_partial = 0.0 - if tidx < TILE_K: - q_val = sQ[tidx] - k_val = sK[tidx] - sum_q_partial = q_val * q_val - sum_k_partial = k_val * k_val - - for offset in [16, 8, 4, 2, 1]: - sum_q_partial += cute.arch.shuffle_sync_bfly( - sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - sum_k_partial += cute.arch.shuffle_sync_bfly( - sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 - ) - - if in_warp_tid == 0: - smem_o[warp_idx] = sum_q_partial - smem_o[warp_idx + 8] = sum_k_partial - cute.arch.barrier() - - if warp_idx == 0: - local_sum_q = 0.0 - local_sum_k = 0.0 - if in_warp_tid < NUM_WARPS_LARGE: - local_sum_q = smem_o[in_warp_tid] - local_sum_k = smem_o[in_warp_tid + 8] - for offset in [4, 2, 1]: - local_sum_q += cute.arch.shuffle_sync_bfly( - local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 - ) - local_sum_k += cute.arch.shuffle_sync_bfly( - local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 - ) - if in_warp_tid == 0: - smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) - smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) - cute.arch.barrier() - - inv_norm_q = smem_o[0] - inv_norm_k = smem_o[1] - - if tidx < TILE_K: - sK[tidx] = sK[tidx] * inv_norm_k - sQ[tidx] = sQ[tidx] * scale * inv_norm_q - cute.arch.barrier() - else: - if tidx < TILE_K: - sQ[tidx] = sQ[tidx] * scale - cute.arch.barrier() - - for v_tile in range(num_v_tiles): - stage = v_tile % NUM_STAGES - - for k_iter in range(NUM_K_ITERS): - flat_idx = tidx + k_iter * NUM_THREADS_LARGE - k_load = flat_idx // TILE_V - v_load = flat_idx % TILE_V - if k_load < TILE_K: - v_global_load = v_tile * TILE_V + v_load - h_val = 0.0 - if v_global_load < v.shape[3]: - h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] - ) - sData[(k_load, v_load, stage)] = h_val - - cute.arch.barrier() - - v_global = v_tile * TILE_V + v_idx - r_v = 0.0 - if v_global < v.shape[3]: - r_v = cutlass.Float32(v[0, i_n, i_hv, v_global]) - - sum_hk = 0.0 - for k_iter in range(NUM_K_ITERS): - k_base = k_iter * ROWS_PER_ITER - k_idx = k_base + k_local - sum_hk += sData[(k_idx, v_idx, stage)] * sG[k_idx] * sK[k_idx] - - for offset in [4, 2, 1]: - sum_hk += cute.arch.shuffle_sync_bfly( - sum_hk, - offset=offset * V_PER_WARP, - mask=-1, - mask_and_clamp=31, - ) - - v_new = (r_v - sum_hk) * r_beta - v_new = cute.arch.shuffle_sync(v_new, v_local) - - sum_hq = 0.0 - for k_iter in range(NUM_K_ITERS): - k_base = k_iter * ROWS_PER_ITER - k_idx = k_base + k_local - h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] - h_new = h_old + sK[k_idx] * v_new - sData[(k_idx, v_idx, stage)] = h_new - sum_hq += h_new * sQ[k_idx] - - for offset in [4, 2, 1]: - sum_hq += cute.arch.shuffle_sync_bfly( - sum_hq, - offset=offset * V_PER_WARP, - mask=-1, - mask_and_clamp=31, - ) - - if k_local == 0 and v_global < v.shape[3]: - o[(0, i_n, i_hv, v_global)] = cutlass.BFloat16(sum_hq) - - cute.arch.barrier() - - for k_iter in range(NUM_K_ITERS): - flat_idx = tidx + k_iter * NUM_THREADS_LARGE - k_write = flat_idx // TILE_V - v_write = flat_idx % TILE_V - if k_write < TILE_K: - v_global_write = v_tile * TILE_V + v_write - if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( - sData[(k_write, v_write, stage)] - ) - - cute.arch.barrier() - - return ( - kda_kernel_small_batch, - kda_kernel_small_batch_varlen, - kda_kernel_large_batch, - kda_kernel_large_batch_varlen, - ) - - -def _create_jit_functions(): - """Create JIT-compiled launcher functions for all KDA kernel variants.""" - - kda_small, kda_small_varlen, kda_large, kda_large_varlen = _define_kernels() - - @cute.jit - def run_small_batch( - cu_seqlens: cute.Tensor, - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - h0_source: cute.Tensor, - h0_indices: cute.Tensor, - o: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - B: cutlass.Constexpr[int], - T: cutlass.Constexpr[int], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - K: cutlass.Constexpr[int], - V: cutlass.Constexpr[int], - use_initial_state: cutlass.Constexpr[bool], - use_qk_l2norm: cutlass.Constexpr[bool], - stream: cuda.CUstream, - ): - del cu_seqlens, B, T, K, use_initial_state - _, hv_dim, v_dim, _ = h0_source.layout.shape - n_indices = h0_indices.layout.shape[0] - batch_size = n_indices * hv_dim - - num_v_tiles_small = cute.ceil_div(v_dim, TILE_V_SMALL) - smem_layout_small = cute.make_layout( - (TILE_K, TILE_V_SMALL, NUM_STAGES), - stride=(TILE_V_SMALL_PADDED, 1, TILE_K * TILE_V_SMALL_PADDED), - ) - smem_bytes_small = ( - 4 * TILE_K * TILE_V_SMALL_PADDED * NUM_STAGES - + 4 * TILE_V_SMALL - + 4 * TILE_K * 2 - + 4 * TILE_K - + 64 - ) - - kda_small( - None, - h0_source, - smem_layout_small, - num_v_tiles_small, - q, - k, - v, - a, - b, - A_log, - dt_bias, - o, - h0_indices, - softplus_beta, - softplus_threshold, - scale, - H, - HV, - use_qk_l2norm, - ).launch( - grid=(batch_size * NUM_BLOCKS_PER_STATE_SMALL, 1, 1), - block=[NUM_THREADS, 1, 1], - smem=smem_bytes_small, - stream=stream, - ) - - @cute.jit - def run_small_batch_varlen( - cu_seqlens: cute.Tensor, - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - h0_source: cute.Tensor, - h0_indices: cute.Tensor, - o: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - B: cutlass.Constexpr[int], - T: cutlass.Constexpr[int], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - K: cutlass.Constexpr[int], - V: cutlass.Constexpr[int], - use_initial_state: cutlass.Constexpr[bool], - use_qk_l2norm: cutlass.Constexpr[bool], - stream: cuda.CUstream, - ): - del cu_seqlens, B, T, K, use_initial_state - _, hv_dim, v_dim, _ = h0_source.layout.shape - n_indices = h0_indices.layout.shape[0] - batch_size = n_indices * hv_dim - - num_v_tiles_small = cute.ceil_div(v_dim, TILE_V_SMALL) - smem_layout_small = cute.make_layout( - (TILE_K, TILE_V_SMALL, NUM_STAGES), - stride=(TILE_V_SMALL_PADDED, 1, TILE_K * TILE_V_SMALL_PADDED), - ) - smem_bytes_small = ( - 4 * TILE_K * TILE_V_SMALL_PADDED * NUM_STAGES - + 4 * TILE_V_SMALL - + 4 * TILE_K * 2 - + 4 * TILE_K - + 64 - ) - - kda_small_varlen( - None, - h0_source, - smem_layout_small, - num_v_tiles_small, - q, - k, - v, - a, - b, - A_log, - dt_bias, - o, - h0_indices, - softplus_beta, - softplus_threshold, - scale, - H, - HV, - use_qk_l2norm, - ).launch( - grid=(batch_size * NUM_BLOCKS_PER_STATE_SMALL, 1, 1), - block=[NUM_THREADS, 1, 1], - smem=smem_bytes_small, - stream=stream, - ) - - @cute.jit - def run_large_batch( - cu_seqlens: cute.Tensor, - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - h0_source: cute.Tensor, - h0_indices: cute.Tensor, - o: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - B: cutlass.Constexpr[int], - T: cutlass.Constexpr[int], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - K: cutlass.Constexpr[int], - V: cutlass.Constexpr[int], - use_initial_state: cutlass.Constexpr[bool], - use_qk_l2norm: cutlass.Constexpr[bool], - stream: cuda.CUstream, - ): - del cu_seqlens, B, T, K, use_initial_state - _, hv_dim, v_dim, _ = h0_source.layout.shape - n_indices = h0_indices.layout.shape[0] - batch_size = n_indices * hv_dim - - num_v_tiles = cute.ceil_div(v_dim, TILE_V) - smem_layout = cute.make_layout( - (TILE_K, TILE_V, NUM_STAGES), - stride=(TILE_V_PADDED, 1, TILE_K * TILE_V_PADDED), - ) - smem_bytes = ( - 4 * TILE_K * TILE_V_PADDED * NUM_STAGES - + 4 * TILE_V - + 4 * TILE_K * 2 - + 4 * TILE_K - + 64 - ) - - kda_large( - None, - h0_source, - smem_layout, - num_v_tiles, - q, - k, - v, - a, - b, - A_log, - dt_bias, - o, - h0_indices, - softplus_beta, - softplus_threshold, - scale, - H, - HV, - use_qk_l2norm, - ).launch( - grid=(batch_size, 1, 1), - block=[NUM_THREADS_LARGE, 1, 1], - smem=smem_bytes, - stream=stream, - ) - - @cute.jit - def run_large_batch_varlen( - cu_seqlens: cute.Tensor, - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - a: cute.Tensor, - b: cute.Tensor, - A_log: cute.Tensor, - dt_bias: cute.Tensor, - h0_source: cute.Tensor, - h0_indices: cute.Tensor, - o: cute.Tensor, - softplus_beta: cutlass.Constexpr[float], - softplus_threshold: cutlass.Constexpr[float], - scale: cutlass.Constexpr[float], - B: cutlass.Constexpr[int], - T: cutlass.Constexpr[int], - H: cutlass.Constexpr[int], - HV: cutlass.Constexpr[int], - K: cutlass.Constexpr[int], - V: cutlass.Constexpr[int], - use_initial_state: cutlass.Constexpr[bool], - use_qk_l2norm: cutlass.Constexpr[bool], - stream: cuda.CUstream, - ): - del cu_seqlens, B, T, K, use_initial_state - _, hv_dim, v_dim, _ = h0_source.layout.shape - n_indices = h0_indices.layout.shape[0] - batch_size = n_indices * hv_dim - - num_v_tiles = cute.ceil_div(v_dim, TILE_V) - smem_layout = cute.make_layout( - (TILE_K, TILE_V, NUM_STAGES), - stride=(TILE_V_PADDED, 1, TILE_K * TILE_V_PADDED), - ) - smem_bytes = ( - 4 * TILE_K * TILE_V_PADDED * NUM_STAGES - + 4 * TILE_V - + 4 * TILE_K * 2 - + 4 * TILE_K - + 64 - ) - - kda_large_varlen( - None, - h0_source, - smem_layout, - num_v_tiles, - q, - k, - v, - a, - b, - A_log, - dt_bias, - o, - h0_indices, - softplus_beta, - softplus_threshold, - scale, - H, - HV, - use_qk_l2norm, - ).launch( - grid=(batch_size, 1, 1), - block=[NUM_THREADS_LARGE, 1, 1], - smem=smem_bytes, - stream=stream, - ) - - return ( - run_small_batch, - run_small_batch_varlen, - run_large_batch, - run_large_batch_varlen, - ) - - -_jit_functions = None - - -def _get_jit_functions(): - global _jit_functions - if _jit_functions is None: - _jit_functions = _create_jit_functions() - return _jit_functions - - -def _get_compiled_kernel( - N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode, - scale, use_qk_l2norm, softplus_beta, softplus_threshold, -): - """Get or compile the KDA kernel for given dimensions and Constexpr values.""" - global _compiled_kernels - - key = ( - N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode, - scale, use_qk_l2norm, softplus_beta, softplus_threshold, - ) - if key in _compiled_kernels: - return _compiled_kernels[key] - - cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda") - - if is_varlen_decode: - q = torch.zeros(1, N, H, K, dtype=torch.bfloat16, device="cuda") - k = torch.zeros(1, N, H, K, dtype=torch.bfloat16, device="cuda") - v = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") - a = torch.zeros(N, HV, K, dtype=torch.bfloat16, device="cuda") - b = torch.zeros(N, HV, dtype=torch.bfloat16, device="cuda") - o = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") - else: - q = torch.zeros(N, 1, H, K, dtype=torch.bfloat16, device="cuda") - k = torch.zeros(N, 1, H, K, dtype=torch.bfloat16, device="cuda") - v = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") - a = torch.zeros(N, 1, HV, K, dtype=torch.bfloat16, device="cuda") - b = torch.zeros(N, 1, HV, dtype=torch.bfloat16, device="cuda") - o = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") - - A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") - dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") - h0_source = torch.zeros(pool_size, HV, V, K, dtype=torch.float32, device="cuda") - h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") - - cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16) - q_tensor = from_dlpack(q, assumed_align=16) - k_tensor = from_dlpack(k, assumed_align=16) - v_tensor = from_dlpack(v, assumed_align=16) - a_tensor = from_dlpack(a, assumed_align=16) - b_tensor = from_dlpack(b, assumed_align=16) - A_log_tensor = from_dlpack(A_log, assumed_align=16) - dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) - h0_source_tensor = from_dlpack(h0_source, assumed_align=16) - h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16) - o_tensor = from_dlpack(o, assumed_align=16) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - run_small, run_small_varlen, run_large, run_large_varlen = _get_jit_functions() - if use_small_batch: - kernel_func = run_small_varlen if is_varlen_decode else run_small - else: - kernel_func = run_large_varlen if is_varlen_decode else run_large - - compiled_kernel = cute.compile( - kernel_func, - cu_seqlens_tensor, - q_tensor, - k_tensor, - v_tensor, - a_tensor, - b_tensor, - A_log_tensor, - dt_bias_tensor, - h0_source_tensor, - h0_indices_tensor, - o_tensor, - softplus_beta=softplus_beta, - softplus_threshold=softplus_threshold, - scale=scale, - B=1 if is_varlen_decode else N, - T=N if is_varlen_decode else 1, - H=H, - K=K, - V=V, - HV=HV, - use_initial_state=True, - use_qk_l2norm=use_qk_l2norm, - stream=stream, - ) - - _compiled_kernels[key] = compiled_kernel - logger.info( - "CuTe DSL KDA kernel compiled: " - f"N={N}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, " - f"small_batch={use_small_batch}, varlen={is_varlen_decode}" - ) - return compiled_kernel - - -def _normalize_A_log(A_log: torch.Tensor, HV: int) -> torch.Tensor: - if A_log.numel() != HV: - raise ValueError(f"Unexpected A_log shape: {A_log.shape}; expected numel={HV}") - return A_log.reshape(HV).contiguous() - - -def _normalize_dt_bias(dt_bias: torch.Tensor, HV: int, K: int) -> torch.Tensor: - if dt_bias.numel() != HV * K: - raise ValueError( - f"Unexpected dt_bias shape: {dt_bias.shape}; expected numel={HV * K}" - ) - return dt_bias.reshape(HV, K).contiguous() - - -def _normalize_kda_a(a, *, is_varlen_decode, N, HV, K): - """Normalize `a` to match the compile-time shape expected by the kernel. - - varlen kernel compiled shape: (N, HV, K) -- 3D - dense kernel compiled shape: (N, 1, HV, K) -- 4D - """ - if is_varlen_decode: - # Target: (N, HV, K) -- 3D - if a.dim() == 2 and a.shape == (N, HV * K): - return a.view(N, HV, K) - if a.dim() == 3 and a.shape == (N, HV, K): - return a # already correct - if a.dim() == 4 and a.shape == (1, N, HV, K): - return a.squeeze(0) # remove leading dim - raise ValueError(f"Unexpected a shape for varlen: {a.shape}") - else: - # Target: (N, 1, HV, K) -- 4D - if a.dim() == 2 and a.shape == (N, HV * K): - return a.view(N, 1, HV, K) - if a.dim() == 3 and a.shape == (N, HV, K): - return a.unsqueeze(1) - if a.dim() == 4 and a.shape == (N, 1, HV, K): - return a - raise ValueError(f"Unexpected a shape for dense: {a.shape}") - - -def kda_decode( - A_log: torch.Tensor, - dt_bias: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - initial_state_source: torch.Tensor, - initial_state_indices: torch.Tensor, - cu_seqlens: Optional[torch.Tensor] = None, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = True, - softplus_beta: float = 1.0, - softplus_threshold: float = 20.0, -) -> torch.Tensor: - """CuTe DSL implementation of fused sigmoid gating KDA update. - - State layout contract: - initial_state_source.shape == (pool_size, HV, V, K) - - Dense decode: - q/k: (N, 1, H, K) - v: (N, 1, HV, V) - a: (N, 1, HV, K) - b: (N, 1, HV) - - Varlen decode: - q/k: (1, N, H, K) - v: (1, N, HV, V) - a: (N, HV, K) or (1, N, HV, K) - b: (N, HV) or (1, N, HV) - """ - - A_log = A_log.contiguous() - - B_q, T_q, H, K = q.shape - HV = v.shape[2] - V = v.shape[3] - N = initial_state_indices.shape[0] - - assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}" - assert ( - V % TILE_V_SMALL == 0 - ), f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" - assert ( - V % TILE_V == 0 - ), f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}" - assert (V // TILE_V_SMALL) % NUM_BLOCKS_PER_STATE_SMALL == 0, ( - "Small-batch KDA kernel requires num_v_tiles_small divisible by " - f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}" - ) - - is_varlen_decode = B_q == 1 and T_q == N and N > 1 - if scale is None: - scale = K**-0.5 - else: - assert scale > 0, f"scale must be positive, got {scale}" - - use_small_batch = N < SMALL_BATCH_THRESHOLD - - if initial_state_source.dim() == 1: - pool_size = initial_state_source.numel() // (HV * V * K) - h0_source = initial_state_source.view(pool_size, HV, V, K) - elif initial_state_source.dim() == 4: - pool_size = initial_state_source.shape[0] - h0_source = initial_state_source - else: - raise ValueError( - f"Unexpected initial_state_source shape: {initial_state_source.shape}" - ) - - a = _normalize_kda_a(a, is_varlen_decode=is_varlen_decode, N=N, HV=HV, K=K) - - if is_varlen_decode: - # varlen b compiled: (N, HV) -- 2D - if b.dim() == 3: - b = b.squeeze(0) # (1, N, HV) -> (N, HV) - # b should be 2D (N, HV) - o = q.new_empty(1, N, HV, V, dtype=torch.bfloat16) - else: - # dense b compiled: (N, 1, HV) -- 3D - if b.dim() == 2: - b = b.unsqueeze(1) - # b should be 3D (N, 1, HV) - o = q.new_empty(N, 1, HV, V, dtype=torch.bfloat16) - - q, k, v, a, b = [t.contiguous() for t in (q, k, v, a, b)] - dt_bias = dt_bias.contiguous() - - global _cu_seqlens_cache - if cu_seqlens is not None: - cu_seqlens_to_use = cu_seqlens - else: - cache_key = (N, str(q.device)) - if cache_key not in _cu_seqlens_cache: - _cu_seqlens_cache[cache_key] = torch.arange( - N + 1, dtype=torch.int32, device=q.device - ) - cu_seqlens_to_use = _cu_seqlens_cache[cache_key] - - A_log = _normalize_A_log(A_log, HV) - dt_bias = _normalize_dt_bias(dt_bias, HV, K) - - h0_source = h0_source.contiguous() - - initial_state_indices = initial_state_indices.contiguous() - if cu_seqlens is not None: - cu_seqlens = cu_seqlens.contiguous() - - cu_seqlens_tensor = from_dlpack( - cu_seqlens_to_use.detach(), assumed_align=16 - ).mark_layout_dynamic(leading_dim=0) - q_tensor = from_dlpack(q.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=q.ndim - 1 - ) - k_tensor = from_dlpack(k.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=k.ndim - 1 - ) - v_tensor = from_dlpack(v.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=v.ndim - 1 - ) - a_tensor = from_dlpack(a.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=a.ndim - 1 - ) - b_tensor = from_dlpack(b.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=b.ndim - 1 - ) - A_log_tensor = from_dlpack(A_log.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=0 - ) - dt_bias_tensor = from_dlpack( - dt_bias.detach(), assumed_align=16 - ).mark_layout_dynamic(leading_dim=dt_bias.ndim - 1) - h0_source_tensor = from_dlpack( - h0_source.detach(), assumed_align=16 - ).mark_layout_dynamic(leading_dim=h0_source.ndim - 1) - h0_indices_tensor = from_dlpack( - initial_state_indices.detach(), assumed_align=16 - ).mark_layout_dynamic(leading_dim=0) - o_tensor = from_dlpack(o.detach(), assumed_align=16).mark_layout_dynamic( - leading_dim=o.ndim - 1 - ) - - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - compiled_kernel = _get_compiled_kernel( - N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode, - scale=scale, use_qk_l2norm=use_qk_l2norm_in_kernel, - softplus_beta=softplus_beta, softplus_threshold=softplus_threshold, - ) - - compiled_kernel( - cu_seqlens_tensor, - q_tensor, - k_tensor, - v_tensor, - a_tensor, - b_tensor, - A_log_tensor, - dt_bias_tensor, - h0_source_tensor, - h0_indices_tensor, - o_tensor, - stream, - ) - - return o diff --git a/cula/lightning/__init__.py b/cula/lightning/__init__.py index 24889372..973f7707 100644 --- a/cula/lightning/__init__.py +++ b/cula/lightning/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from cula.lightning.la_decode import linear_attention_decode +from cula.ops.la_decode import linear_attention_decode from cula.ops.lightning_attn import ( LinearAttentionChunkwiseDecay, lightning_attn_fwd, diff --git a/cula/ops/__init__.py b/cula/ops/__init__.py index 44ee07fb..d3d6e7a6 100644 --- a/cula/ops/__init__.py +++ b/cula/ops/__init__.py @@ -12,3 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode +from cula.ops.la_decode import linear_attention_decode + +__all__ = [ + "kda_decode", + "fused_sigmoid_gating_delta_rule_update", + "linear_attention_decode", +] + diff --git a/cula/ops/kda_decode.py b/cula/ops/kda_decode.py new file mode 100644 index 00000000..361ae53d --- /dev/null +++ b/cula/ops/kda_decode.py @@ -0,0 +1,2091 @@ +"""CuTe DSL Fused Sigmoid Gating Delta Rule Kernel for KDA Decode. + +This version uses the production / FLA-compatible VK state layout: + state.shape == (pool_size, HV, V, K) + +The kernel still computes on a logical (K, V) matrix in shared memory. Global +state loads/stores therefore explicitly map: + global(V, K) <-> shared(K, V) + +Notes: +- This is a correctness-first implementation for decode. +- It keeps the original small-batch / large-batch split. +- It preserves the previous PAD semantics: if pool_idx < 0 the block does not + load / update / write output or state, consistent with the earlier CuTe path. +""" + +import logging + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +logger = logging.getLogger(__name__) + +# ----------------------------------------------------------------------------- +# Runtime behavior +# ----------------------------------------------------------------------------- +# The dominant setup cost in this decode path is kernel compilation. We keep a +# lightweight compile cache so each shape/configuration is compiled once, then +# launched repeatedly. Runtime dispatch uses TVM-FFI so torch tensors can be +# passed directly without rebuilding CuTe wrappers on every call. +_compiled_kernels: dict[tuple, object] = {} +_cu_seqlens_cache: dict[tuple, torch.Tensor] = {} +_stream_cache: dict[tuple, cuda.CUstream] = {} + +# ----------------------------------------------------------------------------- +# Kernel tuning constants +# ----------------------------------------------------------------------------- +# The current decode path is tuned around K=128, which is the primary target +# workload in this project. Small-batch and large-batch paths use different +# CTA organizations to balance launch overhead and throughput. +TILE_K = 128 +TILE_V = 32 +TILE_V_PADDED = 36 +TILE_V_SMALL = 16 +TILE_V_SMALL_PADDED = 20 +# Decode does not currently overlap state prefetch with compute, so a single +# shared-memory stage avoids unnecessary SMEM pressure and improves occupancy. +NUM_STAGES = 1 +NUM_THREADS = 128 +# One CTA per state is still the best default for larger decode batches, but +# very small N*H cases underfill the GPU. For those micro-batches we compile a +# dedicated split-state variant that launches multiple CTAs per state. +NUM_BLOCKS_PER_STATE_SMALL = 1 +MAX_NUM_BLOCKS_PER_STATE_SMALL = 8 +N4_NUM_BLOCKS_PER_STATE_SMALL = 4 +NUM_THREADS_LARGE = 256 +NUM_WARPS_LARGE = 8 +V_PER_WARP = 4 +ROWS_PER_ITER = 8 +NUM_K_ITERS = TILE_K // ROWS_PER_ITER +SMALL_BATCH_THRESHOLD = 1024 +MICRO_BATCH_NH_THRESHOLD = 512 +DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD = 8 +N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD = 64 +DENSE_SMALL_HV_PARALLEL_MAX_N = 16 + + +def _get_cached_cute_tensor(tensor: torch.Tensor, *, leading_dim: int, assumed_align: int = 16): + """Wrap a torch.Tensor as a CuTe tensor for compilation-time use.""" + return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic(leading_dim=leading_dim) + + +def _get_cached_stream(device: torch.device): + """Convert the active torch stream to a cached cuda.bindings CUstream.""" + stream_id = int(torch.cuda.current_stream(device=device).cuda_stream) + cache_key = (str(device), stream_id) + if cache_key not in _stream_cache: + _stream_cache[cache_key] = cuda.CUstream(stream_id) + return _stream_cache[cache_key] + + +def _get_cached_dispatch_bundle( + cu_seqlens: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + h0_source: torch.Tensor, + initial_state_indices: torch.Tensor, + o: torch.Tensor, +): + """Build the full set of CuTe tensor handles required for one kernel launch. + + This helper is now only needed for compile-time preparation. Runtime launch + uses TVM-FFI and passes torch tensors directly. + """ + return ( + _get_cached_cute_tensor(cu_seqlens, leading_dim=0), + _get_cached_cute_tensor(q, leading_dim=q.ndim - 1), + _get_cached_cute_tensor(k, leading_dim=k.ndim - 1), + _get_cached_cute_tensor(v, leading_dim=v.ndim - 1), + _get_cached_cute_tensor(a, leading_dim=a.ndim - 1), + _get_cached_cute_tensor(b, leading_dim=b.ndim - 1), + _get_cached_cute_tensor(A_log, leading_dim=0), + _get_cached_cute_tensor(dt_bias, leading_dim=dt_bias.ndim - 1), + _get_cached_cute_tensor(h0_source, leading_dim=h0_source.ndim - 1), + _get_cached_cute_tensor(initial_state_indices, leading_dim=0), + _get_cached_cute_tensor(o, leading_dim=o.ndim - 1), + ) + +def _select_small_blocks_per_state(N: int, H: int, HV: int, V: int) -> int: + del HV + num_v_tiles_small = V // TILE_V_SMALL + if N <= 4: + # For N=4, the path is launch-overhead dominated. Splitting all the way + # to 8 CTAs per state over-fragments the work and hurts latency. + return min(N4_NUM_BLOCKS_PER_STATE_SMALL, num_v_tiles_small) + if N * H <= MICRO_BATCH_NH_THRESHOLD: + return min(4, num_v_tiles_small) + return NUM_BLOCKS_PER_STATE_SMALL + + +def _try_fast_dense_decode( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor | None, + initial_state_indices: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + scale: float | None, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float, + softplus_threshold: float, + out: torch.Tensor | None, + state_layout: str | None, +): + """Fast path for the common dense decode case used by the benchmark. + + This bypasses the broader compatibility/normalization logic when inputs are + already in the exact kernel-ready layout and dtype, which materially lowers + Python-side overhead for tiny N. + """ + if initial_state_source is None or initial_state_indices is None: + return None + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + return None + + B_q, T_q, H, K = q.shape + if T_q != 1 or k.shape != q.shape or v.shape[0] != B_q or v.shape[1] != 1: + return None + + N = initial_state_indices.shape[0] + HV = v.shape[2] + V = v.shape[3] + if B_q != N or K != TILE_K or V % TILE_V_SMALL != 0 or V % TILE_V != 0: + return None + + if ( + q.device.type != "cuda" + or q.dtype != torch.bfloat16 + or k.dtype != torch.bfloat16 + or v.dtype != torch.bfloat16 + or initial_state_source.dtype != torch.float32 + or initial_state_indices.dtype != torch.int32 + or A_log.dtype != torch.float32 + or dt_bias.dtype != torch.float32 + ): + return None + + if not ( + q.is_contiguous() + and k.is_contiguous() + and v.is_contiguous() + and initial_state_source.is_contiguous() + and initial_state_indices.is_contiguous() + and A_log.is_contiguous() + and dt_bias.is_contiguous() + ): + return None + + if A_log.numel() != HV or dt_bias.shape != (HV, K): + return None + + normalized_layout = "vk" if state_layout is None else str(state_layout).strip().lower() + if normalized_layout == "vk": + if initial_state_source.ndim != 4 or initial_state_source.shape[1:] != (HV, V, K): + return None + state_layout_is_kv = False + elif normalized_layout == "kv": + if initial_state_source.ndim != 4 or initial_state_source.shape[1:] != (HV, K, V): + return None + state_layout_is_kv = True + else: + return None + + if not a.is_contiguous() or a.device != q.device or a.dtype != torch.bfloat16: + return None + if a.dim() == 4 and a.shape == (N, 1, HV, K): + a_kernel = a + elif a.dim() == 3 and a.shape == (N, 1, HV * K): + a_kernel = a.view(N, 1, HV, K) + elif a.dim() == 3 and a.shape == (N, HV, K): + a_kernel = a.unsqueeze(1) + else: + return None + + if b.device != q.device or b.dtype != torch.bfloat16 or not b.is_contiguous(): + return None + if b.dim() == 3 and b.shape == (N, 1, HV): + b_kernel = b + elif b.dim() == 2 and b.shape == (N, HV): + b_kernel = b.unsqueeze(1) + else: + return None + + if scale is None: + scale = K**-0.5 + elif scale <= 0: + return None + + o = _prepare_output_tensor(q, out, (N, 1, HV, V)) + + if cu_seqlens is not None: + if cu_seqlens.dtype != torch.int32 or cu_seqlens.numel() != N + 1: + return None + cu_seqlens_to_use = cu_seqlens.contiguous() + else: + cache_key = (N, str(q.device)) + if cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cache_key] = torch.arange(N + 1, dtype=torch.int32, device=q.device) + cu_seqlens_to_use = _cu_seqlens_cache[cache_key] + + use_small_batch = N < SMALL_BATCH_THRESHOLD + dense_small_hv_parallel_head_threshold = ( + N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD + ) + dense_small_hv_parallel = ( + use_small_batch and H <= dense_small_hv_parallel_head_threshold and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + ) + num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) + + compiled_kernel = _get_compiled_kernel( + N, + H, + HV, + K, + V, + initial_state_source.shape[0], + use_small_batch, + False, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=False, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + compiled_kernel( + cu_seqlens_to_use, + q, + k, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + initial_state_source, + initial_state_indices, + o, + _get_cached_stream(q.device), + ) + return o + + +def _define_kernels(): + """Define CuTe DSL kernels for KDA normal and varlen decode modes.""" + + NUM_WARPS_SMALL = 4 + V_PER_WARP_SMALL = TILE_V_SMALL // NUM_WARPS_SMALL + ROWS_PER_ITER_SMALL = 32 // V_PER_WARP_SMALL + NUM_K_ITERS_SMALL = TILE_K // ROWS_PER_ITER_SMALL + + @cute.kernel + def kda_kernel_small_batch( + tiled_copy_load: cute.TiledCopy, + h0_source: cute.Tensor, + smem_layout_staged: cute.Layout, + num_v_tiles: cutlass.Constexpr[int], + num_blocks_per_state_small: cutlass.Constexpr[int], + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + precomputed_decay_beta: cutlass.Constexpr[bool], + dense_small_hv_parallel: cutlass.Constexpr[bool], + ): + """Small-batch dense KDA kernel for q/k/v shaped as (N, 1, ...). + + High-level flow: + 1. Each CTA handles one (token, value-head) pair across several V tiles. + 2. q, k, and the gating decay term g are staged into shared memory. + 3. Optional q/k L2 normalization is computed at block scope. + 4. Each V tile runs one delta-rule update and writes back state/output. + """ + del tiled_copy_load + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + block_idx, _, _ = cute.arch.block_idx() + + batch_idx = block_idx // num_blocks_per_state_small + batch_inner = block_idx % num_blocks_per_state_small + num_v_tiles_per_block = num_v_tiles // num_blocks_per_state_small + start_v_tile = batch_inner * num_v_tiles_per_block + + num_value_heads_per_q = HV // H + i_n = 0 + i_h = 0 + i_hv_base = 0 + num_hv_iters = 1 + if dense_small_hv_parallel: + i_n = batch_idx // HV + i_hv_base = batch_idx % HV + i_h = i_hv_base // num_value_heads_per_q + num_hv_iters = 1 + else: + i_n = batch_idx // H + i_h = batch_idx % H + i_hv_base = i_h * num_value_heads_per_q + num_hv_iters = num_value_heads_per_q + + pool_idx = h0_indices[i_n] + + if pool_idx >= 0: + k_local = in_warp_tid // V_PER_WARP_SMALL + v_local = in_warp_tid % V_PER_WARP_SMALL + v_base = warp_idx * V_PER_WARP_SMALL + v_idx = v_base + v_local + + smem = cutlass.utils.SmemAllocator() + sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) + smem_o_layout = cute.make_layout((TILE_V_SMALL,), stride=(1,)) + smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) + smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_gk_layout = cute.make_layout((TILE_K,), stride=(1,)) + sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) + sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) + sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) + sGK = smem.allocate_tensor(cutlass.Float32, smem_gk_layout, 128) + + if tidx < TILE_K: + sK[tidx] = cutlass.Float32(k[i_n, 0, i_h, tidx]) + sQ[tidx] = cutlass.Float32(q[i_n, 0, i_h, tidx]) + cute.arch.barrier() + + if use_qk_l2norm: + sum_q_partial = 0.0 + sum_k_partial = 0.0 + if warp_idx == 0: + for norm_iter in range(4): + norm_idx = in_warp_tid + norm_iter * 32 + q_val = sQ[norm_idx] + k_val = sK[norm_idx] + sum_q_partial += q_val * q_val + sum_k_partial += k_val * k_val + + for offset in [16, 8, 4, 2, 1]: + sum_q_partial += cute.arch.shuffle_sync_bfly(sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31) + sum_k_partial += cute.arch.shuffle_sync_bfly(sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31) + + if in_warp_tid == 0: + smem_o[0] = cute.rsqrt(sum_q_partial + 1e-6) + smem_o[1] = cute.rsqrt(sum_k_partial + 1e-6) + cute.arch.barrier() + + inv_norm_q = smem_o[0] + inv_norm_k = smem_o[1] + + if tidx < TILE_K: + sK[tidx] = sK[tidx] * inv_norm_k + sQ[tidx] = sQ[tidx] * scale * inv_norm_q + cute.arch.barrier() + else: + if tidx < TILE_K: + sQ[tidx] = sQ[tidx] * scale + cute.arch.barrier() + + for hv_offset in range(num_hv_iters): + i_hv = i_hv_base + hv_offset + + if precomputed_decay_beta: + if tidx < TILE_K: + sG[tidx] = cutlass.Float32(a[i_n, 0, i_hv, tidx]) + else: + r_exp_A = 0.0 + if in_warp_tid == 0: + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv])) + r_exp_A = cute.arch.shuffle_sync(r_exp_A, 0) + if tidx < TILE_K: + r_a_k = cutlass.Float32(a[i_n, 0, i_hv, tidx]) + r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) + x = r_a_k + r_dt_bias_k + beta_x = softplus_beta * x + softplus_x = 0.0 + if beta_x <= softplus_threshold: + exp_beta_x = cute.exp(beta_x) + log_input = cutlass.Float32(1.0 + exp_beta_x) + log_result = cutlass.Float32(cute.log(log_input)) + softplus_x = cutlass.Float32((cutlass.Float32(1.0) / softplus_beta) * log_result) + else: + softplus_x = x + sG[tidx] = cute.exp(-r_exp_A * softplus_x) + + r_beta = 0.0 + if in_warp_tid == 0: + r_b = cutlass.Float32(b[i_n, 0, i_hv]) + if precomputed_decay_beta: + r_beta = r_b + else: + r_beta = 1.0 / (1.0 + cute.exp(-r_b)) + r_beta = cute.arch.shuffle_sync(r_beta, 0) + + if tidx < TILE_K: + sGK[tidx] = sG[tidx] * sK[tidx] + cute.arch.barrier() + + kv_v_load = 0 + kv_k_load_base = 0 + kv_k_load_step = 0 + vk_k_load = 0 + vk_v_load_base = 0 + vk_v_load_step = 0 + if state_layout_is_kv: + kv_v_load = tidx % TILE_V_SMALL + kv_k_load_base = tidx // TILE_V_SMALL + kv_k_load_step = NUM_THREADS // TILE_V_SMALL + else: + vk_k_load = tidx % TILE_K + vk_v_load_base = tidx // TILE_K + vk_v_load_step = NUM_THREADS // TILE_K + + for v_tile_offset in range(num_v_tiles_per_block): + stage = v_tile_offset % NUM_STAGES + v_tile = start_v_tile + v_tile_offset + v_global_base = v_tile * TILE_V_SMALL + + for k_iter in range(NUM_K_ITERS_SMALL): + k_load = 0 + v_load = 0 + if state_layout_is_kv: + k_load = kv_k_load_base + k_iter * kv_k_load_step + v_load = kv_v_load + else: + k_load = vk_k_load + v_load = vk_v_load_base + k_iter * vk_v_load_step + v_global_load = v_global_base + v_load + h_val = 0.0 + if v_global_load < v.shape[3]: + if state_layout_is_kv: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, k_load, v_global_load)]) + else: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, v_global_load, k_load)]) + sData[(k_load, v_load, stage)] = h_val + + cute.arch.barrier() + + v_global = v_tile * TILE_V_SMALL + v_idx + r_v = 0.0 + if v_global < v.shape[3]: + r_v = cutlass.Float32(v[i_n, 0, i_hv, v_global]) + + sum_hk = 0.0 + for k_iter in range(NUM_K_ITERS_SMALL): + k_base = k_iter * ROWS_PER_ITER_SMALL + k_idx = k_base + k_local + sum_hk += sData[(k_idx, v_idx, stage)] * sGK[k_idx] + + for offset in [4, 2, 1]: + sum_hk += cute.arch.shuffle_sync_bfly( + sum_hk, + offset=offset * V_PER_WARP_SMALL, + mask=-1, + mask_and_clamp=31, + ) + + v_new = (r_v - sum_hk) * r_beta + v_new = cute.arch.shuffle_sync(v_new, v_local) + + sum_hq = 0.0 + for k_iter in range(NUM_K_ITERS_SMALL): + k_base = k_iter * ROWS_PER_ITER_SMALL + k_idx = k_base + k_local + h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] + h_new = h_old + sK[k_idx] * v_new + sData[(k_idx, v_idx, stage)] = h_new + sum_hq += h_new * sQ[k_idx] + + for offset in [4, 2, 1]: + sum_hq += cute.arch.shuffle_sync_bfly( + sum_hq, + offset=offset * V_PER_WARP_SMALL, + mask=-1, + mask_and_clamp=31, + ) + + if k_local == 0 and v_global < v.shape[3]: + o[(i_n, 0, i_hv, v_global)] = cutlass.BFloat16(sum_hq) + + cute.arch.barrier() + + for k_iter in cutlass.range(NUM_K_ITERS_SMALL, unroll=2): + k_write = 0 + v_write = 0 + if state_layout_is_kv: + k_write = kv_k_load_base + k_iter * kv_k_load_step + v_write = kv_v_load + else: + k_write = vk_k_load + v_write = vk_v_load_base + k_iter * vk_v_load_step + v_global_write = v_global_base + v_write + if v_global_write < v.shape[3]: + if state_layout_is_kv: + h0_source[(pool_idx, i_hv, k_write, v_global_write)] = sData[(k_write, v_write, stage)] + else: + h0_source[(pool_idx, i_hv, v_global_write, k_write)] = sData[(k_write, v_write, stage)] + + cute.arch.barrier() + + @cute.kernel + def kda_kernel_small_batch_varlen( + tiled_copy_load: cute.TiledCopy, + h0_source: cute.Tensor, + smem_layout_staged: cute.Layout, + num_v_tiles: cutlass.Constexpr[int], + num_blocks_per_state_small: cutlass.Constexpr[int], + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + precomputed_decay_beta: cutlass.Constexpr[bool], + ): + """Small batch KDA kernel for varlen decode: q/k/v shapes (1, N, ...).""" + del tiled_copy_load + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + block_idx, _, _ = cute.arch.block_idx() + + batch_idx = block_idx // num_blocks_per_state_small + batch_inner = block_idx % num_blocks_per_state_small + num_v_tiles_per_block = num_v_tiles // num_blocks_per_state_small + start_v_tile = batch_inner * num_v_tiles_per_block + + i_n = batch_idx // HV + i_hv = batch_idx % HV + i_h = i_hv // (HV // H) + + pool_idx = h0_indices[i_n] + + if pool_idx >= 0: + k_local = in_warp_tid // V_PER_WARP_SMALL + v_local = in_warp_tid % V_PER_WARP_SMALL + v_base = warp_idx * V_PER_WARP_SMALL + v_idx = v_base + v_local + + smem = cutlass.utils.SmemAllocator() + sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) + smem_o_layout = cute.make_layout((TILE_V_SMALL,), stride=(1,)) + smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) + smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_gk_layout = cute.make_layout((TILE_K,), stride=(1,)) + sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) + sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) + sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) + sGK = smem.allocate_tensor(cutlass.Float32, smem_gk_layout, 128) + + if tidx < TILE_K: + sK[tidx] = cutlass.Float32(k[0, i_n, i_h, tidx]) + sQ[tidx] = cutlass.Float32(q[0, i_n, i_h, tidx]) + + if precomputed_decay_beta: + if tidx < TILE_K: + sG[tidx] = cutlass.Float32(a[i_n, i_hv, tidx]) + else: + r_exp_A = 0.0 + if in_warp_tid == 0: + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv])) + r_exp_A = cute.arch.shuffle_sync(r_exp_A, 0) + if tidx < TILE_K: + r_a_k = cutlass.Float32(a[i_n, i_hv, tidx]) + r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) + x = r_a_k + r_dt_bias_k + beta_x = softplus_beta * x + softplus_x = 0.0 + if beta_x <= softplus_threshold: + exp_beta_x = cute.exp(beta_x) + log_input = cutlass.Float32(1.0 + exp_beta_x) + log_result = cutlass.Float32(cute.log(log_input)) + softplus_x = cutlass.Float32((cutlass.Float32(1.0) / softplus_beta) * log_result) + else: + softplus_x = x + sG[tidx] = cute.exp(-r_exp_A * softplus_x) + + r_beta = 0.0 + if in_warp_tid == 0: + r_b = cutlass.Float32(b[i_n, i_hv]) + if precomputed_decay_beta: + r_beta = r_b + else: + r_beta = 1.0 / (1.0 + cute.exp(-r_b)) + r_beta = cute.arch.shuffle_sync(r_beta, 0) + + cute.arch.barrier() + + if use_qk_l2norm: + sum_q_partial = 0.0 + sum_k_partial = 0.0 + if warp_idx == 0: + for norm_iter in range(4): + norm_idx = in_warp_tid + norm_iter * 32 + q_val = sQ[norm_idx] + k_val = sK[norm_idx] + sum_q_partial += q_val * q_val + sum_k_partial += k_val * k_val + + for offset in [16, 8, 4, 2, 1]: + sum_q_partial += cute.arch.shuffle_sync_bfly(sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31) + sum_k_partial += cute.arch.shuffle_sync_bfly(sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31) + + if in_warp_tid == 0: + smem_o[0] = cute.rsqrt(sum_q_partial + 1e-6) + smem_o[1] = cute.rsqrt(sum_k_partial + 1e-6) + cute.arch.barrier() + + inv_norm_q = smem_o[0] + inv_norm_k = smem_o[1] + + if tidx < TILE_K: + sK[tidx] = sK[tidx] * inv_norm_k + sQ[tidx] = sQ[tidx] * scale * inv_norm_q + cute.arch.barrier() + else: + if tidx < TILE_K: + sQ[tidx] = sQ[tidx] * scale + cute.arch.barrier() + + if tidx < TILE_K: + sGK[tidx] = sG[tidx] * sK[tidx] + cute.arch.barrier() + + kv_v_load = 0 + kv_k_load_base = 0 + kv_k_load_step = 0 + vk_k_load = 0 + vk_v_load_base = 0 + vk_v_load_step = 0 + if state_layout_is_kv: + kv_v_load = tidx % TILE_V_SMALL + kv_k_load_base = tidx // TILE_V_SMALL + kv_k_load_step = NUM_THREADS // TILE_V_SMALL + else: + vk_k_load = tidx % TILE_K + vk_v_load_base = tidx // TILE_K + vk_v_load_step = NUM_THREADS // TILE_K + + for v_tile_offset in range(num_v_tiles_per_block): + stage = v_tile_offset % NUM_STAGES + v_tile = start_v_tile + v_tile_offset + v_global_base = v_tile * TILE_V_SMALL + + for k_iter in range(NUM_K_ITERS_SMALL): + k_load = 0 + v_load = 0 + if state_layout_is_kv: + k_load = kv_k_load_base + k_iter * kv_k_load_step + v_load = kv_v_load + else: + k_load = vk_k_load + v_load = vk_v_load_base + k_iter * vk_v_load_step + v_global_load = v_global_base + v_load + h_val = 0.0 + if v_global_load < v.shape[3]: + if state_layout_is_kv: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, k_load, v_global_load)]) + else: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, v_global_load, k_load)]) + sData[(k_load, v_load, stage)] = h_val + + cute.arch.barrier() + + v_global = v_tile * TILE_V_SMALL + v_idx + r_v = 0.0 + if v_global < v.shape[3]: + r_v = cutlass.Float32(v[0, i_n, i_hv, v_global]) + + sum_hk = 0.0 + for k_iter in range(NUM_K_ITERS_SMALL): + k_base = k_iter * ROWS_PER_ITER_SMALL + k_idx = k_base + k_local + sum_hk += sData[(k_idx, v_idx, stage)] * sGK[k_idx] + + for offset in [4, 2, 1]: + sum_hk += cute.arch.shuffle_sync_bfly( + sum_hk, + offset=offset * V_PER_WARP_SMALL, + mask=-1, + mask_and_clamp=31, + ) + + v_new = (r_v - sum_hk) * r_beta + v_new = cute.arch.shuffle_sync(v_new, v_local) + + sum_hq = 0.0 + for k_iter in range(NUM_K_ITERS_SMALL): + k_base = k_iter * ROWS_PER_ITER_SMALL + k_idx = k_base + k_local + h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] + h_new = h_old + sK[k_idx] * v_new + sData[(k_idx, v_idx, stage)] = h_new + sum_hq += h_new * sQ[k_idx] + + for offset in [4, 2, 1]: + sum_hq += cute.arch.shuffle_sync_bfly( + sum_hq, + offset=offset * V_PER_WARP_SMALL, + mask=-1, + mask_and_clamp=31, + ) + + if k_local == 0 and v_global < v.shape[3]: + o[(0, i_n, i_hv, v_global)] = cutlass.BFloat16(sum_hq) + + cute.arch.barrier() + + for k_iter in cutlass.range(NUM_K_ITERS_SMALL, unroll=2): + k_write = 0 + v_write = 0 + if state_layout_is_kv: + k_write = kv_k_load_base + k_iter * kv_k_load_step + v_write = kv_v_load + else: + k_write = vk_k_load + v_write = vk_v_load_base + k_iter * vk_v_load_step + v_global_write = v_global_base + v_write + if v_global_write < v.shape[3]: + if state_layout_is_kv: + h0_source[(pool_idx, i_hv, k_write, v_global_write)] = sData[(k_write, v_write, stage)] + else: + h0_source[(pool_idx, i_hv, v_global_write, k_write)] = sData[(k_write, v_write, stage)] + + cute.arch.barrier() + + @cute.kernel + def kda_kernel_large_batch( + tiled_copy_load: cute.TiledCopy, + h0_source: cute.Tensor, + smem_layout_staged: cute.Layout, + num_v_tiles: cutlass.Constexpr[int], + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + ): + """Large batch KDA kernel for dense decode: q/k/v shapes (N, 1, ...).""" + del tiled_copy_load + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + batch_idx, _, _ = cute.arch.block_idx() + + i_nhv = batch_idx // num_v_tiles + v_tile = batch_idx % num_v_tiles + i_n = i_nhv // HV + i_hv = i_nhv % HV + i_h = i_hv // (HV // H) + + pool_idx = h0_indices[i_n] + + if pool_idx >= 0: + k_local = in_warp_tid // V_PER_WARP + v_local = in_warp_tid % V_PER_WARP + v_base = warp_idx * V_PER_WARP + v_idx = v_base + v_local + + smem = cutlass.utils.SmemAllocator() + sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) + smem_o_layout = cute.make_layout((TILE_V,), stride=(1,)) + smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) + smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) + sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) + sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) + sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) + + if tidx < TILE_K: + sK[tidx] = cutlass.Float32(k[i_n, 0, i_h, tidx]) + sQ[tidx] = cutlass.Float32(q[i_n, 0, i_h, tidx]) + + r_exp_A = 0.0 + if in_warp_tid == 0: + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv])) + r_exp_A = cute.arch.shuffle_sync(r_exp_A, 0) + if tidx < TILE_K: + r_a_k = cutlass.Float32(a[i_n, 0, i_hv, tidx]) + r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) + x = r_a_k + r_dt_bias_k + beta_x = softplus_beta * x + softplus_x = 0.0 + if beta_x <= softplus_threshold: + exp_beta_x = cute.exp(beta_x) + log_input = cutlass.Float32(1.0 + exp_beta_x) + log_result = cutlass.Float32(cute.log(log_input)) + softplus_x = cutlass.Float32((cutlass.Float32(1.0) / softplus_beta) * log_result) + else: + softplus_x = x + sG[tidx] = cute.exp(-r_exp_A * softplus_x) + + r_beta = 0.0 + if in_warp_tid == 0: + r_b = cutlass.Float32(b[i_n, 0, i_hv]) + r_beta = 1.0 / (1.0 + cute.exp(-r_b)) + r_beta = cute.arch.shuffle_sync(r_beta, 0) + + cute.arch.barrier() + + if use_qk_l2norm: + sum_q_partial = 0.0 + sum_k_partial = 0.0 + if tidx < TILE_K: + q_val = sQ[tidx] + k_val = sK[tidx] + sum_q_partial = q_val * q_val + sum_k_partial = k_val * k_val + + for offset in [16, 8, 4, 2, 1]: + sum_q_partial += cute.arch.shuffle_sync_bfly(sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31) + sum_k_partial += cute.arch.shuffle_sync_bfly(sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31) + + if in_warp_tid == 0: + smem_o[warp_idx] = sum_q_partial + smem_o[warp_idx + 8] = sum_k_partial + cute.arch.barrier() + + if warp_idx == 0: + local_sum_q = 0.0 + local_sum_k = 0.0 + if in_warp_tid < NUM_WARPS_LARGE: + local_sum_q = smem_o[in_warp_tid] + local_sum_k = smem_o[in_warp_tid + 8] + for offset in [4, 2, 1]: + local_sum_q += cute.arch.shuffle_sync_bfly(local_sum_q, offset=offset, mask=-1, mask_and_clamp=31) + local_sum_k += cute.arch.shuffle_sync_bfly(local_sum_k, offset=offset, mask=-1, mask_and_clamp=31) + if in_warp_tid == 0: + smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) + smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) + cute.arch.barrier() + + inv_norm_q = smem_o[0] + inv_norm_k = smem_o[1] + + if tidx < TILE_K: + sK[tidx] = sK[tidx] * inv_norm_k + sQ[tidx] = sQ[tidx] * scale * inv_norm_q + cute.arch.barrier() + else: + if tidx < TILE_K: + sQ[tidx] = sQ[tidx] * scale + cute.arch.barrier() + + stage = 0 + + for k_iter in range(NUM_K_ITERS): + flat_idx = tidx + k_iter * NUM_THREADS_LARGE + k_load = 0 + v_load = 0 + if state_layout_is_kv: + k_load = flat_idx // TILE_V + v_load = flat_idx % TILE_V + else: + k_load = flat_idx % TILE_K + v_load = flat_idx // TILE_K + v_global_load = v_tile * TILE_V + v_load + h_val = 0.0 + if v_global_load < v.shape[3]: + if state_layout_is_kv: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, k_load, v_global_load)]) + else: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, v_global_load, k_load)]) + sData[(k_load, v_load, stage)] = h_val + + cute.arch.barrier() + + v_global = v_tile * TILE_V + v_idx + r_v = 0.0 + if v_global < v.shape[3]: + r_v = cutlass.Float32(v[i_n, 0, i_hv, v_global]) + + sum_hk = 0.0 + for k_iter in range(NUM_K_ITERS): + k_base = k_iter * ROWS_PER_ITER + k_idx = k_base + k_local + sum_hk += sData[(k_idx, v_idx, stage)] * sG[k_idx] * sK[k_idx] + + for offset in [4, 2, 1]: + sum_hk += cute.arch.shuffle_sync_bfly( + sum_hk, + offset=offset * V_PER_WARP, + mask=-1, + mask_and_clamp=31, + ) + + v_new = (r_v - sum_hk) * r_beta + v_new = cute.arch.shuffle_sync(v_new, v_local) + + sum_hq = 0.0 + for k_iter in range(NUM_K_ITERS): + k_base = k_iter * ROWS_PER_ITER + k_idx = k_base + k_local + h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] + h_new = h_old + sK[k_idx] * v_new + sData[(k_idx, v_idx, stage)] = h_new + sum_hq += h_new * sQ[k_idx] + + for offset in [4, 2, 1]: + sum_hq += cute.arch.shuffle_sync_bfly( + sum_hq, + offset=offset * V_PER_WARP, + mask=-1, + mask_and_clamp=31, + ) + + if k_local == 0 and v_global < v.shape[3]: + o[(i_n, 0, i_hv, v_global)] = cutlass.BFloat16(sum_hq) + + cute.arch.barrier() + + for k_iter in cutlass.range(NUM_K_ITERS, unroll=2): + flat_idx = tidx + k_iter * NUM_THREADS_LARGE + k_write = 0 + v_write = 0 + if state_layout_is_kv: + k_write = flat_idx // TILE_V + v_write = flat_idx % TILE_V + else: + k_write = flat_idx % TILE_K + v_write = flat_idx // TILE_K + v_global_write = v_tile * TILE_V + v_write + if v_global_write < v.shape[3]: + if state_layout_is_kv: + h0_source[(pool_idx, i_hv, k_write, v_global_write)] = sData[(k_write, v_write, stage)] + else: + h0_source[(pool_idx, i_hv, v_global_write, k_write)] = sData[(k_write, v_write, stage)] + + @cute.kernel + def kda_kernel_large_batch_varlen( + tiled_copy_load: cute.TiledCopy, + h0_source: cute.Tensor, + smem_layout_staged: cute.Layout, + num_v_tiles: cutlass.Constexpr[int], + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + ): + """Large batch KDA kernel for varlen decode: q/k/v shapes (1, N, ...).""" + del tiled_copy_load + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + batch_idx, _, _ = cute.arch.block_idx() + + i_nhv = batch_idx // num_v_tiles + v_tile = batch_idx % num_v_tiles + i_n = i_nhv // HV + i_hv = i_nhv % HV + i_h = i_hv // (HV // H) + + pool_idx = h0_indices[i_n] + + if pool_idx >= 0: + k_local = in_warp_tid // V_PER_WARP + v_local = in_warp_tid % V_PER_WARP + v_base = warp_idx * V_PER_WARP + v_idx = v_base + v_local + + smem = cutlass.utils.SmemAllocator() + sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) + smem_o_layout = cute.make_layout((TILE_V,), stride=(1,)) + smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) + smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) + smem_g_layout = cute.make_layout((TILE_K,), stride=(1,)) + sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) + sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) + sG = smem.allocate_tensor(cutlass.Float32, smem_g_layout, 128) + + if tidx < TILE_K: + sK[tidx] = cutlass.Float32(k[0, i_n, i_h, tidx]) + sQ[tidx] = cutlass.Float32(q[0, i_n, i_h, tidx]) + + r_exp_A = 0.0 + if in_warp_tid == 0: + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv])) + r_exp_A = cute.arch.shuffle_sync(r_exp_A, 0) + if tidx < TILE_K: + r_a_k = cutlass.Float32(a[i_n, i_hv, tidx]) + r_dt_bias_k = cutlass.Float32(dt_bias[i_hv, tidx]) + x = r_a_k + r_dt_bias_k + beta_x = softplus_beta * x + softplus_x = 0.0 + if beta_x <= softplus_threshold: + exp_beta_x = cute.exp(beta_x) + log_input = cutlass.Float32(1.0 + exp_beta_x) + log_result = cutlass.Float32(cute.log(log_input)) + softplus_x = cutlass.Float32((cutlass.Float32(1.0) / softplus_beta) * log_result) + else: + softplus_x = x + sG[tidx] = cute.exp(-r_exp_A * softplus_x) + + r_beta = 0.0 + if in_warp_tid == 0: + r_b = cutlass.Float32(b[i_n, i_hv]) + r_beta = 1.0 / (1.0 + cute.exp(-r_b)) + r_beta = cute.arch.shuffle_sync(r_beta, 0) + + cute.arch.barrier() + + if use_qk_l2norm: + sum_q_partial = 0.0 + sum_k_partial = 0.0 + if tidx < TILE_K: + q_val = sQ[tidx] + k_val = sK[tidx] + sum_q_partial = q_val * q_val + sum_k_partial = k_val * k_val + + for offset in [16, 8, 4, 2, 1]: + sum_q_partial += cute.arch.shuffle_sync_bfly(sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31) + sum_k_partial += cute.arch.shuffle_sync_bfly(sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31) + + if in_warp_tid == 0: + smem_o[warp_idx] = sum_q_partial + smem_o[warp_idx + 8] = sum_k_partial + cute.arch.barrier() + + if warp_idx == 0: + local_sum_q = 0.0 + local_sum_k = 0.0 + if in_warp_tid < NUM_WARPS_LARGE: + local_sum_q = smem_o[in_warp_tid] + local_sum_k = smem_o[in_warp_tid + 8] + for offset in [4, 2, 1]: + local_sum_q += cute.arch.shuffle_sync_bfly(local_sum_q, offset=offset, mask=-1, mask_and_clamp=31) + local_sum_k += cute.arch.shuffle_sync_bfly(local_sum_k, offset=offset, mask=-1, mask_and_clamp=31) + if in_warp_tid == 0: + smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) + smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) + cute.arch.barrier() + + inv_norm_q = smem_o[0] + inv_norm_k = smem_o[1] + + if tidx < TILE_K: + sK[tidx] = sK[tidx] * inv_norm_k + sQ[tidx] = sQ[tidx] * scale * inv_norm_q + cute.arch.barrier() + else: + if tidx < TILE_K: + sQ[tidx] = sQ[tidx] * scale + cute.arch.barrier() + + stage = 0 + + for k_iter in range(NUM_K_ITERS): + flat_idx = tidx + k_iter * NUM_THREADS_LARGE + k_load = 0 + v_load = 0 + if state_layout_is_kv: + k_load = flat_idx // TILE_V + v_load = flat_idx % TILE_V + else: + k_load = flat_idx % TILE_K + v_load = flat_idx // TILE_K + v_global_load = v_tile * TILE_V + v_load + h_val = 0.0 + if v_global_load < v.shape[3]: + if state_layout_is_kv: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, k_load, v_global_load)]) + else: + h_val = cutlass.Float32(h0_source[(pool_idx, i_hv, v_global_load, k_load)]) + sData[(k_load, v_load, stage)] = h_val + + cute.arch.barrier() + + v_global = v_tile * TILE_V + v_idx + r_v = 0.0 + if v_global < v.shape[3]: + r_v = cutlass.Float32(v[0, i_n, i_hv, v_global]) + + sum_hk = 0.0 + for k_iter in range(NUM_K_ITERS): + k_base = k_iter * ROWS_PER_ITER + k_idx = k_base + k_local + sum_hk += sData[(k_idx, v_idx, stage)] * sG[k_idx] * sK[k_idx] + + for offset in [4, 2, 1]: + sum_hk += cute.arch.shuffle_sync_bfly( + sum_hk, + offset=offset * V_PER_WARP, + mask=-1, + mask_and_clamp=31, + ) + + v_new = (r_v - sum_hk) * r_beta + v_new = cute.arch.shuffle_sync(v_new, v_local) + + sum_hq = 0.0 + for k_iter in range(NUM_K_ITERS): + k_base = k_iter * ROWS_PER_ITER + k_idx = k_base + k_local + h_old = sData[(k_idx, v_idx, stage)] * sG[k_idx] + h_new = h_old + sK[k_idx] * v_new + sData[(k_idx, v_idx, stage)] = h_new + sum_hq += h_new * sQ[k_idx] + + for offset in [4, 2, 1]: + sum_hq += cute.arch.shuffle_sync_bfly( + sum_hq, + offset=offset * V_PER_WARP, + mask=-1, + mask_and_clamp=31, + ) + + if k_local == 0 and v_global < v.shape[3]: + o[(0, i_n, i_hv, v_global)] = cutlass.BFloat16(sum_hq) + + cute.arch.barrier() + + for k_iter in cutlass.range(NUM_K_ITERS, unroll=2): + flat_idx = tidx + k_iter * NUM_THREADS_LARGE + k_write = 0 + v_write = 0 + if state_layout_is_kv: + k_write = flat_idx // TILE_V + v_write = flat_idx % TILE_V + else: + k_write = flat_idx % TILE_K + v_write = flat_idx // TILE_K + v_global_write = v_tile * TILE_V + v_write + if v_global_write < v.shape[3]: + if state_layout_is_kv: + h0_source[(pool_idx, i_hv, k_write, v_global_write)] = sData[(k_write, v_write, stage)] + else: + h0_source[(pool_idx, i_hv, v_global_write, k_write)] = sData[(k_write, v_write, stage)] + + return ( + kda_kernel_small_batch, + kda_kernel_small_batch_varlen, + kda_kernel_large_batch, + kda_kernel_large_batch_varlen, + ) + + +def _create_jit_functions(): + """Create JIT-compiled launcher functions for all KDA kernel variants.""" + + kda_small, kda_small_varlen, kda_large, kda_large_varlen = _define_kernels() + + @cute.jit + def run_small_batch( + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + h0_source: cute.Tensor, + h0_indices: cute.Tensor, + o: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_initial_state: cutlass.Constexpr[bool], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + precomputed_decay_beta: cutlass.Constexpr[bool], + num_blocks_per_state_small: cutlass.Constexpr[int], + dense_small_hv_parallel: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + del cu_seqlens, B, T, K, use_initial_state + n_indices = h0_indices.layout.shape[0] + batch_size = n_indices * (HV if dense_small_hv_parallel else H) + + num_v_tiles_small = cute.ceil_div(V, TILE_V_SMALL) + smem_layout_small = cute.make_layout( + (TILE_K, TILE_V_SMALL, NUM_STAGES), + stride=(TILE_V_SMALL_PADDED, 1, TILE_K * TILE_V_SMALL_PADDED), + ) + smem_bytes_small = 4 * TILE_K * TILE_V_SMALL_PADDED * NUM_STAGES + 4 * TILE_V_SMALL + 4 * TILE_K * 4 + 64 + + kda_small( + None, + h0_source, + smem_layout_small, + num_v_tiles_small, + num_blocks_per_state_small, + q, + k, + v, + a, + b, + A_log, + dt_bias, + o, + h0_indices, + softplus_beta, + softplus_threshold, + scale, + H, + HV, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + dense_small_hv_parallel, + ).launch( + grid=(batch_size * num_blocks_per_state_small, 1, 1), + block=[NUM_THREADS, 1, 1], + smem=smem_bytes_small, + stream=stream, + ) + + @cute.jit + def run_small_batch_varlen( + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + h0_source: cute.Tensor, + h0_indices: cute.Tensor, + o: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_initial_state: cutlass.Constexpr[bool], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + precomputed_decay_beta: cutlass.Constexpr[bool], + num_blocks_per_state_small: cutlass.Constexpr[int], + dense_small_hv_parallel: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + del cu_seqlens, B, T, K, use_initial_state, dense_small_hv_parallel + n_indices = h0_indices.layout.shape[0] + batch_size = n_indices * HV + + num_v_tiles_small = cute.ceil_div(V, TILE_V_SMALL) + smem_layout_small = cute.make_layout( + (TILE_K, TILE_V_SMALL, NUM_STAGES), + stride=(TILE_V_SMALL_PADDED, 1, TILE_K * TILE_V_SMALL_PADDED), + ) + smem_bytes_small = 4 * TILE_K * TILE_V_SMALL_PADDED * NUM_STAGES + 4 * TILE_V_SMALL + 4 * TILE_K * 4 + 64 + + kda_small_varlen( + None, + h0_source, + smem_layout_small, + num_v_tiles_small, + num_blocks_per_state_small, + q, + k, + v, + a, + b, + A_log, + dt_bias, + o, + h0_indices, + softplus_beta, + softplus_threshold, + scale, + H, + HV, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + ).launch( + grid=(batch_size * num_blocks_per_state_small, 1, 1), + block=[NUM_THREADS, 1, 1], + smem=smem_bytes_small, + stream=stream, + ) + + @cute.jit + def run_large_batch( + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + h0_source: cute.Tensor, + h0_indices: cute.Tensor, + o: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_initial_state: cutlass.Constexpr[bool], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + precomputed_decay_beta: cutlass.Constexpr[bool], + num_blocks_per_state_small: cutlass.Constexpr[int], + dense_small_hv_parallel: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + del cu_seqlens, B, T, K, use_initial_state, precomputed_decay_beta, num_blocks_per_state_small, dense_small_hv_parallel + n_indices = h0_indices.layout.shape[0] + batch_size = n_indices * HV + + num_v_tiles = cute.ceil_div(V, TILE_V) + smem_layout = cute.make_layout( + (TILE_K, TILE_V, NUM_STAGES), + stride=(TILE_V_PADDED, 1, TILE_K * TILE_V_PADDED), + ) + smem_bytes = 4 * TILE_K * TILE_V_PADDED * NUM_STAGES + 4 * TILE_V + 4 * TILE_K * 2 + 4 * TILE_K + 64 + + kda_large( + None, + h0_source, + smem_layout, + num_v_tiles, + q, + k, + v, + a, + b, + A_log, + dt_bias, + o, + h0_indices, + softplus_beta, + softplus_threshold, + scale, + H, + HV, + use_qk_l2norm, + state_layout_is_kv, + ).launch( + grid=(batch_size * num_v_tiles, 1, 1), + block=[NUM_THREADS_LARGE, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + @cute.jit + def run_large_batch_varlen( + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + h0_source: cute.Tensor, + h0_indices: cute.Tensor, + o: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_initial_state: cutlass.Constexpr[bool], + use_qk_l2norm: cutlass.Constexpr[bool], + state_layout_is_kv: cutlass.Constexpr[bool], + precomputed_decay_beta: cutlass.Constexpr[bool], + num_blocks_per_state_small: cutlass.Constexpr[int], + dense_small_hv_parallel: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + del cu_seqlens, B, T, K, use_initial_state, precomputed_decay_beta, num_blocks_per_state_small, dense_small_hv_parallel + n_indices = h0_indices.layout.shape[0] + batch_size = n_indices * HV + + num_v_tiles = cute.ceil_div(V, TILE_V) + smem_layout = cute.make_layout( + (TILE_K, TILE_V, NUM_STAGES), + stride=(TILE_V_PADDED, 1, TILE_K * TILE_V_PADDED), + ) + smem_bytes = 4 * TILE_K * TILE_V_PADDED * NUM_STAGES + 4 * TILE_V + 4 * TILE_K * 2 + 4 * TILE_K + 64 + + kda_large_varlen( + None, + h0_source, + smem_layout, + num_v_tiles, + q, + k, + v, + a, + b, + A_log, + dt_bias, + o, + h0_indices, + softplus_beta, + softplus_threshold, + scale, + H, + HV, + use_qk_l2norm, + state_layout_is_kv, + ).launch( + grid=(batch_size * num_v_tiles, 1, 1), + block=[NUM_THREADS_LARGE, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + return ( + run_small_batch, + run_small_batch_varlen, + run_large_batch, + run_large_batch_varlen, + ) + + +_jit_functions = None + + +def _get_jit_functions(): + global _jit_functions + if _jit_functions is None: + _jit_functions = _create_jit_functions() + return _jit_functions + + +def _get_compiled_kernel( + N, + H, + HV, + K, + V, + pool_size, + use_small_batch, + is_varlen_decode, + scale, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + num_blocks_per_state_small, + dense_small_hv_parallel, + softplus_beta, + softplus_threshold, +): + """Get or lazily compile one CuteDSL decode kernel variant. + + Compile-time specialization is still important here, so we cache the result + by shape, layout, and constexpr options. The compiled function is emitted + with TVM-FFI enabled so runtime calls can pass torch tensors directly. + """ + global _compiled_kernels + + key = ( + N, + H, + HV, + K, + V, + pool_size, + use_small_batch, + is_varlen_decode, + scale, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + num_blocks_per_state_small, + dense_small_hv_parallel, + softplus_beta, + softplus_threshold, + ) + if key in _compiled_kernels: + return _compiled_kernels[key] + + cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda") + + if is_varlen_decode: + q = torch.zeros(1, N, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(1, N, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") + else: + q = torch.zeros(N, 1, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(N, 1, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, 1, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, 1, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") + + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + if state_layout_is_kv: + h0_source = torch.zeros(pool_size, HV, K, V, dtype=torch.float32, device="cuda") + else: + h0_source = torch.zeros(pool_size, HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + + cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16) + q_tensor = from_dlpack(q, assumed_align=16) + k_tensor = from_dlpack(k, assumed_align=16) + v_tensor = from_dlpack(v, assumed_align=16) + a_tensor = from_dlpack(a, assumed_align=16) + b_tensor = from_dlpack(b, assumed_align=16) + A_log_tensor = from_dlpack(A_log, assumed_align=16) + dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) + h0_source_tensor = from_dlpack(h0_source, assumed_align=16) + h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16) + o_tensor = from_dlpack(o, assumed_align=16) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + run_small, run_small_varlen, run_large, run_large_varlen = _get_jit_functions() + if use_small_batch: + kernel_func = run_small_varlen if is_varlen_decode else run_small + else: + kernel_func = run_large_varlen if is_varlen_decode else run_large + + compiled_kernel = cute.compile( + kernel_func, + cu_seqlens_tensor, + q_tensor, + k_tensor, + v_tensor, + a_tensor, + b_tensor, + A_log_tensor, + dt_bias_tensor, + h0_source_tensor, + h0_indices_tensor, + o_tensor, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + B=1 if is_varlen_decode else N, + T=N if is_varlen_decode else 1, + H=H, + K=K, + V=V, + HV=HV, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=precomputed_decay_beta, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + stream=stream, + options="--enable-tvm-ffi --opt-level 1", + ) + + _compiled_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA kernel compiled: " + f"N={N}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, " + f"small_batch={use_small_batch}, varlen={is_varlen_decode}" + ) + return compiled_kernel + + +def _normalize_A_log(A_log: torch.Tensor, HV: int) -> torch.Tensor: + if A_log.numel() != HV: + raise ValueError(f"Unexpected A_log shape: {A_log.shape}; expected numel={HV}") + return A_log.reshape(HV).contiguous() + + +def _normalize_dt_bias(dt_bias: torch.Tensor, HV: int, K: int) -> torch.Tensor: + if dt_bias.numel() != HV * K: + raise ValueError(f"Unexpected dt_bias shape: {dt_bias.shape}; expected numel={HV * K}") + return dt_bias.reshape(HV, K).contiguous() + + +def _canonicalize_state_layout(state_layout: str | None) -> str: + """Accept only the two explicit state layouts used by the kernel. + + Internal meaning: + - "vk": state shape (..., V, K) + - "kv": state shape (..., K, V) + """ + if state_layout is None: + return "vk" + + normalized = str(state_layout).strip().lower() + if normalized not in ("vk", "kv"): + raise ValueError(f"Unsupported state_layout={state_layout}; expected only 'vk' or 'kv'") + return normalized + + +def _normalize_kda_a(a, *, is_varlen_decode, N, HV, K): + """Normalize `a` to match the compile-time shape expected by the kernel. + + Supports both cuLA-native layouts and the public flattened compatibility layouts. + + varlen kernel compiled shape: (N, HV, K) -- 3D + dense kernel compiled shape: (N, 1, HV, K) -- 4D + """ + if is_varlen_decode: + # Target: (N, HV, K) -- 3D + if a.dim() == 2 and a.shape == (N, HV * K): + return a.view(N, HV, K) + if a.dim() == 3 and a.shape == (N, HV, K): + return a + if a.dim() == 3 and a.shape == (N, 1, HV * K): + return a.view(N, HV, K) + if a.dim() == 3 and a.shape == (1, N, HV * K): + return a.view(N, HV, K) + if a.dim() == 4 and a.shape == (1, N, HV, K): + return a.squeeze(0) + if a.dim() == 4 and a.shape == (1, N, 1, HV * K): + return a.view(N, HV, K) + raise ValueError(f"Unexpected a shape for varlen: {a.shape}") + else: + # Target: (N, 1, HV, K) -- 4D + if a.dim() == 2 and a.shape == (N, HV * K): + return a.view(N, 1, HV, K) + if a.dim() == 3 and a.shape == (N, HV, K): + return a.unsqueeze(1) + if a.dim() == 3 and a.shape == (N, 1, HV * K): + return a.view(N, 1, HV, K) + if a.dim() == 4 and a.shape == (N, 1, HV, K): + return a + raise ValueError(f"Unexpected a shape for dense: {a.shape}") + + +def _normalize_state_source(initial_state_source, *, N, HV, K, V, device, state_layout="vk"): + """Validate that the incoming state already matches the requested layout.""" + if initial_state_source is None: + if state_layout == "vk": + h0_source = torch.zeros(N, HV, V, K, dtype=torch.float32, device=device) + return h0_source, N, False + h0_source = torch.zeros(N, HV, K, V, dtype=torch.float32, device=device) + return h0_source, N, True + + if initial_state_source.dim() != 4: + raise ValueError(f"Unexpected initial_state_source shape: {initial_state_source.shape}; expected a 4D state tensor") + + if initial_state_source.shape[1] != HV: + raise ValueError(f"Unexpected initial_state_source shape: {initial_state_source.shape}; expected HV={HV}") + + if state_layout == "vk": + if initial_state_source.shape[2:] != (V, K): + raise ValueError( + f"State layout mismatch for state_layout='vk': got {initial_state_source.shape}, expected (..., {HV}, {V}, {K})" + ) + return initial_state_source, initial_state_source.shape[0], False + + if initial_state_source.shape[2:] != (K, V): + raise ValueError( + f"State layout mismatch for state_layout='kv': got {initial_state_source.shape}, expected (..., {HV}, {K}, {V})" + ) + return initial_state_source, initial_state_source.shape[0], True + + +def _normalize_state_indices(initial_state_indices, *, N, pool_size, device): + """Normalize state indices for decode. + + For compatibility callers, missing indices default to a sequential mapping. + """ + if initial_state_indices is None: + if pool_size < N: + raise ValueError(f"initial_state_source only has pool_size={pool_size}, but N={N}") + return torch.arange(N, device=device, dtype=torch.int32) + + indices = initial_state_indices.to(device=device, dtype=torch.int32) + if indices.numel() != N: + raise ValueError(f"Unexpected initial_state_indices shape: {initial_state_indices.shape}; expected numel={N}") + return indices.contiguous() + + +def _prepare_output_tensor(q: torch.Tensor, out: torch.Tensor | None, shape: tuple[int, ...]) -> torch.Tensor: + if out is None: + return q.new_empty(shape, dtype=torch.bfloat16) + if out.shape != shape: + raise ValueError(f"Unexpected out shape: {out.shape}; expected {shape}") + if out.device != q.device: + raise ValueError(f"Unexpected out device: {out.device}; expected {q.device}") + if out.dtype != torch.bfloat16: + raise ValueError(f"Unexpected out dtype: {out.dtype}; expected torch.bfloat16") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + return out + + +def fused_sigmoid_gating_delta_rule_update( + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + softplus_beta: float, + softplus_threshold: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor | None, + initial_state_indices: torch.Tensor | None, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + is_kda: bool = False, + out: torch.Tensor | None = None, + state_layout: str = "vk", +): + """Public cuLA decode API backed by CuTe DSL. + + Supported state layouts: + - "vk": state shape (pool_size, HV, V, K), default and recommended + - "kv": state shape (pool_size, HV, K, V) + + The caller is expected to pass a state tensor that already matches the + selected layout exactly. + """ + if not is_kda: + raise NotImplementedError("cuLA fused decode currently supports only is_kda=True mode") + + return kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + out=out, + state_layout=state_layout, + ) + + +def kda_decode( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + state_layout: str = "vk", +) -> torch.Tensor: + """CuTe DSL implementation of fused sigmoid gating KDA update. + + State layout contract: + - "vk": (pool_size, HV, V, K), default + - "kv": (pool_size, HV, K, V) + + Dense decode: + q/k: (N, 1, H, K) + v: (N, 1, HV, V) + a: (N, 1, HV, K) + b: (N, 1, HV) + + Varlen decode: + q/k: (1, N, H, K) + v: (1, N, HV, V) + a: (N, HV, K) or (1, N, HV, K) + b: (N, HV) or (1, N, HV) + """ + + B_q, T_q, H, K = q.shape + HV = v.shape[2] + V = v.shape[3] + if initial_state_indices is not None: + N = initial_state_indices.shape[0] + else: + N = T_q if B_q == 1 and T_q > 1 else B_q + + if scale is None: + scale = K**-0.5 + else: + assert scale > 0, f"scale must be positive, got {scale}" + + state_layout = _canonicalize_state_layout(state_layout) + + fast_dense_out = _try_fast_dense_decode( + A_log, + dt_bias, + q, + k, + v, + a, + b, + initial_state_source, + initial_state_indices, + cu_seqlens, + scale, + use_qk_l2norm_in_kernel, + softplus_beta, + softplus_threshold, + out, + state_layout, + ) + if fast_dense_out is not None: + return fast_dense_out + + A_log = A_log.contiguous() + + assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}" + assert V % TILE_V_SMALL == 0, f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" + assert V % TILE_V == 0, f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}" + num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) + assert (V // TILE_V_SMALL) % num_blocks_per_state_small == 0, ( + f"Small-batch KDA kernel requires num_v_tiles_small divisible by {num_blocks_per_state_small}, got V={V}" + ) + + is_varlen_decode = B_q == 1 and T_q == N and N > 1 + + # The public API only accepts the two explicit kernel layouts. + + # Small and large batches use different kernel organizations: + # small batches prioritize lower launch overhead, while large batches focus + # on sustained throughput. + use_small_batch = N < SMALL_BATCH_THRESHOLD + dense_small_hv_parallel_head_threshold = ( + N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD + ) + dense_small_hv_parallel = ( + use_small_batch + and (not is_varlen_decode) + and H <= dense_small_hv_parallel_head_threshold + and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + ) + + # fast_path means the incoming state already matches one of the expected + # layouts, so we can skip extra normalization work. + fast_path = False + state_layout_is_kv = False + pool_size = N + h0_source = initial_state_source + + if h0_source is None: + if state_layout == "vk": + h0_source = torch.zeros(N, HV, V, K, dtype=torch.float32, device=q.device) + state_layout_is_kv = False + else: + h0_source = torch.zeros(N, HV, K, V, dtype=torch.float32, device=q.device) + state_layout_is_kv = True + pool_size = N + fast_path = True + elif h0_source.dim() == 4 and h0_source.shape[1] == HV: + pool_size = h0_source.shape[0] + if state_layout == "kv": + if h0_source.shape[2:] != (K, V): + raise ValueError(f"State layout mismatch for state_layout='kv': got {h0_source.shape}, expected (..., {HV}, {K}, {V})") + state_layout_is_kv = True + fast_path = True + else: + if h0_source.shape[2:] != (V, K): + raise ValueError(f"State layout mismatch for state_layout='vk': got {h0_source.shape}, expected (..., {HV}, {V}, {K})") + fast_path = True + + if fast_path: + a_fast = _normalize_kda_a(a, is_varlen_decode=is_varlen_decode, N=N, HV=HV, K=K) + if is_varlen_decode: + if b.dim() == 3: + b = b.squeeze(0) + o = _prepare_output_tensor(q, out, (1, N, HV, V)) + else: + if b.dim() == 2: + b = b.unsqueeze(1) + o = _prepare_output_tensor(q, out, (N, 1, HV, V)) + a = a_fast + if initial_state_indices is None: + if pool_size < N: + fast_path = False + else: + initial_state_indices = torch.arange(N, device=q.device, dtype=torch.int32) + elif ( + initial_state_indices.device != q.device + or initial_state_indices.dtype != torch.int32 + or initial_state_indices.numel() != N + ): + fast_path = False + + if not fast_path: + h0_source, pool_size, state_layout_is_kv = _normalize_state_source( + initial_state_source, + N=N, + HV=HV, + K=K, + V=V, + device=q.device, + state_layout=state_layout, + ) + + a = _normalize_kda_a(a, is_varlen_decode=is_varlen_decode, N=N, HV=HV, K=K) + + if is_varlen_decode: + # varlen b compiled: (N, HV) -- 2D + if b.dim() == 3: + b = b.squeeze(0) # (1, N, HV) -> (N, HV) + # b should be 2D (N, HV) + o = _prepare_output_tensor(q, out, (1, N, HV, V)) + else: + # dense b compiled: (N, 1, HV) -- 3D + if b.dim() == 2: + b = b.unsqueeze(1) + # b should be 3D (N, 1, HV) + o = _prepare_output_tensor(q, out, (N, 1, HV, V)) + + q = q if q.is_contiguous() else q.contiguous() + k = k if k.is_contiguous() else k.contiguous() + v = v if v.is_contiguous() else v.contiguous() + a = a if a.is_contiguous() else a.contiguous() + b = b if b.is_contiguous() else b.contiguous() + dt_bias = dt_bias if dt_bias.is_contiguous() else dt_bias.contiguous() + + if cu_seqlens is not None: + cu_seqlens_to_use = cu_seqlens + else: + cache_key = (N, str(q.device)) + if cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cache_key] = torch.arange(N + 1, dtype=torch.int32, device=q.device) + cu_seqlens_to_use = _cu_seqlens_cache[cache_key] + + A_log = _normalize_A_log(A_log, HV) + dt_bias = _normalize_dt_bias(dt_bias, HV, K) + + precomputed_decay_beta = False + a_kernel, b_kernel = a, b + + if not fast_path: + initial_state_indices = _normalize_state_indices( + initial_state_indices, + N=N, + pool_size=pool_size, + device=q.device, + ) + if cu_seqlens is not None: + cu_seqlens_to_use = cu_seqlens.contiguous() + + stream = _get_cached_stream(q.device) + + compiled_kernel = _get_compiled_kernel( + N, + H, + HV, + K, + V, + pool_size, + use_small_batch, + is_varlen_decode, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=precomputed_decay_beta, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + + # With TVM-FFI enabled at compile time, the runtime launch can pass torch + # tensors directly instead of rebuilding CuTe tensor wrappers for each call. + compiled_kernel( + cu_seqlens_to_use, + q, + k, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + h0_source, + initial_state_indices, + o, + stream, + ) + + return o diff --git a/cula/ops/kda_decode_fla.py b/cula/ops/kda_decode_fla.py new file mode 100644 index 00000000..2b724f58 --- /dev/null +++ b/cula/ops/kda_decode_fla.py @@ -0,0 +1,247 @@ +from typing import Optional + +import torch +import triton +import triton.language as tl + +@triton.jit(do_not_specialize=["T"]) +def fused_sigmoid_gating_delta_rule_update_kernel( + A_log, + a, + dt_bias, + softplus_beta, + softplus_threshold, + q, + k, + v, + b, + o, + h0_source, + h0_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_KDA: tl.constexpr, +): + """ + Fused kernel that combines sigmoid gating computation with recurrent delta rule update. + """ + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + p_b = b + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + # Gating computation pointers + p_A_log = A_log + i_hv + if IS_KDA: + # For KDA, a has shape [T, HV*K] (flattened), so stride is HV*K per token + p_a = a + bos * HV * K + i_hv * K + o_k + p_dt_bias = dt_bias + i_hv * K + o_k + else: + p_a = a + bos * HV + i_hv + p_dt_bias = dt_bias + i_hv + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + # Load inputs + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_b = tl.load(p_b).to(tl.float32) + + # Compute sigmoid gating + # Load gating parameters + b_A_log = tl.load(p_A_log).to(tl.float32) + b_a = tl.load(p_a).to(tl.float32) + b_dt_bias = tl.load(p_dt_bias).to(tl.float32) + + # Compute g = -exp(A_log) * softplus(a + dt_bias) + x = b_a + b_dt_bias + beta_x = softplus_beta * x + # Apply softplus with numerical stability + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + + # Compute beta = sigmoid(b) + b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + + # Apply L2 normalization if enabled + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6)) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) + + b_q = b_q * scale + + # Apply gating to hidden state: h *= exp(g) + if IS_KDA: + b_h *= tl.exp(b_g[:, None]) + else: + b_h *= tl.exp(b_g) + + # Delta rule: v -= sum(h * k, dim=0) + b_v -= tl.sum(b_h * b_k[:, None], 0) + + # Apply beta gating: v *= beta + b_v *= b_beta + + # Update hidden state: h += k[:, None] * v[None, :] + b_h += b_k[:, None] * b_v[None, :] + + # Compute output: o = sum(h * q, dim=0) + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # Update pointers for next timestep + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_b += HV + if IS_KDA: + p_a += HV * K + else: + p_a += HV + + # Store final state back to h0_source with bounds checking + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +def fused_sigmoid_gating_delta_rule_update( + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + softplus_beta: float, + softplus_threshold: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + is_kda: bool = False, +): + """ + Fused triton implementation of sigmoid gating delta rule update. + This function uses a single fused kernel that combines both sigmoid gating computation + and the recurrent delta rule update for better performance. + """ + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + + # Kernel needs valid tensors for h0_source and h0_indices even when not using initial state. + # Create dummy tensors if needed. + use_initial_state = initial_state_source is not None + if initial_state_indices is None: + initial_state_indices = torch.full((N,), -1, dtype=torch.int32, device=q.device) + if initial_state_source is None: + # Create a dummy state buffer with minimal shape + initial_state_source = torch.zeros(1, HV, K, V, dtype=q.dtype, device=q.device) + + o = q.new_empty(NK, *v.shape) + grid = (NK, NV, N * HV) + + fused_sigmoid_gating_delta_rule_update_kernel[grid]( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + q=q, + k=k, + v=v, + b=b, + o=o, + h0_source=initial_state_source, + h0_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + USE_INITIAL_STATE=use_initial_state, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + IS_VARLEN=cu_seqlens is not None, + IS_KDA=is_kda, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o diff --git a/cula/lightning/la_decode.py b/cula/ops/la_decode.py similarity index 100% rename from cula/lightning/la_decode.py rename to cula/ops/la_decode.py diff --git a/tests/test_kda_decode.py b/tests/test_kda_decode.py index d0320334..2f5469f3 100644 --- a/tests/test_kda_decode.py +++ b/tests/test_kda_decode.py @@ -33,22 +33,22 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.kda import kda_decode +from cula.kda import kda_decode, fused_sigmoid_gating_delta_rule_update # --------------------------------------------------------------------------- # PyTorch reference # --------------------------------------------------------------------------- def torch_kda_decode_ref( - q, # (N, H, K) float32 - k, # (N, H, K) float32 (H here is the query/key head count) - v, # (N, HV, V) float32 - a, # (N, HV, K) float32 - b, # (N, HV) float32 - A_log, # (HV,) float32 - dt_bias, # (HV, K) float32 - state, # (N, HV, V, K) float32 - scale, # float + q, # (N, H, K) float32 + k, # (N, H, K) float32 (H here is the query/key head count) + v, # (N, HV, V) float32 + a, # (N, HV, K) float32 + b, # (N, HV) float32 + A_log, # (HV,) float32 + dt_bias, # (HV, K) float32 + state, # (N, HV, V, K) float32 + scale, # float use_l2norm=True, softplus_beta=1.0, softplus_threshold=20.0, @@ -126,13 +126,13 @@ def run_kda_decode_dense(q, k, v, a, b, A_log, dt_bias, state, scale): HV, V = v.shape[1], v.shape[2] # Reshape to dense layout: (N, 1, H, K), etc. - q_4d = q.unsqueeze(1).contiguous() # (N, 1, H, K) - k_4d = k.unsqueeze(1).contiguous() # (N, 1, H, K) — note: H not HV for k - v_4d = v.unsqueeze(1).contiguous() # (N, 1, HV, V) - a_4d = a.unsqueeze(1).contiguous() # (N, 1, HV, K) - b_3d = b.unsqueeze(1).contiguous() # (N, 1, HV) + q_4d = q.unsqueeze(1).contiguous() # (N, 1, H, K) + k_4d = k.unsqueeze(1).contiguous() # (N, 1, H, K) — note: H not HV for k + v_4d = v.unsqueeze(1).contiguous() # (N, 1, HV, V) + a_4d = a.unsqueeze(1).contiguous() # (N, 1, HV, K) + b_3d = b.unsqueeze(1).contiguous() # (N, 1, HV) - state_source = state.clone().contiguous() # (N, HV, V, K) + state_source = state.clone().contiguous() # (N, HV, V, K) indices = torch.arange(N, device=q.device, dtype=torch.int32) o = kda_decode( @@ -158,13 +158,13 @@ def run_kda_decode_varlen(q, k, v, a, b, A_log, dt_bias, state, scale): HV, V = v.shape[1], v.shape[2] # Reshape to varlen layout: (1, N, H, K), etc. - q_4d = q.unsqueeze(0).contiguous() # (1, N, H, K) - k_4d = k.unsqueeze(0).contiguous() # (1, N, H, K) - v_4d = v.unsqueeze(0).contiguous() # (1, N, HV, V) - a_3d = a.contiguous() # (N, HV, K) — varlen uses 3D - b_2d = b.contiguous() # (N, HV) — varlen uses 2D + q_4d = q.unsqueeze(0).contiguous() # (1, N, H, K) + k_4d = k.unsqueeze(0).contiguous() # (1, N, H, K) + v_4d = v.unsqueeze(0).contiguous() # (1, N, HV, V) + a_3d = a.contiguous() # (N, HV, K) — varlen uses 3D + b_2d = b.contiguous() # (N, HV) — varlen uses 2D - state_source = state.clone().contiguous() # (N, HV, V, K) + state_source = state.clone().contiguous() # (N, HV, V, K) indices = torch.arange(N, device=q.device, dtype=torch.int32) cu_seqlens = torch.arange(N + 1, device=q.device, dtype=torch.int32) @@ -192,10 +192,34 @@ def _assert_close(name, ref, actual, atol=3e-2, rtol=2e-2): max_diff = diff.max().item() mean_diff = diff.mean().item() ok = torch.allclose(ref.float(), actual.float(), atol=atol, rtol=rtol) - assert ok, ( - f"{name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}, " - f"atol={atol}, rtol={rtol}" + assert ok, f"{name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}, atol={atol}, rtol={rtol}" + + +def run_kda_decode_triton_compatible(q, k, v, a, b, A_log, dt_bias, state_kv, scale, fused_fn=None): + """Run the Triton-compatible cuLA decode API.""" + N = q.shape[0] + indices = torch.arange(N, device=q.device, dtype=torch.int32) + out = q.new_empty(N, 1, v.shape[1], v.shape[2], dtype=torch.bfloat16) + fused_fn = fused_sigmoid_gating_delta_rule_update if fused_fn is None else fused_fn + o = fused_fn( + A_log=A_log, + a=a.reshape(N, 1, -1).to(torch.bfloat16), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q.unsqueeze(1).to(torch.bfloat16), + k=k.unsqueeze(1).to(torch.bfloat16), + v=v.unsqueeze(1).to(torch.bfloat16), + b=b.unsqueeze(1).to(torch.bfloat16), + initial_state_source=state_kv, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + is_kda=True, + out=out, ) + return o.squeeze(1), state_kv.permute(0, 1, 3, 2).contiguous() # --------------------------------------------------------------------------- @@ -210,13 +234,28 @@ def test_kda_decode_dense(N, H, HV): # Reference (fp32) o_ref, state_ref = torch_kda_decode_ref( - q.float(), k.float(), v.float(), a, b.float(), - A_log, dt_bias, state.clone(), scale, + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state.clone(), + scale, ) # Kernel o_kernel, state_kernel = run_kda_decode_dense( - q, k, v, a, b, A_log, dt_bias, state, scale, + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, ) _assert_close("output", o_ref, o_kernel.float()) @@ -235,13 +274,28 @@ def test_kda_decode_varlen(N, H, HV): # Reference (fp32) o_ref, state_ref = torch_kda_decode_ref( - q.float(), k.float(), v.float(), a, b.float(), - A_log, dt_bias, state.clone(), scale, + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state.clone(), + scale, ) # Kernel o_kernel, state_kernel = run_kda_decode_varlen( - q, k, v, a, b, A_log, dt_bias, state, scale, + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, ) _assert_close("output", o_ref, o_kernel.float()) @@ -258,12 +312,27 @@ def test_kda_decode_large_v(N): q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) o_ref, state_ref = torch_kda_decode_ref( - q.float(), k.float(), v.float(), a, b.float(), - A_log, dt_bias, state.clone(), scale, + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state.clone(), + scale, ) o_kernel, state_kernel = run_kda_decode_dense( - q, k, v, a, b, A_log, dt_bias, state, scale, + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, ) _assert_close("output", o_ref, o_kernel.float()) @@ -280,17 +349,142 @@ def test_kda_decode_zero_state(): state = torch.zeros(N, HV, V, K, device="cuda", dtype=torch.float32) o_ref, state_ref = torch_kda_decode_ref( - q.float(), k.float(), v.float(), a, b.float(), - A_log, dt_bias, state.clone(), scale, + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state.clone(), + scale, ) o_kernel, state_kernel = run_kda_decode_dense( - q, k, v, a, b, A_log, dt_bias, state, scale, + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, ) _assert_close("output", o_ref, o_kernel.float()) _assert_close("state", state_ref, state_kernel) +def test_kda_decode_layout_defaults_to_vk(): + N, H, HV, K, V = 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state_vk = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = torch_kda_decode_ref( + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state_vk.clone(), + scale, + ) + + indices = torch.arange(N, device=q.device, dtype=torch.int32) + out = q.new_empty(N, 1, HV, V, dtype=torch.bfloat16) + state_default = state_vk.clone().contiguous() + state_v_last = state_vk.permute(0, 1, 3, 2).contiguous() + + o_default = fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a.reshape(N, 1, -1).to(torch.bfloat16), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q.unsqueeze(1).to(torch.bfloat16), + k=k.unsqueeze(1).to(torch.bfloat16), + v=v.unsqueeze(1).to(torch.bfloat16), + b=b.unsqueeze(1).to(torch.bfloat16), + initial_state_source=state_default, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + out=out, + ) + + o_v_last = fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a.reshape(N, 1, -1).to(torch.bfloat16), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q.unsqueeze(1).to(torch.bfloat16), + k=k.unsqueeze(1).to(torch.bfloat16), + v=v.unsqueeze(1).to(torch.bfloat16), + b=b.unsqueeze(1).to(torch.bfloat16), + initial_state_source=state_v_last, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="kv", + ) + + _assert_close("default-vk output", o_ref, o_default.squeeze(1).float()) + _assert_close("default-vk state", state_ref, state_default) + _assert_close("kv output", o_ref, o_v_last.squeeze(1).float()) + _assert_close("kv state", state_ref, state_v_last.permute(0, 1, 3, 2).contiguous()) + + +def test_kda_decode_layout_mismatch_raises(): + N, H, HV, K, V = 4, 8, 16, 128, 256 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state_vk = make_inputs(N, H, HV, K, V) + indices = torch.arange(N, device=q.device, dtype=torch.int32) + state_kv = state_vk.permute(0, 1, 3, 2).contiguous() + + with pytest.raises(ValueError, match="State layout mismatch"): + fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a.reshape(N, 1, -1).to(torch.bfloat16), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q.unsqueeze(1).to(torch.bfloat16), + k=k.unsqueeze(1).to(torch.bfloat16), + v=v.unsqueeze(1).to(torch.bfloat16), + b=b.unsqueeze(1).to(torch.bfloat16), + initial_state_source=state_kv, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="vk", + ) + + with pytest.raises(ValueError, match="State layout mismatch"): + fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a.reshape(N, 1, -1).to(torch.bfloat16), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q.unsqueeze(1).to(torch.bfloat16), + k=k.unsqueeze(1).to(torch.bfloat16), + v=v.unsqueeze(1).to(torch.bfloat16), + b=b.unsqueeze(1).to(torch.bfloat16), + initial_state_source=state_vk, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="kv", + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/tests/test_la_decode.py b/tests/test_la_decode.py index d7a3a404..41ee05ed 100644 --- a/tests/test_la_decode.py +++ b/tests/test_la_decode.py @@ -30,7 +30,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.lightning.la_decode import linear_attention_decode +from cula.ops.la_decode import linear_attention_decode try: from fla.ops.common.fused_recurrent import fused_recurrent_fwd From 50dd3f16a1a977a36d54b7688a87b51c959c9ae2 Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Thu, 23 Apr 2026 01:08:32 +0800 Subject: [PATCH 04/34] fix(linter): fix linter & updates kda decode perf (#60) --- BENCHMARK_KDA_DECODE_GB200.md | 56 +++++++++++++++++----------------- benchmarks/bench_kda_decode.py | 24 +++++++-------- cula/ops/__init__.py | 1 - cula/ops/kda_decode.py | 13 +++++--- cula/ops/kda_decode_fla.py | 23 +++----------- tests/test_kda_decode.py | 4 +-- 6 files changed, 55 insertions(+), 66 deletions(-) diff --git a/BENCHMARK_KDA_DECODE_GB200.md b/BENCHMARK_KDA_DECODE_GB200.md index b1ebd4c2..a31b85df 100644 --- a/BENCHMARK_KDA_DECODE_GB200.md +++ b/BENCHMARK_KDA_DECODE_GB200.md @@ -1,15 +1,15 @@ # Benchmark Results - KDA Decode -> Auto-generated by `benchmarks/bench_kda_decode.py` on 2026-04-18 23:38:47. +> Auto-generated by `benchmarks/bench_kda_decode.py` on 2026-04-22 19:27:18. -> **GPU:** NVIDIA GB200 | **CUDA:** 13.0 | **PyTorch:** 2.11.0+cu130 | **Python:** 3.12.3 +> **GPU:** NVIDIA GB200 | **CUDA:** 13.0 | **PyTorch:** 2.9.1+cu130 | **Python:** 3.12.3 > Decode setting: single-token (T=1), K=128, V=128. ## Summary -- v-last speedup (FLA/cuLA): avg=1.02x, min=0.94x, max=1.16x -- k-last speedup (FLA/cuLA): avg=1.15x, min=1.01x, max=1.31x +- v-last speedup (FLA/cuLA): avg=1.11x, min=1.00x, max=1.25x +- k-last speedup (FLA/cuLA): avg=1.24x, min=1.14x, max=1.41x - Batch sizes: [1, 4, 8, 16, 32, 64, 128, 256] - Q/K heads (H): [8, 32, 64] - V heads (HV): 128 @@ -21,14 +21,14 @@ | N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | |---|------------------:|------------------:|----------------:|---------------:|---------------:| -| 1 | 0.0423 | 0.0408 | 0.0490 | **1.16x** | **1.20x** | -| 4 | 0.0533 | 0.0495 | 0.0503 | **0.94x** | **1.02x** | -| 8 | 0.0409 | 0.0379 | 0.0415 | **1.01x** | **1.09x** | -| 16 | 0.0705 | 0.0623 | 0.0693 | **0.98x** | **1.11x** | -| 32 | 0.1382 | 0.1194 | 0.1313 | **0.95x** | **1.10x** | -| 64 | 0.2630 | 0.2322 | 0.2505 | **0.95x** | **1.08x** | -| 128 | 0.5008 | 0.4275 | 0.4892 | **0.98x** | **1.14x** | -| 256 | 0.9801 | 0.8387 | 0.9666 | **0.99x** | **1.15x** | +| 1 | 0.0518 | 0.0510 | 0.0616 | **1.19x** | **1.21x** | +| 4 | 0.0586 | 0.0544 | 0.0619 | **1.06x** | **1.14x** | +| 8 | 0.0512 | 0.0494 | 0.0586 | **1.14x** | **1.19x** | +| 16 | 0.0705 | 0.0611 | 0.0726 | **1.03x** | **1.19x** | +| 32 | 0.1360 | 0.1191 | 0.1384 | **1.02x** | **1.16x** | +| 64 | 0.2620 | 0.2301 | 0.2656 | **1.01x** | **1.15x** | +| 128 | 0.5000 | 0.4243 | 0.5206 | **1.04x** | **1.23x** | +| 256 | 0.9759 | 0.8331 | 1.0317 | **1.06x** | **1.24x** | ### Accuracy (Output) @@ -62,14 +62,14 @@ | N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | |---|------------------:|------------------:|----------------:|---------------:|---------------:| -| 1 | 0.0455 | 0.0441 | 0.0499 | **1.09x** | **1.13x** | -| 4 | 0.0532 | 0.0499 | 0.0505 | **0.95x** | **1.01x** | -| 8 | 0.0407 | 0.0363 | 0.0431 | **1.06x** | **1.19x** | -| 16 | 0.0746 | 0.0664 | 0.0705 | **0.95x** | **1.06x** | -| 32 | 0.1314 | 0.1135 | 0.1312 | **1.00x** | **1.16x** | -| 64 | 0.2538 | 0.2227 | 0.2516 | **0.99x** | **1.13x** | -| 128 | 0.4659 | 0.3944 | 0.4962 | **1.06x** | **1.26x** | -| 256 | 0.9040 | 0.7630 | 0.9710 | **1.07x** | **1.27x** | +| 1 | 0.0522 | 0.0502 | 0.0605 | **1.16x** | **1.20x** | +| 4 | 0.0585 | 0.0546 | 0.0647 | **1.11x** | **1.19x** | +| 8 | 0.0541 | 0.0509 | 0.0667 | **1.23x** | **1.31x** | +| 16 | 0.0736 | 0.0646 | 0.0739 | **1.00x** | **1.14x** | +| 32 | 0.1308 | 0.1123 | 0.1398 | **1.07x** | **1.24x** | +| 64 | 0.2528 | 0.2207 | 0.2662 | **1.05x** | **1.21x** | +| 128 | 0.4669 | 0.3911 | 0.5248 | **1.12x** | **1.34x** | +| 256 | 0.9017 | 0.7582 | 1.0352 | **1.15x** | **1.37x** | ### Accuracy (Output) @@ -103,14 +103,14 @@ | N | cuLA v-last (ms) | cuLA k-last (ms) | FLA v-last (ms) | v-last speedup | k-last speedup | |---|------------------:|------------------:|----------------:|---------------:|---------------:| -| 1 | 0.0453 | 0.0444 | 0.0505 | **1.11x** | **1.14x** | -| 4 | 0.0544 | 0.0502 | 0.0518 | **0.95x** | **1.03x** | -| 8 | 0.0401 | 0.0359 | 0.0392 | **0.98x** | **1.09x** | -| 16 | 0.0685 | 0.0605 | 0.0699 | **1.02x** | **1.15x** | -| 32 | 0.1324 | 0.1156 | 0.1316 | **0.99x** | **1.14x** | -| 64 | 0.2394 | 0.2020 | 0.2520 | **1.05x** | **1.25x** | -| 128 | 0.4578 | 0.3845 | 0.4923 | **1.08x** | **1.28x** | -| 256 | 0.8939 | 0.7465 | 0.9784 | **1.09x** | **1.31x** | +| 1 | 0.0521 | 0.0524 | 0.0651 | **1.25x** | **1.24x** | +| 4 | 0.0568 | 0.0526 | 0.0620 | **1.09x** | **1.18x** | +| 8 | 0.0469 | 0.0457 | 0.0548 | **1.17x** | **1.20x** | +| 16 | 0.0679 | 0.0592 | 0.0736 | **1.08x** | **1.24x** | +| 32 | 0.1309 | 0.1138 | 0.1400 | **1.07x** | **1.23x** | +| 64 | 0.2401 | 0.2011 | 0.2678 | **1.12x** | **1.33x** | +| 128 | 0.4587 | 0.3828 | 0.5244 | **1.14x** | **1.37x** | +| 256 | 0.8923 | 0.7435 | 1.0453 | **1.17x** | **1.41x** | ### Accuracy (Output) diff --git a/benchmarks/bench_kda_decode.py b/benchmarks/bench_kda_decode.py index 0658623c..ed5209c4 100644 --- a/benchmarks/bench_kda_decode.py +++ b/benchmarks/bench_kda_decode.py @@ -165,7 +165,7 @@ def write_markdown_report(args, gpu_name: str, sections: list[tuple[int, int, li def summary(vals): if not vals: return "n/a" - return f"avg={sum(vals)/len(vals):.2f}x, min={min(vals):.2f}x, max={max(vals):.2f}x" + return f"avg={sum(vals) / len(vals):.2f}x, min={min(vals):.2f}x, max={max(vals):.2f}x" lines = [] lines.append("# Benchmark Results - KDA Decode") @@ -340,9 +340,15 @@ def setup_fla_v_last(): state_bench_fla_v_last.copy_(state_init_v_last) with torch.no_grad(): - t_cula_v_last = benchmark_fn(lambda: call_cula_v_last(state_bench_cula_v_last), setup_fn=setup_cula_v_last, warmup=w, rep=r) - t_cula_k_last = benchmark_fn(lambda: call_cula_k_last(state_bench_cula_k_last), setup_fn=setup_cula_k_last, warmup=w, rep=r) - t_fla_v_last = benchmark_fn(lambda: call_fla_v_last(state_bench_fla_v_last), setup_fn=setup_fla_v_last, warmup=w, rep=r) + t_cula_v_last = benchmark_fn( + lambda: call_cula_v_last(state_bench_cula_v_last), setup_fn=setup_cula_v_last, warmup=w, rep=r + ) + t_cula_k_last = benchmark_fn( + lambda: call_cula_k_last(state_bench_cula_k_last), setup_fn=setup_cula_k_last, warmup=w, rep=r + ) + t_fla_v_last = benchmark_fn( + lambda: call_fla_v_last(state_bench_fla_v_last), setup_fn=setup_fla_v_last, warmup=w, rep=r + ) return { "N": N, @@ -421,10 +427,7 @@ def print_section(h_dim: int, v_dim: int): ) print() - hdr_out = ( - f"{'N':>5} | {'cuLA v out RMSE':>16} | {'rel':>10} | " - f"{'cuLA k out RMSE':>16} | {'rel':>10}" - ) + hdr_out = f"{'N':>5} | {'cuLA v out RMSE':>16} | {'rel':>10} | {'cuLA k out RMSE':>16} | {'rel':>10}" print(hdr_out) print("-" * len(hdr_out)) for res in results: @@ -434,10 +437,7 @@ def print_section(h_dim: int, v_dim: int): ) print() - hdr_state = ( - f"{'N':>5} | {'cuLA v state RMSE':>18} | {'rel':>10} | " - f"{'cuLA k state RMSE':>18} | {'rel':>10}" - ) + hdr_state = f"{'N':>5} | {'cuLA v state RMSE':>18} | {'rel':>10} | {'cuLA k state RMSE':>18} | {'rel':>10}" print(hdr_state) print("-" * len(hdr_state)) for res in results: diff --git a/cula/ops/__init__.py b/cula/ops/__init__.py index d3d6e7a6..6450488b 100644 --- a/cula/ops/__init__.py +++ b/cula/ops/__init__.py @@ -20,4 +20,3 @@ "fused_sigmoid_gating_delta_rule_update", "linear_attention_decode", ] - diff --git a/cula/ops/kda_decode.py b/cula/ops/kda_decode.py index 361ae53d..d84c77bf 100644 --- a/cula/ops/kda_decode.py +++ b/cula/ops/kda_decode.py @@ -114,6 +114,7 @@ def _get_cached_dispatch_bundle( _get_cached_cute_tensor(o, leading_dim=o.ndim - 1), ) + def _select_small_blocks_per_state(N: int, H: int, HV: int, V: int) -> int: del HV num_v_tiles_small = V // TILE_V_SMALL @@ -245,7 +246,7 @@ def _try_fast_dense_decode( N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD ) dense_small_hv_parallel = ( - use_small_batch and H <= dense_small_hv_parallel_head_threshold and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + use_small_batch and dense_small_hv_parallel_head_threshold >= H and N <= DENSE_SMALL_HV_PARALLEL_MAX_N ) num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) @@ -1938,7 +1939,7 @@ def kda_decode( dense_small_hv_parallel = ( use_small_batch and (not is_varlen_decode) - and H <= dense_small_hv_parallel_head_threshold + and dense_small_hv_parallel_head_threshold >= H and N <= DENSE_SMALL_HV_PARALLEL_MAX_N ) @@ -1962,12 +1963,16 @@ def kda_decode( pool_size = h0_source.shape[0] if state_layout == "kv": if h0_source.shape[2:] != (K, V): - raise ValueError(f"State layout mismatch for state_layout='kv': got {h0_source.shape}, expected (..., {HV}, {K}, {V})") + raise ValueError( + f"State layout mismatch for state_layout='kv': got {h0_source.shape}, expected (..., {HV}, {K}, {V})" + ) state_layout_is_kv = True fast_path = True else: if h0_source.shape[2:] != (V, K): - raise ValueError(f"State layout mismatch for state_layout='vk': got {h0_source.shape}, expected (..., {HV}, {V}, {K})") + raise ValueError( + f"State layout mismatch for state_layout='vk': got {h0_source.shape}, expected (..., {HV}, {V}, {K})" + ) fast_path = True if fast_path: diff --git a/cula/ops/kda_decode_fla.py b/cula/ops/kda_decode_fla.py index 2b724f58..2c4d0cfc 100644 --- a/cula/ops/kda_decode_fla.py +++ b/cula/ops/kda_decode_fla.py @@ -1,9 +1,8 @@ -from typing import Optional - import torch import triton import triton.language as tl + @triton.jit(do_not_specialize=["T"]) def fused_sigmoid_gating_delta_rule_update_kernel( A_log, @@ -78,13 +77,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( if USE_INITIAL_STATE: idx = tl.load(h0_indices + i_n) if idx >= 0: - p_h0 = ( - h0_source - + idx * HV * K * V - + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] - ) + p_h0 = h0_source + idx * HV * K * V + i_hv * K * V + o_k[:, None] * V + o_v[None, :] b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) for _ in range(0, T): @@ -155,13 +148,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( if USE_INITIAL_STATE: idx = tl.load(h0_indices + i_n) if idx >= 0: - p_h0 = ( - h0_source - + idx * HV * K * V - + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] - ) + p_h0 = h0_source + idx * HV * K * V + i_hv * K * V + o_k[:, None] * V + o_v[None, :] tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) @@ -177,9 +164,9 @@ def fused_sigmoid_gating_delta_rule_update( b: torch.Tensor, initial_state_source: torch.Tensor, initial_state_indices: torch.Tensor, - scale: Optional[float] = None, + scale: float | None = None, use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, + cu_seqlens: torch.Tensor | None = None, is_kda: bool = False, ): """ diff --git a/tests/test_kda_decode.py b/tests/test_kda_decode.py index 2f5469f3..121e6eee 100644 --- a/tests/test_kda_decode.py +++ b/tests/test_kda_decode.py @@ -33,7 +33,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.kda import kda_decode, fused_sigmoid_gating_delta_rule_update +from cula.kda import fused_sigmoid_gating_delta_rule_update, kda_decode # --------------------------------------------------------------------------- @@ -123,7 +123,6 @@ def make_inputs(N, H, HV, K, V, device="cuda", seed=42): def run_kda_decode_dense(q, k, v, a, b, A_log, dt_bias, state, scale): """Run kda_decode in dense layout: (N, 1, H/HV, dim).""" N, H, K = q.shape - HV, V = v.shape[1], v.shape[2] # Reshape to dense layout: (N, 1, H, K), etc. q_4d = q.unsqueeze(1).contiguous() # (N, 1, H, K) @@ -155,7 +154,6 @@ def run_kda_decode_dense(q, k, v, a, b, A_log, dt_bias, state, scale): def run_kda_decode_varlen(q, k, v, a, b, A_log, dt_bias, state, scale): """Run kda_decode in varlen layout: (1, N, H/HV, dim).""" N, H, K = q.shape - HV, V = v.shape[1], v.shape[2] # Reshape to varlen layout: (1, N, H, K), etc. q_4d = q.unsqueeze(0).contiguous() # (1, N, H, K) From 33ec0b3a2e24c61b0da3607cd54e564ccdd89b34 Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Thu, 23 Apr 2026 01:08:57 +0800 Subject: [PATCH 05/34] feat: BHVK (K-last) state layout for Lightning Attention prefill & decode (#56) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * opt: implement pretransposed state layout (BHVK) for lightning attention * test: update test suite for BHVK state layout * Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Refactor tensor initialization by removing transpose * feat: SMEM-mediated coalesced state load/store for BHVK layout Implement cooperative GMEM↔SMEM transpose for VK state access: - 128 CE threads load/store states cooperatively (coalesced GMEM) - 2 strips of 64 V-rows with padded SMEM (stride D+4) to eliminate bank conflicts - Row-based addressing: each iteration covers 4 rows × 32 cols with LDG/STG.128 - Per-thread SMEM read/write uses padded stride for conflict-free bank access Performance (vs FLA): h0_ht avg 1.33x, varlen avg 1.44x, no_state avg 1.46x Precision: RMSE matches FLA exactly (0.2347% O, 0.0198% Ht) All 10 tests pass * fix(linter): fix linter * fix: la_decode test state layout BHKV↔BHVK conversion run_la_decode now transposes state_4d from BHKV to BHVK before passing to the kernel, and transposes output state back to BHKV for comparison. E2E test also converts prefill BHVK output to BHKV before passing to run_la_decode and torch_la_decode_ref. All 20 tests pass. --------- Co-authored-by: Mr_antimatter Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- benchmarks/bench_lightning_attn.py | 7 +- cula/ops/lightning_attn.py | 104 ++++++++++++++++++++++++++--- tests/test_la_decode.py | 53 ++++++++++++--- tests/test_lightning_attn.py | 25 ++++--- 4 files changed, 157 insertions(+), 32 deletions(-) diff --git a/benchmarks/bench_lightning_attn.py b/benchmarks/bench_lightning_attn.py index b8ebfe4f..f4668a69 100644 --- a/benchmarks/bench_lightning_attn.py +++ b/benchmarks/bench_lightning_attn.py @@ -269,6 +269,7 @@ def benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, i output_ht = mode == "h0_ht" h0 = torch.randn(B, H, D, D, dtype=torch.float32, device=DEVICE) * 0.01 if has_h0 else None h0_fla = h0.clone() if h0 is not None else None + h0_cute = h0.transpose(-1, -2).contiguous() if h0 is not None else None # BHVK for CuTe result = {"B": B, "T": T, "H": H, "D": D, "mode": mode} ht_fla = None @@ -286,7 +287,7 @@ def benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, i # --- CuteDSL --- try: - o_cute, ht_cute, cute_ms, compile_ms = run_cutedsl(Q, K, V, decay, h0, output_ht, warmup, iters) + o_cute, ht_cute, cute_ms, compile_ms = run_cutedsl(Q, K, V, decay, h0_cute, output_ht, warmup, iters) result["cutedsl_ms"] = cute_ms result["compile_ms"] = compile_ms except Exception as e: @@ -316,7 +317,9 @@ def benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, i result[f"{label}_o_rmse_ratio"] = rmse / (rms + 1e-8) result[f"{label}_o_maxdiff"] = diff.abs().max().item() if output_ht and ht_naive is not None and ht_test is not None: - diff_ht = ht_naive - ht_test.float() + # CuTe kernel outputs BHVK state; transpose to BHKV for comparison + ht_cmp = ht_test.transpose(-1, -2).float() if label == "cute" else ht_test.float() + diff_ht = ht_naive - ht_cmp ht_rms = ht_naive.pow(2).mean().sqrt().item() ht_rmse = diff_ht.pow(2).mean().sqrt().item() result[f"{label}_ht_rmse_ratio"] = ht_rmse / (ht_rms + 1e-8) diff --git a/cula/ops/lightning_attn.py b/cula/ops/lightning_attn.py index 057cc736..8a6b204e 100644 --- a/cula/ops/lightning_attn.py +++ b/cula/ops/lightning_attn.py @@ -405,14 +405,15 @@ def __call__( v = cute.make_tensor(v_in.iterator, v_layout) o = cute.make_tensor(o_in.iterator, o_layout) - # Initial state / final state: [B, H, D, D] stored as row-major FP32 + # Initial state / final state: [B, H, D, D] in BHVK layout (K-contiguous) + # CuTe shape (V, K, (H, B)) with strides (D, 1, ...) for K-contiguous access. # When has_initial_state / output_final_state is False, None is passed # and the parameter is eliminated at compile time via const_expr guards. # For varlen: state pool is [pool_size, H, D, D]. We use B (=N) as the # pool dimension — strides are correct regardless of actual pool_size. fstate_layout = cute.make_layout( (D, D, (H, B)), - stride=(1, D, (D * D, D * D * H)), + stride=(D, 1, (D * D, D * D * H)), ) if cutlass.const_expr(self.has_initial_state): initial_state = cute.make_tensor(initial_state_in.iterator, fstate_layout) @@ -720,6 +721,15 @@ class SharedStorage: sWorkIdx: cute.struct.MemRange[Int32, 2] # Double-buffered scheduling mbarriers (count=1 each, Load warp elect_one arrives) sched_mbar: cute.struct.MemRange[Int64, 2] + # State transpose buffer for coalesced BHVK state load/store. + # Sized for 64 V-rows × (D+4) K-values in f32 (~33KB for D=128). + # Padding of 4 elements per row eliminates SMEM bank conflicts + # (stride 132 mod 32 = 4 → consecutive rows hit distinct banks). + # Used in 2 strips to load/store the D×D state with coalesced GMEM access. + sStateBuf: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, 64 * (self.K + 4)], + self.buffer_align_bytes, + ] self.shared_storage = SharedStorage @@ -946,6 +956,8 @@ def kernel( sK_weighted = storage.sK_weighted.get_tensor(kv_k_smem_layout_single.outer, swizzle=kv_k_smem_layout_single.inner) # Decay lookup table sDecayLUT = cute.make_tensor(storage.sDecayLUT.data_ptr(), cute.make_layout(self.chunk_size)) + # State transpose buffer (SMEM) for coalesced BHVK GMEM access + state_buf_ptr = storage.sStateBuf.data_ptr() # (((64,2),16),1,4,2):(((1,4096),64),0,1024,8192)> sV = storage.sV.get_tensor(v_smem_layout_staged.outer, swizzle=v_smem_layout_staged.inner) # (MMA, MMA_N, MMA_K, STAGE) @@ -1700,11 +1712,48 @@ def kernel( decay_s_cuda = decay_tensor_cuda[hidx] block_decay = cute.exp(-decay_s_cuda * cutlass.Float32(C), fastmath=self.use_fast_math) - # -------------- Initial State Loading (h0) ---------------- + # -------------- Initial State Loading (h0) via SMEM transpose ----- + # VK GMEM layout is K-contiguous but non-coalesced across threads. + # We load cooperatively (all threads, coalesced) to SMEM, then + # each thread reads its V-row from SMEM. 2 strips of 64 V-rows. if cutlass.const_expr(self.has_initial_state): gState_h0 = initial_state[None, None, (hidx, state_idx)] - gRow_h0 = cute.make_tensor(gState_h0.iterator + local_tidx, cute.make_layout(_D, stride=_D)) - cute.autovec_copy(gRow_h0, init_flat) + _STRIP = 64 + _PAD = 4 + _VEC = 4 + _SMEM_STRIDE = _D + _PAD # 132 — eliminates bank conflicts + _CE_THREADS = self.threads_per_warp * len(self.cuda_warp_ids) + _COLS_PER_THREAD = _D // _VEC # 32 threads cover one row + _ROWS_PER_ITER = _CE_THREADS // _COLS_PER_THREAD # 4 rows + _NUM_ITERS = _STRIP // _ROWS_PER_ITER # 16 + row_in_iter = local_tidx // _COLS_PER_THREAD + col_base = (local_tidx % _COLS_PER_THREAD) * _VEC + for strip_idx in cutlass.range(_D // _STRIP): + strip_base = strip_idx * (_STRIP * _D) + # Cooperative coalesced GMEM → padded SMEM (row-by-row) + for ci in cutlass.range(_NUM_ITERS, unroll_full=True): + row = ci * _ROWS_PER_ITER + row_in_iter + g_vec = cute.make_tensor( + gState_h0.iterator + strip_base + row * _D + col_base, + cute.make_layout(_VEC), + ) + s_vec = cute.make_tensor( + state_buf_ptr + row * _SMEM_STRIDE + col_base, + cute.make_layout(_VEC), + ) + cute.autovec_copy(g_vec, s_vec) + cute.arch.barrier(barrier_id=5, number_of_threads=_CE_THREADS) + # Each thread reads its V-row from padded SMEM + strip_start = strip_idx * _STRIP + if local_tidx >= strip_start: + if local_tidx < strip_start + _STRIP: + local_v = local_tidx - strip_start + s_row = cute.make_tensor( + state_buf_ptr + local_v * _SMEM_STRIDE, + cute.make_layout(_D), + ) + cute.autovec_copy(s_row, init_flat) + cute.arch.barrier(barrier_id=5, number_of_threads=_CE_THREADS) # Store raw h0 as BF16 to kv16 TMEM for SQ MMA at idx=0 tmem_store_rAccKVAsBF16.store(tTR_rKV.load().to(self.io_dtype)) @@ -1877,10 +1926,45 @@ def kernel( gState_ht = initial_state[None, None, (hidx, state_idx)] else: gState_ht = final_state[None, None, (hidx, state_idx)] - gRow_ht = cute.make_tensor(gState_ht.iterator + local_tidx, cute.make_layout(_D, stride=_D)) + # Store via SMEM transpose for coalesced GMEM writes out_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(_D)) - cute.autovec_copy(out_flat, gRow_ht) + _STRIP = 64 + _PAD = 4 + _VEC = 4 + _SMEM_STRIDE = _D + _PAD + _CE_THREADS = self.threads_per_warp * len(self.cuda_warp_ids) + _COLS_PER_THREAD = _D // _VEC + _ROWS_PER_ITER = _CE_THREADS // _COLS_PER_THREAD + _NUM_ITERS = _STRIP // _ROWS_PER_ITER + row_in_iter = local_tidx // _COLS_PER_THREAD + col_base = (local_tidx % _COLS_PER_THREAD) * _VEC + for strip_idx in cutlass.range(_D // _STRIP): + strip_base = strip_idx * (_STRIP * _D) + strip_start = strip_idx * _STRIP + # Each thread writes its V-row to padded SMEM + if local_tidx >= strip_start: + if local_tidx < strip_start + _STRIP: + local_v = local_tidx - strip_start + s_row = cute.make_tensor( + state_buf_ptr + local_v * _SMEM_STRIDE, + cute.make_layout(_D), + ) + cute.autovec_copy(out_flat, s_row) + cute.arch.barrier(barrier_id=5, number_of_threads=_CE_THREADS) + # Cooperative coalesced padded SMEM → GMEM (row-by-row) + for ci in cutlass.range(_NUM_ITERS, unroll_full=True): + row = ci * _ROWS_PER_ITER + row_in_iter + s_vec = cute.make_tensor( + state_buf_ptr + row * _SMEM_STRIDE + col_base, + cute.make_layout(_VEC), + ) + g_vec = cute.make_tensor( + gState_ht.iterator + strip_base + row * _D + col_base, + cute.make_layout(_VEC), + ) + cute.autovec_copy(s_vec, g_vec) + cute.arch.barrier(barrier_id=5, number_of_threads=_CE_THREADS) # Advance k_stage_offset by number of chunks in this WU # so next WU's k_stage_idx stays in sync with the K pipeline. @@ -2895,12 +2979,12 @@ def lightning_attn_fwd( V: (B, S, H, D) bf16 value decay: (H,) f32 per-head decay coefficients scale: attention scale factor (default: 1.0) - initial_state: (B, H, D, D) f32 initial state or None + initial_state: (B, H, D, D) f32 initial state in BHVK layout, or None output_final_state: whether to output final state chunk_size: chunk size (default: 64) Returns: - (O, ht): output tensor (B,S,H,D) bf16, final state (B,H,D,D) f32 or None + (O, ht): output tensor (B,S,H,D) bf16, final state (B,H,D,D) f32 in BHVK layout or None """ B, S, H, D = Q.shape O = torch.zeros_like(Q) @@ -3111,7 +3195,7 @@ def lightning_attn_fwd_varlen( decay: (H,) f32 per-head decay coefficients cu_seqlens: (N+1,) int32 cumulative sequence lengths scale: attention scale factor (default: 1.0) - state_pool: (pool_size, H, D, D) f32 state pool, or None + state_pool: (pool_size, H, D, D) f32 state pool in BHVK layout, or None If None, a zero state pool is allocated with pool_size=N. States are updated in-place (INPLACE_UPDATE). initial_state_indices: (N,) int32 indices into state_pool per sequence. diff --git a/tests/test_la_decode.py b/tests/test_la_decode.py index 41ee05ed..688900f7 100644 --- a/tests/test_la_decode.py +++ b/tests/test_la_decode.py @@ -83,13 +83,9 @@ def make_inputs(B, H, D, device="cuda", seed=42): def run_la_decode(q, k, v, state_4d, decay_scales, scale): """Run la_decode with proper state layout conversion.""" B, H, D, _ = state_4d.shape - # la_decode state layout: [B*H, V, K] (pretransposed) - state_cute = ( - state_4d.clone() - .permute(0, 1, 3, 2) # [B, H, V, K] - .reshape(B * H, D, D) - .contiguous() - ) + # la_decode kernel expects BHVK layout: [B*H, V, K] + # Reference/test state is BHKV: [B, H, K, V] → transpose to BHVK + state_cute = state_4d.clone().transpose(-1, -2).contiguous().reshape(B * H, D, D) out = torch.zeros(B, H, D, device=q.device, dtype=torch.bfloat16) s_offsets = torch.arange(B, device=q.device, dtype=torch.int32) @@ -111,8 +107,8 @@ def run_la_decode(q, k, v, state_4d, decay_scales, scale): K_SPLIT_DIM=D, V_SPLIT_DIM=D, ) - # Convert state back to [B, H, K, V] - state_out = state_cute.reshape(B, H, D, D).permute(0, 1, 3, 2).contiguous() + # Convert output state back from BHVK to BHKV for comparison + state_out = state_cute.reshape(B, H, D, D).transpose(-1, -2).contiguous() return out, state_out @@ -230,6 +226,45 @@ def test_vs_fla(B): max_ref = torch.abs(o_fla.float()).max().item() assert rmse / (max_ref + 1e-8) < 0.005, f"B={B}: vs fla mismatch, rel_rmse={rmse / (max_ref + 1e-8):.6f}" + # --------------------------------------------------------------------------- + + +# End-to-End Prefill -> Decode Test +# --------------------------------------------------------------------------- +def test_prefill_decode_e2e(): + """Verify prefill output state passes directly into decode without transpose.""" + from cula.ops.lightning_attn import lightning_attn_fwd + + B, S, H, D = 2, 64, 8, 128 + scale = D**-0.5 + decay_scales = 0.5 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + # Dummy prefill tokens + q_pre = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) + k_pre = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) + v_pre = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) + + # 1. Run Prefill (Generates BHVK ht) + _, ht = lightning_attn_fwd(q_pre, k_pre, v_pre, decay_scales, scale=scale, output_final_state=True) + + # Prefill outputs BHVK; convert to BHKV for reference and run_la_decode helper + ht_kv = ht.transpose(-1, -2).contiguous() + + # Dummy decode tokens + q_dec = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k_dec = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v_dec = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + + # 2. Run Decode (run_la_decode handles BHKV→BHVK internally) + out_dec, state_new = run_la_decode(q_dec, k_dec, v_dec, ht_kv, decay_scales, scale) + + # 3. Check against PyTorch reference (BHKV state) + out_ref, state_new_ref = torch_la_decode_ref(q_dec, k_dec, v_dec, ht_kv, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((out_dec.float() - out_ref.float()) ** 2)).item() + max_ref = torch.abs(out_ref.float()).max().item() + assert rmse / (max_ref + 1e-8) < 0.01, "E2E Output mismatch" + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/tests/test_lightning_attn.py b/tests/test_lightning_attn.py index 42fe2415..8958f54c 100644 --- a/tests/test_lightning_attn.py +++ b/tests/test_lightning_attn.py @@ -307,6 +307,7 @@ def test_initial_and_final_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, at V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) h0 = torch.randn(B, H, D, D, device="cuda", dtype=torch.float32) * 0.01 + h0_vk = h0.transpose(-1, -2).contiguous() # BHVK for CuTe kernel O_ref, ht_ref = pytorch_reference( Q, @@ -325,14 +326,14 @@ def test_initial_and_final_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, at V, decay, chunk_size=C, - initial_state=h0.clone(), + initial_state=h0_vk.clone(), output_final_state=True, ) p1 = _compare("output", O_cute, O_ref_bf16, atol=atol, rtol=rtol, verbose=verbose) p2 = True if ht_ref is not None and ht_cute is not None: - p2 = _compare("state", ht_cute, ht_ref, atol=atol, rtol=rtol, verbose=verbose) + p2 = _compare("state", ht_cute.transpose(-1, -2), ht_ref, atol=atol, rtol=rtol, verbose=verbose) passed = p1 and p2 print(f" {'✓ PASSED' if passed else '✗ FAILED'}") @@ -396,11 +397,12 @@ def test_against_fla_with_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, ato K = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 h0 = torch.randn(B, H, D, D, device="cuda", dtype=torch.float32) * 0.01 + h0_vk = h0.transpose(-1, -2).contiguous() # BHVK for CuTe kernel decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) g_gamma = -decay - # FLA + # FLA (expects BHKV state) O_fla, ht_fla = chunk_simple_gla( Q, K, @@ -412,7 +414,7 @@ def test_against_fla_with_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, ato head_first=False, ) - # Ours + # Ours (expects BHVK state) O_cute, ht_cute = run_cute_kernel( Q, K, @@ -420,14 +422,14 @@ def test_against_fla_with_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, ato decay, scale=1.0, chunk_size=C, - initial_state=h0.clone(), + initial_state=h0_vk.clone(), output_final_state=True, ) p1 = _compare("output", O_cute, O_fla, atol=atol, rtol=rtol, verbose=verbose) p2 = True if ht_fla is not None and ht_cute is not None: - p2 = _compare("state", ht_cute, ht_fla, atol=atol, rtol=rtol, verbose=verbose) + p2 = _compare("state", ht_cute.transpose(-1, -2), ht_fla, atol=atol, rtol=rtol, verbose=verbose) passed = p1 and p2 print(f" {'✓ PASSED' if passed else '✗ FAILED'}") @@ -523,9 +525,9 @@ def test_varlen_with_initial_state(seq_lens=None, H=4, D=128, C=64, decay_val=0. V = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) - # State pool with 3 slots, use indices [2, 0] + # State pool with 3 slots, use indices [2, 0] — BHVK layout for CuTe pool_size = 3 - state_pool = torch.randn(pool_size, H, D, D, dtype=torch.float32, device="cuda") * 0.01 + state_pool = torch.randn(pool_size, H, D, D, dtype=torch.float32, device="cuda").transpose(-1, -2).contiguous() * 0.01 indices = torch.tensor([2, 0], dtype=torch.int32, device="cuda") O_var, sp = run_cute_kernel_varlen( @@ -591,6 +593,7 @@ def test_varlen_against_pytorch_ref( decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) state_pool = torch.randn(N, H, D, D, dtype=torch.float32, device="cuda") * 0.01 + state_pool_vk = state_pool.transpose(-1, -2).contiguous() # BHVK for CuTe O_var, sp = run_cute_kernel_varlen( Q, @@ -599,7 +602,7 @@ def test_varlen_against_pytorch_ref( decay, cu_seqlens, chunk_size=C, - state_pool=state_pool.clone(), + state_pool=state_pool_vk.clone(), ) all_pass = True @@ -609,7 +612,7 @@ def test_varlen_against_pytorch_ref( Qi = Q[:, bos:eos].contiguous() Ki = K[:, bos:eos].contiguous() Vi = V[:, bos:eos].contiguous() - h0_i = state_pool[i : i + 1].clone() + h0_i = state_pool[i : i + 1].clone() # BHKV for pytorch_reference O_ref_i, ht_ref_i = pytorch_reference( Qi, @@ -624,7 +627,7 @@ def test_varlen_against_pytorch_ref( O_ref_bf16 = O_ref_i.to(torch.bfloat16) po = _compare(f"O[{i}]", O_var[:, bos:eos], O_ref_bf16, atol=atol, rtol=rtol, verbose=verbose) - ps = _compare(f"ht[{i}]", sp[i], ht_ref_i, atol=atol, rtol=rtol, verbose=verbose) + ps = _compare(f"ht[{i}]", sp[i].transpose(-1, -2), ht_ref_i, atol=atol, rtol=rtol, verbose=verbose) all_pass = all_pass and po and ps print(f" {'✓ PASSED' if all_pass else '✗ FAILED'}") From 2c8c6fa294aca03901c44feff0cdeff403fbdb88 Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Thu, 23 Apr 2026 10:32:09 +0800 Subject: [PATCH 06/34] perf(la_decode): boost small-batch (B<=32) by 33% via 8-warp CTA (#59) * perf(la_decode): boost small-batch (B<=32) by 33% via 8-warp CTA The small-batch path (B<=32) splits each batch's state into 8 sub-blocks to grow grid for more SMs. With NUM_THREADS=128 (4 warps), per-CTA only 4 warps were active and SM occupancy stayed at 38%. Profiling showed B<=32 was occupancy-bound (low warps/SM), not HBM bound (DRAM only ~34% peak). Bumping the small-batch path to NUM_THREADS=256 (8 warps/CTA) doubles per-CTA warps without changing grid, lifting occupancy 38%->85% at B=32. Each warp now processes 1 row of TILE_V instead of 2, balancing the work. The big-batch path (B>32) keeps NUM_THREADS=128 since adding warps there hurts due to redundant Q/K/V loads per warp. NCU verified at B=32: gpu_time: 27.78 -> 20.90 us (1.33x) occupancy: 38.4% -> 84.6% short_scoreboard: 5.0 -> 2.81 End-to-end timings (us): B | before | after 1 | ~28 | 5.92 8 | ~28 | 9.41 16 | ~28 | 13.76 32 | 27.78 | 20.90 64 | 27.78 | 27.68 (no change, big-path) 128 | 49.92 | 49.86 256 | 91.20 | 92.16 All 19 tests in tests/test_la_decode.py pass. * perf(la_decode): revert small-batch to 4 warps (sweep showed regression) Earlier commit 36eb38b bundled two changes: (1) bumped small-batch path NUM_THREADS to 256 (8 warps), and (2) refactored constexpr loops, batched LDS, and v_src indexing. A subsequent NUM_THREADS sweep across B in {1,8,16,32,64,128,256} shows that change (1) is now a regression in current code state -- the latency optimizations from (2) reduced LDS stalls so much that 4 warps already saturate them, and 8 warps just adds barrier and scheduling overhead. Sweep results (3 runs, min, us, B=32): small=8 warps (prev): 18.48 small=4 warps (now): 16.54 (-10.5%) End-to-end timings vs prior commit (us): B | prev | now | delta 1 | 5.92 | 4.85 | +22% 8 | 9.41 | 8.20 | +15% 16 | 13.76 | 10.30 | +34% 32 | 20.90 | 16.44 | +27% 64 | 27.68 | 24.77 | tie (+/- noise) 128 | 49.86 | 49.92 | tie 256 | 92.16 | 91.43 | tie Other configs swept (no win): NUM_STAGES in {3,4}, big-path 8 warps, NUM_BLOCKS_PER_STATE in {4,16}. All 19 tests in tests/test_la_decode.py pass. * perf(la_decode): split TILE_V/NUM_STAGES per path (small=TV4/S3/4w, big=TV32/S3/8w) Sweep on GB200 (CUDA_VISIBLE_DEVICES=1/2, min over 6 runs, B*H=B*16): B baseline (TV8/S2, 4w/4w) per-path (TV4-TV32/S3, 4w/8w) speedup 1 10.73 us 10.08 us +6.5% 8 10.67 us 10.34 us +3.2% 16 10.41 us 10.47 us ~tie 32 10.97 us 10.54 us +4.1% 64 14.54 us 14.52 us ~tie 128 46.86 us 43.47 us +7.8% 256 87.28 us 82.09 us +6.3% Sweep showed the small-batch path (B<=32, NUM_BLOCKS_PER_STATE=8 expansion) prefers the smaller tile (TV=4, more cp.async hiding via 4 v-tiles per CTA with NUM_STAGES=3) while the big-batch path (B>32, whole state per CTA) prefers the larger tile (TV=32, fewer v-tiles, 8 warps issuing in parallel, NUM_STAGES=3 still fits). Refactored the two globals into per-path constants TILE_V_SMALL/TILE_V_BIG and NUM_STAGES_SMALL/NUM_STAGES_BIG; the kernels now take TILE_V/NUM_STAGES as constexpr parameters and v_grp = v_tiles // (32 // TILE_V) so the swizzled-store path stays valid for any TILE_V in {4,8,16,32}. Tests: pytest tests/test_la_decode.py -> 19/19 pass. Follow-up to a1895b0 and dd52b1b. --- cula/ops/la_decode.py | 144 +++++++++++++++++++++++++++++------------- 1 file changed, 100 insertions(+), 44 deletions(-) diff --git a/cula/ops/la_decode.py b/cula/ops/la_decode.py index ce40abf3..bd6f1b8c 100644 --- a/cula/ops/la_decode.py +++ b/cula/ops/la_decode.py @@ -19,10 +19,10 @@ Linear Attention updates during decode phase. Simplified compared to GDN. Architecture Design: -- Uses TMA (Tensor Memory Accelerator) for efficient Global Memory → Shared Memory transfers -- Employs 2-stage pipeline to overlap loading and computation, hiding memory latency -- Each block uses 128 threads (4 warps), with each warp processing one matrix row -- Tile size: 8x128 (TILE_V x TILE_K) +- Uses cp.async for efficient Global Memory → Shared Memory transfers +- Employs N-stage pipeline to overlap loading and computation, hiding memory latency +- Per-path config (sweep-tuned on GB200): small path uses TILE_V=4 / 4 warps, + big path uses TILE_V=32 / 8 warps; both run with NUM_STAGES=3. Computation Flow: 1. Warp 0 handles TMA prefetch, loading data from GMEM to SMEM @@ -50,12 +50,24 @@ # ============================================================================ # Global configuration # ============================================================================ -TILE_V = 8 +# Per-path tile / pipeline / warp configuration was sweep-tuned on GB200. +# Small path (B <= 32, grid expanded by NUM_BLOCKS_PER_STATE): +# smaller TILE_V keeps the grid wide, deeper NUM_STAGES hides cp.async latency. +# Big path (B > 32, grid = B*H): +# larger TILE_V amortizes per-warp setup, 8 warps fill issue slots, deeper stages help. +TILE_V_SMALL = 4 +TILE_V_BIG = 32 TILE_K = 128 -NUM_STAGES = 2 -NUM_THREADS = 128 # 4 warps +NUM_STAGES_SMALL = 3 +NUM_STAGES_BIG = 3 +NUM_THREADS_BIG = 256 # 8 warps for big-batch path (B > 32) +NUM_THREADS_SMALL = 128 # 4 warps for small-batch path (B <= 32) NUM_BLOCKS_PER_STATE = 8 +# Backward-compat aliases (still referenced in some places). +TILE_V = TILE_V_SMALL +NUM_STAGES = NUM_STAGES_SMALL + @cute.kernel def la_decode_kernel_small_batch_pretranspose( @@ -76,6 +88,9 @@ def la_decode_kernel_small_batch_pretranspose( H: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], + NUM_WARPS: cutlass.Constexpr[int] = 4, + TILE_V: cutlass.Constexpr[int] = 8, + NUM_STAGES: cutlass.Constexpr[int] = 2, ): """Each block uses pipeline to load one batch and vectorized writeback""" @@ -138,7 +153,7 @@ def la_decode_kernel_small_batch_pretranspose( cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() - for i in range(vec_size): + for i in cutlass.range_constexpr(vec_size): r_q[i] = cutlass.Float32(q[i_n, i_h, i * 32 + lane_id]) r_k[i] = cutlass.Float32(k[i_n, i_h, i * 32 + lane_id]) r_v[i] = cutlass.Float32(v[i_n, i_hv, i * 32 + lane_id]) @@ -146,7 +161,7 @@ def la_decode_kernel_small_batch_pretranspose( cute.arch.barrier() # Ensure all threads finish writing to sV # Apply scaling in Float32 - for i in range(vec_size): + for i in cutlass.range_constexpr(vec_size): r_q[i] = r_q[i] * scale # =================================================================== @@ -175,17 +190,33 @@ def la_decode_kernel_small_batch_pretranspose( cute.arch.cp_async_commit_group() # Step 3: Compute using data from current stage - for row in range(0, TILE_V, 4): + # v_grp selects which r_v[] entry holds this v_tiles' v values. + # r_v has vec_size=4 entries each holding 32 V positions; so each + # entry covers (32 / TILE_V) consecutive v_tiles. (e.g. TILE_V=8 -> 4) + v_src = cutlass.Float32(0.0) + v_grp = v_tiles // (32 // TILE_V) + if v_grp == 0: + v_src = r_v[0] + elif v_grp == 1: + v_src = r_v[1] + elif v_grp == 2: + v_src = r_v[2] + else: + v_src = r_v[3] + + for row in cutlass.range_constexpr(0, TILE_V, NUM_WARPS): row_offset = tidx // 32 v_idx = v_tiles * TILE_V + row + row_offset - v_row = cute.arch.shuffle_sync(r_v[v_idx // 32], v_idx % 32, mask=-1, mask_and_clamp=31) + v_row = cute.arch.shuffle_sync(v_src, v_idx % 32, mask=-1, mask_and_clamp=31) sum_hq = 0.0 - for i in range(vec_size): + # Batch all SMEM loads first to overlap LDS latency (short_scoreboard fix) + for i in cutlass.range_constexpr(vec_size): r_h[i] = sData[(row + row_offset, i * 32 + lane_id, stage)] - r_h[i] = r_h[i] * r_decay - r_h[i] += r_k[i] * v_row + # Then consume them in FMAs / stores + for i in cutlass.range_constexpr(vec_size): + r_h[i] = r_h[i] * r_decay + r_k[i] * v_row gDst[(0, row + row_offset, i * 32 + lane_id, v_tiles)] = r_h[i] sum_hq += r_h[i] * r_q[i] @@ -224,6 +255,9 @@ def la_decode_kernel_big_batch_pretranspose( H: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], + NUM_WARPS: cutlass.Constexpr[int] = 4, + TILE_V: cutlass.Constexpr[int] = 8, + NUM_STAGES: cutlass.Constexpr[int] = 2, ): """Each block uses pipeline to load one batch and vectorized writeback""" @@ -280,7 +314,7 @@ def la_decode_kernel_big_batch_pretranspose( cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() - for i in range(vec_size): + for i in cutlass.range_constexpr(vec_size): r_q[i] = cutlass.Float32(q[i_n, i_h, i * 32 + lane_id]) r_k[i] = cutlass.Float32(k[i_n, i_h, i * 32 + lane_id]) r_v[i] = cutlass.Float32(v[i_n, i_hv, i * 32 + lane_id]) @@ -291,7 +325,7 @@ def la_decode_kernel_big_batch_pretranspose( # Compute g and beta (scalar values) # =================================================================== # Apply scaling in Float32 - for i in range(vec_size): + for i in cutlass.range_constexpr(vec_size): r_q[i] = r_q[i] * scale r_g = cute.exp(-cutlass.Float32(decay_scales[i_h]), fastmath=USE_FAST_MATH) @@ -321,17 +355,31 @@ def la_decode_kernel_big_batch_pretranspose( cute.arch.cp_async_commit_group() # Step 3: Compute using data from current stage - for row in range(0, TILE_V, 4): + # v_grp selects which r_v[] entry holds this v_tiles' v values. + v_src = cutlass.Float32(0.0) + v_grp = v_tiles // (32 // TILE_V) + if v_grp == 0: + v_src = r_v[0] + elif v_grp == 1: + v_src = r_v[1] + elif v_grp == 2: + v_src = r_v[2] + else: + v_src = r_v[3] + + for row in cutlass.range_constexpr(0, TILE_V, NUM_WARPS): row_offset = tidx // 32 v_idx = v_tiles * TILE_V + row + row_offset - v_row = cute.arch.shuffle_sync(r_v[v_idx // 32], v_idx % 32, mask=-1, mask_and_clamp=31) + v_row = cute.arch.shuffle_sync(v_src, v_idx % 32, mask=-1, mask_and_clamp=31) sum_hq = 0.0 - for i in range(vec_size): + # Batch all SMEM loads first to overlap LDS latency (short_scoreboard fix) + for i in cutlass.range_constexpr(vec_size): r_h[i] = sData[(row + row_offset, i * 32 + lane_id, stage)] - r_h[i] = r_h[i] * r_g - r_h[i] += r_k[i] * v_row + # Then consume them in FMAs / stores + for i in cutlass.range_constexpr(vec_size): + r_h[i] = r_h[i] * r_g + r_k[i] * v_row gDst[(0, row + row_offset, i * 32 + lane_id, v_tiles)] = r_h[i] sum_hq += r_h[i] * r_q[i] @@ -383,31 +431,29 @@ def run_la_decode_kernel_big_batch_pretranspose( num_bits_per_copy=128, # 4 elements per copy ) - # Thread layout: 4 rows × 32 threads/row = 128 threads + # Thread layout: NUM_WARPS_BIG rows × 32 threads/row + NUM_WARPS_BIG = NUM_THREADS_BIG // 32 thread_layout = cute.make_layout( - (4, 32), # 4 rows, 32 threads/row + (NUM_WARPS_BIG, 32), stride=(32, 1), ) val_layout = cute.make_layout((1, 4)) # Each thread handles 4 elements tiled_copy_load = cute.make_tiled_copy_tv(copy_atom, thread_layout, val_layout) - num_v_tiles = cute.ceil_div(v_dim, TILE_V) + num_v_tiles = cute.ceil_div(v_dim, TILE_V_BIG) vec_size = TILE_K // 32 # Each thread in a warp processes this many elements (always 4 for TILE_K=128) - # print(f"Batched CP.ASYNC Load + Store (bypass L1 cache)") - # print(f" {batch_size} batches x {v_dim}x{k_dim} matrices") - # print(f" Tile: {TILE_V}x{TILE_K}, {num_v_tiles} tiles/batch") - # print(f" Threads: {NUM_THREADS} ({NUM_THREADS // 32} warps), vec_size: {vec_size}") - # print(f" Total: {total_data_mb:.1f} MB\n") - - # Create SMEM layout - smem_layout_staged = cute.make_layout((TILE_V, TILE_K, NUM_STAGES), stride=(TILE_K, 1, TILE_V * TILE_K)) + # Create SMEM layout (row-major: v-row major, k-contiguous) + smem_layout_staged = cute.make_layout( + (TILE_V_BIG, TILE_K, NUM_STAGES_BIG), + stride=(TILE_K, 1, TILE_V_BIG * TILE_K), + ) - # sData: TILE_V * TILE_K * NUM_STAGES * 4 bytes (Float32) + # sData: TILE_V_BIG * TILE_K * NUM_STAGES_BIG * 4 bytes (Float32) # sOutput: V * 2 bytes (BFloat16) - smem_bytes = 4 * TILE_V * TILE_K * NUM_STAGES + 2 * v_dim + 32 + smem_bytes = 4 * TILE_V_BIG * TILE_K * NUM_STAGES_BIG + 2 * v_dim + 32 la_decode_kernel_big_batch_pretranspose( tiled_copy_load, @@ -427,9 +473,12 @@ def run_la_decode_kernel_big_batch_pretranspose( H, K, V, + NUM_WARPS_BIG, + TILE_V_BIG, + NUM_STAGES_BIG, ).launch( grid=(batch_size, 1, 1), - block=[NUM_THREADS, 1, 1], + block=[NUM_THREADS_BIG, 1, 1], smem=smem_bytes, stream=stream, ) @@ -466,25 +515,29 @@ def run_la_decode_kernel_small_batch_pretranspose( num_bits_per_copy=128, # 4 elements per copy ) - # Thread layout: 4 rows × 32 threads/row = 128 threads + # Thread layout: NUM_WARPS_SMALL rows × 32 threads/row + NUM_WARPS_SMALL = NUM_THREADS_SMALL // 32 thread_layout = cute.make_layout( - (4, 32), # 4 rows, 32 threads/row + (NUM_WARPS_SMALL, 32), stride=(32, 1), ) val_layout = cute.make_layout((1, 4)) # Each thread handles 4 elements tiled_copy_load = cute.make_tiled_copy_tv(copy_atom, thread_layout, val_layout) - num_v_tiles = cute.ceil_div(v_dim, TILE_V) + num_v_tiles = cute.ceil_div(v_dim, TILE_V_SMALL) vec_size = TILE_K // 32 # Each thread in a warp processes this many elements (always 4 for TILE_K=128) - # Create SMEM layout - smem_layout_staged = cute.make_layout((TILE_V, TILE_K, NUM_STAGES), stride=(TILE_K, 1, TILE_V * TILE_K)) + # Create SMEM layout (row-major: v-row major, k-contiguous) + smem_layout_staged = cute.make_layout( + (TILE_V_SMALL, TILE_K, NUM_STAGES_SMALL), + stride=(TILE_K, 1, TILE_V_SMALL * TILE_K), + ) - # sData: TILE_V * TILE_K * NUM_STAGES * 4 bytes (Float32) - # sOutput: TILE_V * 2 bytes (BFloat16) - smem_bytes = 4 * TILE_V * TILE_K * NUM_STAGES + 2 * v_dim + 32 + # sData: TILE_V_SMALL * TILE_K * NUM_STAGES_SMALL * 4 bytes (Float32) + # sOutput: V * 2 bytes (BFloat16) + smem_bytes = 4 * TILE_V_SMALL * TILE_K * NUM_STAGES_SMALL + 2 * v_dim + 32 la_decode_kernel_small_batch_pretranspose( tiled_copy_load, @@ -504,9 +557,12 @@ def run_la_decode_kernel_small_batch_pretranspose( H, K, V, + NUM_WARPS_SMALL, + TILE_V_SMALL, + NUM_STAGES_SMALL, ).launch( grid=(batch_size * NUM_BLOCKS_PER_STATE, 1, 1), - block=[NUM_THREADS, 1, 1], + block=[NUM_THREADS_SMALL, 1, 1], smem=smem_bytes, stream=stream, ) From fa742a23fcc6d818b0f591532ecad7985e3f64e5 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Fri, 24 Apr 2026 10:27:59 +0800 Subject: [PATCH 07/34] [KDA] Optimize recompute_wu with better register allocation (#61) * optimize recompute_wu with better reg allocation * add cuda130 bench, add reg config for different nvcc version * update cuda129 bench --- BENCHMARK_GB200.md | 188 +++++++++--------- BENCHMARK_GB200_CUDA_130.md | 163 +++++++++++++++ .../sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp | 14 +- .../kda_fwd_recomp_w_u_mainloop_sm100.hpp | 18 +- .../kerutils/include/kerutils/device/common.h | 4 + 5 files changed, 280 insertions(+), 107 deletions(-) create mode 100644 BENCHMARK_GB200_CUDA_130.md diff --git a/BENCHMARK_GB200.md b/BENCHMARK_GB200.md index a797507c..726a5e84 100644 --- a/BENCHMARK_GB200.md +++ b/BENCHMARK_GB200.md @@ -1,6 +1,6 @@ # Benchmark Results -> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-04-05. +> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-04-23. > **GPU:** NVIDIA GB200 | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 @@ -14,44 +14,44 @@ | B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------|---------| -| 1 | 512 | 0.518 | 0.417 | **1.24x** | -| 1 | 1024 | 0.508 | 0.412 | **1.23x** | -| 1 | 4096 | 0.751 | 0.537 | **1.40x** | -| 1 | 8192 | 1.395 | 0.993 | **1.40x** | -| 1 | 16384 | 2.736 | 1.919 | **1.43x** | -| 2 | 512 | 0.510 | 0.418 | **1.22x** | -| 2 | 1024 | 0.516 | 0.418 | **1.23x** | -| 2 | 4096 | 1.392 | 1.000 | **1.39x** | -| 2 | 8192 | 2.726 | 1.934 | **1.41x** | -| 2 | 16384 | 5.346 | 3.832 | **1.40x** | +| 1 | 512 | 0.566 | 0.449 | **1.26x** | +| 1 | 1024 | 0.539 | 0.445 | **1.21x** | +| 1 | 4096 | 0.748 | 0.531 | **1.41x** | +| 1 | 8192 | 1.392 | 0.984 | **1.42x** | +| 1 | 16384 | 2.706 | 1.883 | **1.44x** | +| 2 | 512 | 0.568 | 0.460 | **1.24x** | +| 2 | 1024 | 0.556 | 0.452 | **1.23x** | +| 2 | 4096 | 1.389 | 0.990 | **1.40x** | +| 2 | 8192 | 2.701 | 1.899 | **1.42x** | +| 2 | 16384 | 5.307 | 3.766 | **1.41x** | -Summary (10 configs): **avg=1.34x**, min=1.22x, max=1.43x. +Summary (10 configs): **avg=1.34x**, min=1.21x, max=1.44x. ### Variable-Length (H=64, D=128, bf16) | Config | FLA Triton (ms) | cuLA (ms) | Speedup | |--------|-----------------|-----------|---------| -| uniform 10seqs T=4096 [409..415] avg=409 | 0.788 | 0.583 | **1.35x** | -| random 10seqs T=4096 [24..1201] avg=409 | 0.784 | 0.575 | **1.36x** | -| skewed 10seqs T=4096 [227..2053] avg=409 | 0.782 | 0.573 | **1.36x** | -| uniform 20seqs T=4096 [204..220] avg=204 | 0.863 | 0.635 | **1.36x** | -| random 20seqs T=4096 [5..787] avg=204 | 0.835 | 0.617 | **1.35x** | -| skewed 20seqs T=4096 [107..2063] avg=204 | 0.818 | 0.594 | **1.38x** | -| uniform 10seqs T=8192 [819..821] avg=819 | 1.395 | 1.018 | **1.37x** | -| random 10seqs T=8192 [48..2401] avg=819 | 1.421 | 1.039 | **1.37x** | -| skewed 10seqs T=8192 [455..4097] avg=819 | 1.444 | 1.040 | **1.39x** | -| uniform 20seqs T=8192 [409..421] avg=409 | 1.479 | 1.064 | **1.39x** | -| random 20seqs T=8192 [9..1574] avg=409 | 1.482 | 1.066 | **1.39x** | -| skewed 20seqs T=8192 [215..4107] avg=409 | 1.488 | 1.071 | **1.39x** | -| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.697 | 1.952 | **1.38x** | -| random 10seqs T=16384 [95..4802] avg=1638 | 2.691 | 1.946 | **1.38x** | -| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.712 | 1.960 | **1.38x** | -| uniform 20seqs T=16384 [819..823] avg=819 | 2.698 | 1.948 | **1.39x** | -| random 20seqs T=16384 [19..3147] avg=819 | 2.728 | 1.976 | **1.38x** | -| skewed 20seqs T=16384 [431..8195] avg=819 | 2.703 | 1.947 | **1.39x** | - -Summary (18 configs): **avg=1.38x**, min=1.35x, max=1.39x. +| uniform 10seqs T=4096 [409..415] avg=409 | 0.785 | 0.573 | **1.37x** | +| random 10seqs T=4096 [24..1201] avg=409 | 0.779 | 0.567 | **1.37x** | +| skewed 10seqs T=4096 [227..2053] avg=409 | 0.778 | 0.565 | **1.38x** | +| uniform 20seqs T=4096 [204..220] avg=204 | 0.858 | 0.622 | **1.38x** | +| random 20seqs T=4096 [5..787] avg=204 | 0.830 | 0.606 | **1.37x** | +| skewed 20seqs T=4096 [107..2063] avg=204 | 0.813 | 0.586 | **1.39x** | +| uniform 10seqs T=8192 [819..821] avg=819 | 1.388 | 1.004 | **1.38x** | +| random 10seqs T=8192 [48..2401] avg=819 | 1.412 | 1.024 | **1.38x** | +| skewed 10seqs T=8192 [455..4097] avg=819 | 1.439 | 1.028 | **1.40x** | +| uniform 20seqs T=8192 [409..421] avg=409 | 1.477 | 1.052 | **1.40x** | +| random 20seqs T=8192 [9..1574] avg=409 | 1.474 | 1.056 | **1.40x** | +| skewed 20seqs T=8192 [215..4107] avg=409 | 1.483 | 1.061 | **1.40x** | +| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.670 | 1.915 | **1.39x** | +| random 10seqs T=16384 [95..4802] avg=1638 | 2.679 | 1.916 | **1.40x** | +| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.683 | 1.916 | **1.40x** | +| uniform 20seqs T=16384 [819..823] avg=819 | 2.679 | 1.913 | **1.40x** | +| random 20seqs T=16384 [19..3147] avg=819 | 2.713 | 1.937 | **1.40x** | +| skewed 20seqs T=16384 [431..8195] avg=819 | 2.687 | 1.918 | **1.40x** | + +Summary (18 configs): **avg=1.39x**, min=1.37x, max=1.40x. To reproduce: @@ -66,14 +66,14 @@ python benchmarks/bench_kda.py --mode both | B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------|---------| -| 1 | 1024 | 0.087 | 0.065 | **1.33x** | -| 1 | 4096 | 0.172 | 0.157 | **1.10x** | -| 1 | 8192 | 0.323 | 0.292 | **1.10x** | -| 1 | 16384 | 0.618 | 0.563 | **1.10x** | -| 2 | 1024 | 0.097 | 0.061 | **1.60x** | -| 2 | 4096 | 0.319 | 0.175 | **1.82x** | -| 2 | 8192 | 0.611 | 0.329 | **1.86x** | -| 2 | 16384 | 1.192 | 0.632 | **1.89x** | +| 1 | 1024 | 0.100 | 0.071 | **1.41x** | +| 1 | 4096 | 0.174 | 0.157 | **1.11x** | +| 1 | 8192 | 0.331 | 0.293 | **1.13x** | +| 1 | 16384 | 0.637 | 0.563 | **1.13x** | +| 2 | 1024 | 0.093 | 0.062 | **1.51x** | +| 2 | 4096 | 0.305 | 0.176 | **1.73x** | +| 2 | 8192 | 0.585 | 0.328 | **1.79x** | +| 2 | 16384 | 1.140 | 0.632 | **1.80x** | ### Variable-Length (H=64, D=128, bf16) @@ -81,50 +81,50 @@ Persistent CuTe DSL kernel vs FLA Triton varlen. | N (seqs) | T | cuLA (ms) | FLA Triton (ms) | Speedup | |----------|---|-----------|-----------------|---------| -| 5 | 1020 | 0.079 | 0.156 | **1.96x** | -| 5 | 2045 | 0.102 | 0.183 | **1.79x** | -| 5 | 4095 | 0.153 | 0.241 | **1.58x** | -| 5 | 8190 | 0.252 | 0.390 | **1.55x** | -| 5 | 16380 | 0.448 | 0.685 | **1.53x** | -| 5 | 32765 | 0.843 | 1.281 | **1.52x** | -| 8 | 1024 | 0.076 | 0.142 | **1.88x** | -| 8 | 2048 | 0.101 | 0.168 | **1.67x** | -| 8 | 4096 | 0.145 | 0.236 | **1.63x** | -| 8 | 8192 | 0.228 | 0.385 | **1.69x** | -| 8 | 16384 | 0.399 | 0.668 | **1.68x** | -| 8 | 32768 | 0.741 | 1.236 | **1.67x** | -| 10 | 1020 | 0.092 | 0.148 | **1.61x** | -| 10 | 2040 | 0.120 | 0.178 | **1.49x** | -| 10 | 4090 | 0.164 | 0.248 | **1.51x** | -| 10 | 8190 | 0.248 | 0.391 | **1.58x** | -| 10 | 16380 | 0.419 | 0.676 | **1.62x** | -| 10 | 32760 | 0.763 | 1.253 | **1.64x** | -| 12 | 1020 | 0.104 | 0.148 | **1.42x** | -| 12 | 2040 | 0.122 | 0.172 | **1.41x** | -| 12 | 4092 | 0.174 | 0.251 | **1.45x** | -| 12 | 8184 | 0.256 | 0.392 | **1.53x** | -| 12 | 16380 | 0.429 | 0.682 | **1.59x** | -| 12 | 32760 | 0.769 | 1.244 | **1.62x** | -| 16 | 1024 | 0.103 | 0.140 | **1.35x** | -| 16 | 2048 | 0.130 | 0.171 | **1.32x** | -| 16 | 4096 | 0.171 | 0.243 | **1.43x** | -| 16 | 8192 | 0.245 | 0.385 | **1.57x** | -| 16 | 16384 | 0.402 | 0.668 | **1.66x** | -| 16 | 32768 | 0.722 | 1.233 | **1.71x** | -| 20 | 1020 | 0.140 | 0.149 | **1.07x** | -| 20 | 2040 | 0.170 | 0.185 | **1.09x** | -| 20 | 4080 | 0.211 | 0.272 | **1.29x** | -| 20 | 8180 | 0.292 | 0.412 | **1.41x** | -| 20 | 16380 | 0.448 | 0.682 | **1.52x** | -| 20 | 32760 | 0.770 | 1.247 | **1.62x** | -| 25 | 1000 | 0.167 | 0.157 | **0.94x** | -| 25 | 2025 | 0.197 | 0.210 | **1.07x** | -| 25 | 4075 | 0.232 | 0.268 | **1.16x** | -| 25 | 8175 | 0.318 | 0.429 | **1.35x** | -| 25 | 16375 | 0.479 | 0.704 | **1.47x** | -| 25 | 32750 | 0.794 | 1.257 | **1.58x** | - -Summary (126 configs across uniform/skewed/random): **avg=1.52x**, min=0.94x, max=2.08x. +| 5 | 1020 | 0.085 | 0.164 | **1.92x** | +| 5 | 2045 | 0.112 | 0.189 | **1.68x** | +| 5 | 4095 | 0.164 | 0.250 | **1.53x** | +| 5 | 8190 | 0.265 | 0.400 | **1.51x** | +| 5 | 16380 | 0.465 | 0.692 | **1.49x** | +| 5 | 32765 | 0.860 | 1.278 | **1.49x** | +| 8 | 1024 | 0.085 | 0.153 | **1.79x** | +| 8 | 2048 | 0.113 | 0.188 | **1.67x** | +| 8 | 4096 | 0.158 | 0.245 | **1.55x** | +| 8 | 8192 | 0.243 | 0.390 | **1.60x** | +| 8 | 16384 | 0.412 | 0.679 | **1.65x** | +| 8 | 32768 | 0.759 | 1.252 | **1.65x** | +| 10 | 1020 | 0.106 | 0.161 | **1.52x** | +| 10 | 2040 | 0.135 | 0.192 | **1.43x** | +| 10 | 4090 | 0.182 | 0.254 | **1.39x** | +| 10 | 8190 | 0.267 | 0.398 | **1.49x** | +| 10 | 16380 | 0.441 | 0.684 | **1.55x** | +| 10 | 32760 | 0.791 | 1.264 | **1.60x** | +| 12 | 1020 | 0.119 | 0.174 | **1.46x** | +| 12 | 2040 | 0.146 | 0.191 | **1.31x** | +| 12 | 4092 | 0.192 | 0.259 | **1.35x** | +| 12 | 8184 | 0.278 | 0.398 | **1.43x** | +| 12 | 16380 | 0.456 | 0.703 | **1.54x** | +| 12 | 32760 | 0.796 | 1.269 | **1.60x** | +| 16 | 1024 | 0.124 | 0.168 | **1.36x** | +| 16 | 2048 | 0.151 | 0.195 | **1.29x** | +| 16 | 4096 | 0.189 | 0.261 | **1.38x** | +| 16 | 8192 | 0.270 | 0.408 | **1.51x** | +| 16 | 16384 | 0.426 | 0.688 | **1.61x** | +| 16 | 32768 | 0.743 | 1.251 | **1.68x** | +| 20 | 1020 | 0.164 | 0.175 | **1.07x** | +| 20 | 2040 | 0.193 | 0.205 | **1.06x** | +| 20 | 4080 | 0.236 | 0.277 | **1.17x** | +| 20 | 8180 | 0.320 | 0.417 | **1.30x** | +| 20 | 16380 | 0.483 | 0.702 | **1.45x** | +| 20 | 32760 | 0.806 | 1.414 | **1.75x** | +| 25 | 1000 | 0.194 | 0.181 | **0.93x** | +| 25 | 2025 | 0.223 | 0.218 | **0.98x** | +| 25 | 4075 | 0.262 | 0.275 | **1.05x** | +| 25 | 8175 | 0.350 | 0.444 | **1.27x** | +| 25 | 16375 | 0.522 | 0.716 | **1.37x** | +| 25 | 32750 | 0.836 | 1.266 | **1.51x** | + +Summary (126 configs across uniform/skewed/random): **avg=1.48x**, min=0.93x, max=3.84x. To reproduce: @@ -140,21 +140,21 @@ Single-token decode: la_decode (CuTe DSL) vs fla fused_recurrent (Triton). | B | FLA Triton (ms) | cuLA (ms) | Speedup | |---|-----------------|-----------|---------| -| 1 | 0.0602 | 0.0134 | **4.48x** | -| 4 | 0.0587 | 0.0135 | **4.36x** | -| 16 | 0.0624 | 0.0221 | **2.82x** | -| 64 | 0.0996 | 0.0924 | **1.08x** | -| 256 | 0.3493 | 0.3396 | **1.03x** | +| 1 | 0.0645 | 0.0129 | **5.01x** | +| 4 | 0.0661 | 0.0129 | **5.11x** | +| 16 | 0.0664 | 0.0210 | **3.16x** | +| 64 | 0.0997 | 0.0843 | **1.18x** | +| 256 | 0.3489 | 0.3134 | **1.11x** | #### Wrapper (Full Call Path) | B | FLA Triton (ms) | cuLA (ms) | Speedup | |---|-----------------|-----------|---------| -| 1 | 0.0771 | 0.0178 | **4.33x** | -| 4 | 0.0759 | 0.0176 | **4.32x** | -| 16 | 0.0822 | 0.0224 | **3.68x** | -| 64 | 0.0995 | 0.0931 | **1.07x** | -| 256 | 0.3485 | 0.3401 | **1.02x** | +| 1 | 0.0864 | 0.0179 | **4.82x** | +| 4 | 0.0882 | 0.0179 | **4.92x** | +| 16 | 0.0872 | 0.0218 | **4.01x** | +| 64 | 0.0997 | 0.0845 | **1.18x** | +| 256 | 0.3500 | 0.3121 | **1.12x** | To reproduce: diff --git a/BENCHMARK_GB200_CUDA_130.md b/BENCHMARK_GB200_CUDA_130.md new file mode 100644 index 00000000..a34b0a28 --- /dev/null +++ b/BENCHMARK_GB200_CUDA_130.md @@ -0,0 +1,163 @@ +# Benchmark Results + +> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-04-23. + +> **GPU:** NVIDIA GB200 | **CUDA:** 13.0 | **PyTorch:** 2.9.1+cu130 + +> FLA baseline: [flash-linear-attention v0.4.2](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.4.2) + + + +## KDA (Kimi Delta Attention) + +### Fixed-Length (H=64, D=128, bf16) + +| B | T | FLA Triton (ms) | cuLA (ms) | Speedup | +|---|---|-----------------|-----------|---------| +| 1 | 512 | 0.602 | 0.492 | **1.22x** | +| 1 | 1024 | 0.633 | 0.521 | **1.22x** | +| 1 | 4096 | 0.750 | 0.539 | **1.39x** | +| 1 | 8192 | 1.393 | 1.002 | **1.39x** | +| 1 | 16384 | 2.707 | 1.916 | **1.41x** | +| 2 | 512 | 0.610 | 0.523 | **1.17x** | +| 2 | 1024 | 0.644 | 0.524 | **1.23x** | +| 2 | 4096 | 1.388 | 1.005 | **1.38x** | +| 2 | 8192 | 2.704 | 1.933 | **1.40x** | +| 2 | 16384 | 5.303 | 3.821 | **1.39x** | + +Summary (10 configs): **avg=1.32x**, min=1.17x, max=1.41x. + + +### Variable-Length (H=64, D=128, bf16) + +| Config | FLA Triton (ms) | cuLA (ms) | Speedup | +|--------|-----------------|-----------|---------| +| uniform 10seqs T=4096 [409..415] avg=409 | 0.787 | 0.582 | **1.35x** | +| random 10seqs T=4096 [24..1201] avg=409 | 0.782 | 0.576 | **1.36x** | +| skewed 10seqs T=4096 [227..2053] avg=409 | 0.777 | 0.575 | **1.35x** | +| uniform 20seqs T=4096 [204..220] avg=204 | 0.858 | 0.633 | **1.36x** | +| random 20seqs T=4096 [5..787] avg=204 | 0.831 | 0.616 | **1.35x** | +| skewed 20seqs T=4096 [107..2063] avg=204 | 0.813 | 0.596 | **1.36x** | +| uniform 10seqs T=8192 [819..821] avg=819 | 1.389 | 1.022 | **1.36x** | +| random 10seqs T=8192 [48..2401] avg=819 | 1.413 | 1.041 | **1.36x** | +| skewed 10seqs T=8192 [455..4097] avg=819 | 1.440 | 1.045 | **1.38x** | +| uniform 20seqs T=8192 [409..421] avg=409 | 1.476 | 1.069 | **1.38x** | +| random 20seqs T=8192 [9..1574] avg=409 | 1.476 | 1.073 | **1.38x** | +| skewed 20seqs T=8192 [215..4107] avg=409 | 1.484 | 1.077 | **1.38x** | +| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.671 | 1.946 | **1.37x** | +| random 10seqs T=16384 [95..4802] avg=1638 | 2.680 | 1.946 | **1.38x** | +| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.684 | 1.950 | **1.38x** | +| uniform 20seqs T=16384 [819..823] avg=819 | 2.677 | 1.947 | **1.38x** | +| random 20seqs T=16384 [19..3147] avg=819 | 2.713 | 1.971 | **1.38x** | +| skewed 20seqs T=16384 [431..8195] avg=819 | 2.689 | 1.950 | **1.38x** | + +Summary (18 configs): **avg=1.37x**, min=1.35x, max=1.38x. + + +To reproduce: + +```bash +python benchmarks/bench_kda.py --mode both +``` + +## Lightning Attention + +### Prefill (H=64, D=128, bf16) + +| B | T | FLA Triton (ms) | cuLA (ms) | Speedup | +|---|---|-----------------|-----------|---------| +| 1 | 1024 | 0.108 | 0.069 | **1.57x** | +| 1 | 4096 | 0.174 | 0.157 | **1.11x** | +| 1 | 8192 | 0.331 | 0.293 | **1.13x** | +| 1 | 16384 | 0.638 | 0.563 | **1.13x** | +| 2 | 1024 | 0.094 | 0.063 | **1.49x** | +| 2 | 4096 | 0.305 | 0.176 | **1.73x** | +| 2 | 8192 | 0.585 | 0.328 | **1.78x** | +| 2 | 16384 | 1.139 | 0.632 | **1.80x** | + +### Variable-Length (H=64, D=128, bf16) + +Persistent CuTe DSL kernel vs FLA Triton varlen. + +| N (seqs) | T | cuLA (ms) | FLA Triton (ms) | Speedup | +|----------|---|-----------|-----------------|---------| +| 5 | 1020 | 0.090 | 0.187 | **2.09x** | +| 5 | 2045 | 0.113 | 0.211 | **1.87x** | +| 5 | 4095 | 0.164 | 0.258 | **1.58x** | +| 5 | 8190 | 0.265 | 0.412 | **1.56x** | +| 5 | 16380 | 0.465 | 0.705 | **1.52x** | +| 5 | 32765 | 0.859 | 1.284 | **1.49x** | +| 8 | 1024 | 0.090 | 0.172 | **1.92x** | +| 8 | 2048 | 0.114 | 0.198 | **1.74x** | +| 8 | 4096 | 0.158 | 0.252 | **1.60x** | +| 8 | 8192 | 0.243 | 0.399 | **1.64x** | +| 8 | 16384 | 0.413 | 0.689 | **1.67x** | +| 8 | 32768 | 0.758 | 1.259 | **1.66x** | +| 10 | 1020 | 0.107 | 0.171 | **1.60x** | +| 10 | 2040 | 0.135 | 0.200 | **1.48x** | +| 10 | 4090 | 0.182 | 0.268 | **1.47x** | +| 10 | 8190 | 0.266 | 0.407 | **1.53x** | +| 10 | 16380 | 0.440 | 0.694 | **1.58x** | +| 10 | 32760 | 0.791 | 1.275 | **1.61x** | +| 12 | 1020 | 0.120 | 0.176 | **1.47x** | +| 12 | 2040 | 0.145 | 0.194 | **1.34x** | +| 12 | 4092 | 0.192 | 0.265 | **1.38x** | +| 12 | 8184 | 0.279 | 0.404 | **1.45x** | +| 12 | 16380 | 0.455 | 0.700 | **1.54x** | +| 12 | 32760 | 0.795 | 1.267 | **1.59x** | +| 16 | 1024 | 0.124 | 0.166 | **1.34x** | +| 16 | 2048 | 0.150 | 0.187 | **1.25x** | +| 16 | 4096 | 0.189 | 0.258 | **1.37x** | +| 16 | 8192 | 0.268 | 0.401 | **1.49x** | +| 16 | 16384 | 0.426 | 0.688 | **1.61x** | +| 16 | 32768 | 0.742 | 1.251 | **1.68x** | +| 20 | 1020 | 0.163 | 0.170 | **1.04x** | +| 20 | 2040 | 0.192 | 0.202 | **1.05x** | +| 20 | 4080 | 0.237 | 0.287 | **1.21x** | +| 20 | 8180 | 0.321 | 0.431 | **1.34x** | +| 20 | 16380 | 0.482 | 0.701 | **1.45x** | +| 20 | 32760 | 0.806 | 1.267 | **1.57x** | +| 25 | 1000 | 0.195 | 0.182 | **0.93x** | +| 25 | 2025 | 0.223 | 0.224 | **1.01x** | +| 25 | 4075 | 0.263 | 0.277 | **1.05x** | +| 25 | 8175 | 0.348 | 0.444 | **1.27x** | +| 25 | 16375 | 0.522 | 0.717 | **1.37x** | +| 25 | 32750 | 0.835 | 1.275 | **1.53x** | + +Summary (126 configs across uniform/skewed/random): **avg=1.48x**, min=0.93x, max=2.16x. + +To reproduce: + +```bash +python benchmarks/bench_lightning_attn.py --modes no_state varlen +``` + +### Decode (H=64, D=128, bf16, T=1) + +Single-token decode: la_decode (CuTe DSL) vs fla fused_recurrent (Triton). + +#### Kernel-Only + +| B | FLA Triton (ms) | cuLA (ms) | Speedup | +|---|-----------------|-----------|---------| +| 1 | 0.0734 | 0.0125 | **5.89x** | +| 4 | 0.0706 | 0.0132 | **5.37x** | +| 16 | 0.0750 | 0.0209 | **3.59x** | +| 64 | 0.0996 | 0.0843 | **1.18x** | +| 256 | 0.3497 | 0.3121 | **1.12x** | + +#### Wrapper (Full Call Path) + +| B | FLA Triton (ms) | cuLA (ms) | Speedup | +|---|-----------------|-----------|---------| +| 1 | 0.0988 | 0.0189 | **5.23x** | +| 4 | 0.0933 | 0.0182 | **5.12x** | +| 16 | 0.0990 | 0.0209 | **4.74x** | +| 64 | 0.1040 | 0.0844 | **1.23x** | +| 256 | 0.3500 | 0.3134 | **1.12x** | + +To reproduce: + +```bash +python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128 +``` diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp index 5d3ed679..6bfa78cd 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp @@ -88,12 +88,18 @@ struct KdaChunkFwdRecompWUKernelSm100 { static constexpr int NumLoadAuxThreads = 64; // warp 10-11 // ===================== Kernel-only Constants ===================== - static constexpr int NumPrologueRegs = 208; // WG0: element-wise + R2T Akk - static constexpr int NumEpilogueRegs = 216; // WG1: T2R acc + R2G store + kg - static constexpr int NumLoadRegs = 80; // WG2: TMA load + MMA + Aux - static constexpr bool StoreQG = Mainloop::StoreQG; + // NOTE: NVCC 12.9 and 13.0 have performance diffs on the same register config based on our testing +#if CUDA_VERSION_CHECK >= 13000 + static constexpr int NumPrologueRegs = StoreQG ? 232 : 224; // WG0: element-wise + R2T Akk + static constexpr int NumEpilogueRegs = StoreQG ? 192 : 200; // WG1: T2R acc + R2G store + kg +#else + static constexpr int NumPrologueRegs = StoreQG ? 216 : 224; // WG0: element-wise + R2T Akk + static constexpr int NumEpilogueRegs = StoreQG ? 208 : 200; // WG1: T2R acc + R2G store + kg +#endif + static constexpr int NumLoadRegs = 80; // WG2: TMA load + MMA + Aux + // ===================== Warp Roles ===================== enum class WarpRole { Prologue, // WG0: warp 0-3, element-wise K_proc/V_proc → signal MMA diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index e627e0dc..718e0753 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -51,7 +51,7 @@ struct KdaChunkFwdRecompWUMainloopSm100 { static constexpr bool StoreQG = StoreQG_; using ElementBeta = ElementBeta_; - // TODO: double buffer for TMEM acc + // TODO: try optimization with tcgen05.mma.ws enum class TmemAllocation : uint32_t { W = 0, // W, acc, single buffer, [0, 64], U = W + 16 * 65536, // U, acc, [0, 64] +lane16 @@ -736,26 +736,26 @@ struct KdaChunkFwdRecompWUMainloopSm100 { __nv_bfloat16* out_row_base = out_ptr_base + (token_offset_cur + row) * params.d * params.h + head_idx * params.d; - constexpr int HalfK = TileK / 2; + constexpr int QuarK = TileK / 4; ku::tcgen05_after_thread_sync(); #pragma unroll - for (int half = 0; half < 2; ++half) { - float res_half[HalfK]; - ku::tmem_ld_32dp32bNx( - uint32_t(TmemAllocation::W) + buf_mma_idx * 256 + half * HalfK, res_half); + for (int quar = 0; quar < 4; ++quar) { + float res_quar[QuarK]; + ku::tmem_ld_32dp32bNx( + uint32_t(TmemAllocation::W) + buf_mma_idx * 256 + quar * QuarK, res_quar); cutlass::arch::fence_view_async_tmem_load(); if (row < sub_seq_len) { #pragma unroll - for (int i = 0; i < HalfK / 16; ++i) { + for (int i = 0; i < QuarK / 16; ++i) { ku::bf16x16 out; #pragma unroll for (int j = 0; j < 8; ++j) { reinterpret_cast<__nv_bfloat162*>(&out)[j] = - __float22bfloat162_rn(reinterpret_cast(&res_half[i * 16])[j]); + __float22bfloat162_rn(reinterpret_cast(&res_quar[i * 16])[j]); } - store_256b(&out, out_row_base + half * HalfK + i * 16); + store_256b(&out, out_row_base + quar * QuarK + i * 16); } } } diff --git a/csrc/kerutils/include/kerutils/device/common.h b/csrc/kerutils/include/kerutils/device/common.h index 4039bca2..4e338145 100644 --- a/csrc/kerutils/include/kerutils/device/common.h +++ b/csrc/kerutils/include/kerutils/device/common.h @@ -92,3 +92,7 @@ static_assert(false, "kerutils doesn't support SM architectures below SM80"); #define KERUTILS_ENABLE_SM100 #define KERUTILS_ENABLE_SM100A #endif + +#ifndef CUDA_VERSION_CHECK +#define CUDA_VERSION_CHECK (__CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10) +#endif From 9d8eed6e0892a757db4a30cb2e9cd428209cc1f1 Mon Sep 17 00:00:00 2001 From: yechenzhi Date: Wed, 6 May 2026 13:48:26 +0800 Subject: [PATCH 08/34] fix output_final_state wrapper issue (#63) --- cula/kda/hopper_fused_fwd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cula/kda/hopper_fused_fwd.py b/cula/kda/hopper_fused_fwd.py index 9a5af2b3..152cfd31 100644 --- a/cula/kda/hopper_fused_fwd.py +++ b/cula/kda/hopper_fused_fwd.py @@ -121,7 +121,7 @@ def forward( # reshape back o = rearrange(o, "(b t) h d -> b t h d", b=batch_size) - return o.to(q.dtype), final_state + return o.to(q.dtype), final_state if output_final_state else None @staticmethod @input_guard From a6d957293cda4d60a580c077c4f3a186148ba36b Mon Sep 17 00:00:00 2001 From: yechenzhi Date: Sat, 9 May 2026 10:36:13 +0800 Subject: [PATCH 09/34] [KDA] fix internal output_final_state wrapper issue in SM90 (#66) * fix output_final_state wrapper issue * fix internal output state issue * fix signature * explicit output_final_state overwrite output_state_ buffer --- csrc/api/kda_sm90.cu | 31 ++++++++---- csrc/api/pybind.cu | 3 +- csrc/kda/sm90/collective/mainloop_kda_fwd.hpp | 4 +- cula/kda/hopper_fused_fwd.py | 6 ++- tests/test_kda_fused_fwd.py | 47 +++++++++++++++++++ 5 files changed, 77 insertions(+), 14 deletions(-) diff --git a/csrc/api/kda_sm90.cu b/csrc/api/kda_sm90.cu index 7acd0685..e8f3a545 100644 --- a/csrc/api/kda_sm90.cu +++ b/csrc/api/kda_sm90.cu @@ -21,7 +21,7 @@ using OptionalTensor = std::optional; -std::tuple +std::tuple kda_fwd_prefill( OptionalTensor output_, OptionalTensor output_state_, @@ -34,6 +34,7 @@ kda_fwd_prefill( torch::Tensor const& cu_seqlens, torch::Tensor workspace_buffer, float scale, + bool output_final_state, bool safe_gate) { // Q, K, V: [packed_seq, H, D] (already packed by Python layer) auto packed_seq = q.size(0); @@ -52,12 +53,17 @@ kda_fwd_prefill( {packed_seq, num_heads, head_size}, torch::TensorOptions().dtype(q.dtype()).device(q.device())); - // Allocate output state if not provided - torch::Tensor output_state = output_state_.has_value() - ? output_state_.value() - : torch::zeros( - {num_seqs, num_heads, head_size, head_size}, - torch::TensorOptions().dtype(torch::kFloat32).device(q.device())); + // output_final_state controls the API side effect. If it is false, ignore + // even an explicitly provided output_state_ buffer so the kernel skips the + // final-state store. + OptionalTensor output_state = std::nullopt; + if (output_final_state) { + output_state = output_state_.has_value() + ? output_state_.value() + : torch::zeros( + {num_seqs, num_heads, head_size, head_size}, + torch::TensorOptions().dtype(torch::kFloat32).device(q.device())); + } // Validate dtypes TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bfloat16"); @@ -70,7 +76,10 @@ kda_fwd_prefill( TORCH_CHECK(k.is_contiguous(), "k must be contiguous"); TORCH_CHECK(v.is_contiguous(), "v must be contiguous"); TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); - TORCH_CHECK(output_state.is_contiguous(), "output_state must be contiguous"); + if (output_state.has_value()) { + TORCH_CHECK(output_state->dtype() == torch::kFloat32, "output_state must be float32"); + TORCH_CHECK(output_state->is_contiguous(), "output_state must be contiguous"); + } TORCH_CHECK(cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous"); TORCH_CHECK(workspace_buffer.is_contiguous(), "workspace_buffer must be contiguous"); @@ -87,6 +96,8 @@ kda_fwd_prefill( "alpha shape must be [packed_seq, num_heads, head_size]"); alpha_ptr = alpha.data_ptr(); } + + float* output_state_ptr = output_state.has_value() ? output_state->data_ptr() : nullptr; if (beta_.has_value()) { auto& beta = beta_.value(); TORCH_CHECK( @@ -121,7 +132,7 @@ kda_fwd_prefill( kda::sm90::launch_kda_fwd_prefill_kernel( stream, reinterpret_cast(output.data_ptr()), - output_state.data_ptr(), + output_state_ptr, reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), reinterpret_cast(v.data_ptr()), @@ -142,7 +153,7 @@ kda_fwd_prefill( kda::sm90::launch_kda_fwd_prefill_kernel( stream, reinterpret_cast(output.data_ptr()), - output_state.data_ptr(), + output_state_ptr, reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), reinterpret_cast(v.data_ptr()), diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index ba2deb6b..d14a41c5 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -51,7 +51,7 @@ ChunkKDAFwdRecompWU( #endif #if defined(CULA_SM90A_ENABLED) -std::tuple +std::tuple> kda_fwd_prefill( std::optional output_, std::optional output_state_, @@ -64,6 +64,7 @@ kda_fwd_prefill( torch::Tensor const& cu_seqlens, torch::Tensor workspace_buffer, float scale, + bool output_final_state, bool safe_gate); #endif diff --git a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp index d22814db..1f781693 100644 --- a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp +++ b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp @@ -1369,7 +1369,9 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::true_type{}); } - kv_store(); + if (params.ptr_output_state != nullptr) { + kv_store(); + } } template diff --git a/cula/kda/hopper_fused_fwd.py b/cula/kda/hopper_fused_fwd.py index 152cfd31..cc42827a 100644 --- a/cula/kda/hopper_fused_fwd.py +++ b/cula/kda/hopper_fused_fwd.py @@ -102,7 +102,8 @@ def forward( workspace_buffer = _get_cache_buf("hopper_kda_fwd_workspace", workspace_size, q.device) # call the C++ kernel - # Signature: kda_fwd_prefill(output_, output_state_, q, k, v, input_state_, alpha_, beta_, cu_seqlens, workspace, scale, safe_gate) + # Signature: kda_fwd_prefill(output_, output_state_, q, k, v, input_state_, alpha_, beta_, cu_seqlens, + # workspace, scale, output_final_state, safe_gate) o, final_state = cula_cuda.kda_fwd_prefill( None, # output_ (auto-allocate) None, # output_state_ (auto-allocate) @@ -115,13 +116,14 @@ def forward( cu_seqlens, workspace_buffer, scale, + output_final_state, safe_gate, ) # reshape back o = rearrange(o, "(b t) h d -> b t h d", b=batch_size) - return o.to(q.dtype), final_state if output_final_state else None + return o.to(q.dtype), final_state @staticmethod @input_guard diff --git a/tests/test_kda_fused_fwd.py b/tests/test_kda_fused_fwd.py index b9b59f94..9c325524 100644 --- a/tests/test_kda_fused_fwd.py +++ b/tests/test_kda_fused_fwd.py @@ -173,6 +173,53 @@ def test_safe_gate_chunk( assert_close("ht", ref_ht_fla_trans, tri_ht, 0.005) +def test_safe_gate_chunk_no_final_state(): + cula_kda_fused_fwd = get_kda_fused_fwd(device) + + B, T, H, D = 1, 63, 1, 128 + dtype = torch.bfloat16 + + torch.manual_seed(42) + q = torch.rand(B, T, H, D, dtype=dtype, device=device) + k = torch.rand(B, T, H, D, dtype=dtype, device=device) + v = torch.rand(B, T, H, D, dtype=dtype, device=device) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=device)).clamp(-5, 0) + beta = torch.randn(B, T, H, dtype=torch.float32, device=device).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32, device=device) + h0_vk = h0.transpose(-1, -2).contiguous() + + q = F.normalize(q, p=2, dim=-1) + k = F.normalize(k, p=2, dim=-1) + + tri_no_state, tri_ht_no_state = cula_kda_fused_fwd( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0_vk.clone(), + output_final_state=False, + safe_gate=True, + lower_bound=-5.0, + ) + + tri_with_state, tri_ht_with_state = cula_kda_fused_fwd( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0_vk.clone(), + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + ) + + assert tri_ht_no_state is None + assert tri_ht_with_state is not None + assert_close("o", tri_with_state, tri_no_state, 0.005) + + @pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) @pytest.mark.parametrize( ("H", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), From 82899768597e65d2232ab5830c791b402cac2f58 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Wed, 13 May 2026 12:33:33 +0800 Subject: [PATCH 10/34] [Fix] change to umma pipelines for SM100 KDA (#68) * fix umma pipeline in chunk_intra and recomp_wu * update bench * fix deter check with ncu/sanitizer mode Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * update iters * add more compile option for profiling --------- Co-authored-by: boyu.zbw Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- BENCHMARK_GB200_CUDA_130.md | 186 +++++++++--------- benchmarks/bench_kda.py | 47 +++++ csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp | 7 +- .../sm100/kda_fwd_intra_mainloop_sm100.hpp | 3 +- .../sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp | 12 +- .../kda_fwd_recomp_w_u_mainloop_sm100.hpp | 7 +- cula/ops/chunk_delta_h.py | 4 +- cula/ops/fwd_o.py | 4 +- 8 files changed, 161 insertions(+), 109 deletions(-) diff --git a/BENCHMARK_GB200_CUDA_130.md b/BENCHMARK_GB200_CUDA_130.md index a34b0a28..b14d8f9b 100644 --- a/BENCHMARK_GB200_CUDA_130.md +++ b/BENCHMARK_GB200_CUDA_130.md @@ -1,6 +1,6 @@ # Benchmark Results -> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-04-23. +> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-05-12. > **GPU:** NVIDIA GB200 | **CUDA:** 13.0 | **PyTorch:** 2.9.1+cu130 @@ -14,44 +14,44 @@ | B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------|---------| -| 1 | 512 | 0.602 | 0.492 | **1.22x** | -| 1 | 1024 | 0.633 | 0.521 | **1.22x** | -| 1 | 4096 | 0.750 | 0.539 | **1.39x** | -| 1 | 8192 | 1.393 | 1.002 | **1.39x** | -| 1 | 16384 | 2.707 | 1.916 | **1.41x** | -| 2 | 512 | 0.610 | 0.523 | **1.17x** | -| 2 | 1024 | 0.644 | 0.524 | **1.23x** | -| 2 | 4096 | 1.388 | 1.005 | **1.38x** | -| 2 | 8192 | 2.704 | 1.933 | **1.40x** | -| 2 | 16384 | 5.303 | 3.821 | **1.39x** | +| 1 | 512 | 0.582 | 0.483 | **1.21x** | +| 1 | 1024 | 0.579 | 0.493 | **1.17x** | +| 1 | 4096 | 0.749 | 0.541 | **1.38x** | +| 1 | 8192 | 1.393 | 1.009 | **1.38x** | +| 1 | 16384 | 2.706 | 1.931 | **1.40x** | +| 2 | 512 | 0.595 | 0.510 | **1.17x** | +| 2 | 1024 | 0.619 | 0.498 | **1.24x** | +| 2 | 4096 | 1.394 | 1.016 | **1.37x** | +| 2 | 8192 | 2.701 | 1.949 | **1.39x** | +| 2 | 16384 | 5.297 | 3.875 | **1.37x** | -Summary (10 configs): **avg=1.32x**, min=1.17x, max=1.41x. +Summary (10 configs): **avg=1.31x**, min=1.17x, max=1.40x. ### Variable-Length (H=64, D=128, bf16) | Config | FLA Triton (ms) | cuLA (ms) | Speedup | |--------|-----------------|-----------|---------| -| uniform 10seqs T=4096 [409..415] avg=409 | 0.787 | 0.582 | **1.35x** | -| random 10seqs T=4096 [24..1201] avg=409 | 0.782 | 0.576 | **1.36x** | -| skewed 10seqs T=4096 [227..2053] avg=409 | 0.777 | 0.575 | **1.35x** | -| uniform 20seqs T=4096 [204..220] avg=204 | 0.858 | 0.633 | **1.36x** | -| random 20seqs T=4096 [5..787] avg=204 | 0.831 | 0.616 | **1.35x** | -| skewed 20seqs T=4096 [107..2063] avg=204 | 0.813 | 0.596 | **1.36x** | -| uniform 10seqs T=8192 [819..821] avg=819 | 1.389 | 1.022 | **1.36x** | -| random 10seqs T=8192 [48..2401] avg=819 | 1.413 | 1.041 | **1.36x** | -| skewed 10seqs T=8192 [455..4097] avg=819 | 1.440 | 1.045 | **1.38x** | -| uniform 20seqs T=8192 [409..421] avg=409 | 1.476 | 1.069 | **1.38x** | -| random 20seqs T=8192 [9..1574] avg=409 | 1.476 | 1.073 | **1.38x** | -| skewed 20seqs T=8192 [215..4107] avg=409 | 1.484 | 1.077 | **1.38x** | -| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.671 | 1.946 | **1.37x** | -| random 10seqs T=16384 [95..4802] avg=1638 | 2.680 | 1.946 | **1.38x** | -| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.684 | 1.950 | **1.38x** | -| uniform 20seqs T=16384 [819..823] avg=819 | 2.677 | 1.947 | **1.38x** | -| random 20seqs T=16384 [19..3147] avg=819 | 2.713 | 1.971 | **1.38x** | -| skewed 20seqs T=16384 [431..8195] avg=819 | 2.689 | 1.950 | **1.38x** | - -Summary (18 configs): **avg=1.37x**, min=1.35x, max=1.38x. +| uniform 10seqs T=4096 [409..415] avg=409 | 0.783 | 0.585 | **1.34x** | +| random 10seqs T=4096 [24..1201] avg=409 | 0.777 | 0.579 | **1.34x** | +| skewed 10seqs T=4096 [227..2053] avg=409 | 0.776 | 0.578 | **1.34x** | +| uniform 20seqs T=4096 [204..220] avg=204 | 0.855 | 0.633 | **1.35x** | +| random 20seqs T=4096 [5..787] avg=204 | 0.828 | 0.619 | **1.34x** | +| skewed 20seqs T=4096 [107..2063] avg=204 | 0.811 | 0.597 | **1.36x** | +| uniform 10seqs T=8192 [819..821] avg=819 | 1.386 | 1.028 | **1.35x** | +| random 10seqs T=8192 [48..2401] avg=819 | 1.414 | 1.048 | **1.35x** | +| skewed 10seqs T=8192 [455..4097] avg=819 | 1.441 | 1.049 | **1.37x** | +| uniform 20seqs T=8192 [409..421] avg=409 | 1.476 | 1.074 | **1.37x** | +| random 20seqs T=8192 [9..1574] avg=409 | 1.475 | 1.079 | **1.37x** | +| skewed 20seqs T=8192 [215..4107] avg=409 | 1.482 | 1.081 | **1.37x** | +| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.671 | 1.963 | **1.36x** | +| random 10seqs T=16384 [95..4802] avg=1638 | 2.684 | 1.965 | **1.37x** | +| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.688 | 1.972 | **1.36x** | +| uniform 20seqs T=16384 [819..823] avg=819 | 2.680 | 1.966 | **1.36x** | +| random 20seqs T=16384 [19..3147] avg=819 | 2.712 | 1.990 | **1.36x** | +| skewed 20seqs T=16384 [431..8195] avg=819 | 2.691 | 1.970 | **1.37x** | + +Summary (18 configs): **avg=1.36x**, min=1.34x, max=1.37x. To reproduce: @@ -66,14 +66,14 @@ python benchmarks/bench_kda.py --mode both | B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------|---------| -| 1 | 1024 | 0.108 | 0.069 | **1.57x** | -| 1 | 4096 | 0.174 | 0.157 | **1.11x** | -| 1 | 8192 | 0.331 | 0.293 | **1.13x** | -| 1 | 16384 | 0.638 | 0.563 | **1.13x** | -| 2 | 1024 | 0.094 | 0.063 | **1.49x** | -| 2 | 4096 | 0.305 | 0.176 | **1.73x** | -| 2 | 8192 | 0.585 | 0.328 | **1.78x** | -| 2 | 16384 | 1.139 | 0.632 | **1.80x** | +| 1 | 1024 | 0.087 | 0.070 | **1.24x** | +| 1 | 4096 | 0.175 | 0.157 | **1.11x** | +| 1 | 8192 | 0.330 | 0.292 | **1.13x** | +| 1 | 16384 | 0.628 | 0.563 | **1.12x** | +| 2 | 1024 | 0.099 | 0.064 | **1.53x** | +| 2 | 4096 | 0.327 | 0.175 | **1.87x** | +| 2 | 8192 | 0.631 | 0.327 | **1.93x** | +| 2 | 16384 | 1.249 | 0.632 | **1.98x** | ### Variable-Length (H=64, D=128, bf16) @@ -81,50 +81,50 @@ Persistent CuTe DSL kernel vs FLA Triton varlen. | N (seqs) | T | cuLA (ms) | FLA Triton (ms) | Speedup | |----------|---|-----------|-----------------|---------| -| 5 | 1020 | 0.090 | 0.187 | **2.09x** | -| 5 | 2045 | 0.113 | 0.211 | **1.87x** | -| 5 | 4095 | 0.164 | 0.258 | **1.58x** | -| 5 | 8190 | 0.265 | 0.412 | **1.56x** | -| 5 | 16380 | 0.465 | 0.705 | **1.52x** | -| 5 | 32765 | 0.859 | 1.284 | **1.49x** | -| 8 | 1024 | 0.090 | 0.172 | **1.92x** | -| 8 | 2048 | 0.114 | 0.198 | **1.74x** | -| 8 | 4096 | 0.158 | 0.252 | **1.60x** | -| 8 | 8192 | 0.243 | 0.399 | **1.64x** | -| 8 | 16384 | 0.413 | 0.689 | **1.67x** | -| 8 | 32768 | 0.758 | 1.259 | **1.66x** | -| 10 | 1020 | 0.107 | 0.171 | **1.60x** | -| 10 | 2040 | 0.135 | 0.200 | **1.48x** | -| 10 | 4090 | 0.182 | 0.268 | **1.47x** | -| 10 | 8190 | 0.266 | 0.407 | **1.53x** | -| 10 | 16380 | 0.440 | 0.694 | **1.58x** | -| 10 | 32760 | 0.791 | 1.275 | **1.61x** | -| 12 | 1020 | 0.120 | 0.176 | **1.47x** | -| 12 | 2040 | 0.145 | 0.194 | **1.34x** | -| 12 | 4092 | 0.192 | 0.265 | **1.38x** | -| 12 | 8184 | 0.279 | 0.404 | **1.45x** | -| 12 | 16380 | 0.455 | 0.700 | **1.54x** | -| 12 | 32760 | 0.795 | 1.267 | **1.59x** | -| 16 | 1024 | 0.124 | 0.166 | **1.34x** | -| 16 | 2048 | 0.150 | 0.187 | **1.25x** | -| 16 | 4096 | 0.189 | 0.258 | **1.37x** | -| 16 | 8192 | 0.268 | 0.401 | **1.49x** | -| 16 | 16384 | 0.426 | 0.688 | **1.61x** | -| 16 | 32768 | 0.742 | 1.251 | **1.68x** | -| 20 | 1020 | 0.163 | 0.170 | **1.04x** | -| 20 | 2040 | 0.192 | 0.202 | **1.05x** | -| 20 | 4080 | 0.237 | 0.287 | **1.21x** | -| 20 | 8180 | 0.321 | 0.431 | **1.34x** | -| 20 | 16380 | 0.482 | 0.701 | **1.45x** | -| 20 | 32760 | 0.806 | 1.267 | **1.57x** | -| 25 | 1000 | 0.195 | 0.182 | **0.93x** | -| 25 | 2025 | 0.223 | 0.224 | **1.01x** | -| 25 | 4075 | 0.263 | 0.277 | **1.05x** | -| 25 | 8175 | 0.348 | 0.444 | **1.27x** | -| 25 | 16375 | 0.522 | 0.717 | **1.37x** | -| 25 | 32750 | 0.835 | 1.275 | **1.53x** | - -Summary (126 configs across uniform/skewed/random): **avg=1.48x**, min=0.93x, max=2.16x. +| 5 | 1020 | 0.089 | 0.171 | **1.91x** | +| 5 | 2045 | 0.111 | 0.189 | **1.71x** | +| 5 | 4095 | 0.163 | 0.249 | **1.53x** | +| 5 | 8190 | 0.264 | 0.399 | **1.51x** | +| 5 | 16380 | 0.463 | 0.702 | **1.52x** | +| 5 | 32765 | 0.858 | 1.283 | **1.49x** | +| 8 | 1024 | 0.086 | 0.156 | **1.82x** | +| 8 | 2048 | 0.111 | 0.183 | **1.65x** | +| 8 | 4096 | 0.157 | 0.250 | **1.59x** | +| 8 | 8192 | 0.243 | 0.402 | **1.66x** | +| 8 | 16384 | 0.413 | 0.688 | **1.67x** | +| 8 | 32768 | 0.756 | 1.252 | **1.66x** | +| 10 | 1020 | 0.104 | 0.162 | **1.56x** | +| 10 | 2040 | 0.133 | 0.200 | **1.51x** | +| 10 | 4090 | 0.179 | 0.269 | **1.50x** | +| 10 | 8190 | 0.267 | 0.414 | **1.55x** | +| 10 | 16380 | 0.439 | 0.693 | **1.58x** | +| 10 | 32760 | 0.788 | 1.260 | **1.60x** | +| 12 | 1020 | 0.119 | 0.175 | **1.47x** | +| 12 | 2040 | 0.143 | 0.197 | **1.38x** | +| 12 | 4092 | 0.189 | 0.265 | **1.40x** | +| 12 | 8184 | 0.281 | 0.405 | **1.44x** | +| 12 | 16380 | 0.452 | 0.703 | **1.55x** | +| 12 | 32760 | 0.793 | 1.259 | **1.59x** | +| 16 | 1024 | 0.121 | 0.157 | **1.30x** | +| 16 | 2048 | 0.149 | 0.183 | **1.23x** | +| 16 | 4096 | 0.187 | 0.256 | **1.37x** | +| 16 | 8192 | 0.267 | 0.398 | **1.49x** | +| 16 | 16384 | 0.424 | 0.686 | **1.62x** | +| 16 | 32768 | 0.740 | 1.247 | **1.68x** | +| 20 | 1020 | 0.162 | 0.174 | **1.07x** | +| 20 | 2040 | 0.191 | 0.207 | **1.08x** | +| 20 | 4080 | 0.233 | 0.288 | **1.24x** | +| 20 | 8180 | 0.319 | 0.424 | **1.33x** | +| 20 | 16380 | 0.478 | 0.703 | **1.47x** | +| 20 | 32760 | 0.800 | 1.261 | **1.58x** | +| 25 | 1000 | 0.193 | 0.176 | **0.91x** | +| 25 | 2025 | 0.221 | 0.227 | **1.03x** | +| 25 | 4075 | 0.258 | 0.286 | **1.11x** | +| 25 | 8175 | 0.347 | 0.445 | **1.28x** | +| 25 | 16375 | 0.517 | 0.720 | **1.39x** | +| 25 | 32750 | 0.831 | 1.270 | **1.53x** | + +Summary (126 configs across uniform/skewed/random): **avg=1.48x**, min=0.91x, max=2.01x. To reproduce: @@ -140,21 +140,21 @@ Single-token decode: la_decode (CuTe DSL) vs fla fused_recurrent (Triton). | B | FLA Triton (ms) | cuLA (ms) | Speedup | |---|-----------------|-----------|---------| -| 1 | 0.0734 | 0.0125 | **5.89x** | -| 4 | 0.0706 | 0.0132 | **5.37x** | -| 16 | 0.0750 | 0.0209 | **3.59x** | +| 1 | 0.0740 | 0.0134 | **5.53x** | +| 4 | 0.0698 | 0.0130 | **5.39x** | +| 16 | 0.0731 | 0.0209 | **3.50x** | | 64 | 0.0996 | 0.0843 | **1.18x** | -| 256 | 0.3497 | 0.3121 | **1.12x** | +| 256 | 0.3501 | 0.3126 | **1.12x** | #### Wrapper (Full Call Path) | B | FLA Triton (ms) | cuLA (ms) | Speedup | |---|-----------------|-----------|---------| -| 1 | 0.0988 | 0.0189 | **5.23x** | -| 4 | 0.0933 | 0.0182 | **5.12x** | -| 16 | 0.0990 | 0.0209 | **4.74x** | -| 64 | 0.1040 | 0.0844 | **1.23x** | -| 256 | 0.3500 | 0.3134 | **1.12x** | +| 1 | 0.0958 | 0.0189 | **5.08x** | +| 4 | 0.0920 | 0.0186 | **4.95x** | +| 16 | 0.0934 | 0.0211 | **4.43x** | +| 64 | 0.0990 | 0.0850 | **1.17x** | +| 256 | 0.3492 | 0.3133 | **1.11x** | To reproduce: diff --git a/benchmarks/bench_kda.py b/benchmarks/bench_kda.py index 0521f11b..dc31d11d 100644 --- a/benchmarks/bench_kda.py +++ b/benchmarks/bench_kda.py @@ -66,6 +66,12 @@ # ============================================================ # Helpers # ============================================================ +def generate_balanced_seqlens(total_tokens, num_seqs): + base = total_tokens // num_seqs + remainder = total_tokens % num_seqs + return [base] * (num_seqs - 1) + [base + remainder] + + def time_kernel(fn, warmup=None, n_iters=None): if warmup is None: warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP @@ -118,6 +124,44 @@ def run_kda(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, low ) +def check_determinism(H=4, total_T=8192, num_seqs=10, iters=10000): + """Run the kernel multiple times and check that outputs are identical.""" + device = torch.device("cuda") + D = 128 + seq_lens = generate_balanced_seqlens(total_T, num_seqs) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + + inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + init_state=init_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + + ref_out, ref_state = run_kda(**common, fn=cula_chunk_kda) + + for i in range(iters): + out, state = run_kda(**common, fn=cula_chunk_kda) + assert torch.isnan(out).sum() == 0, f"Output contains NaNs at iter {i}" + assert torch.isnan(state).sum() == 0, f"State contains NaNs at iter {i}" + assert torch.isfinite(out).all(), f"Output contains infs at iter {i}" + assert torch.isfinite(state).all(), f"State contains infs at iter {i}" + assert torch.equal(out, ref_out), f"Output mismatch at iter {i}" + assert torch.equal(state, ref_state), f"State mismatch at iter {i}" + + # ============================================================ # Fixed-length benchmark # ============================================================ @@ -380,6 +424,9 @@ def main(): fixed_res, varlen_res = [], [] + if not (args.ncu or args.sanitizer): + check_determinism(H=H, iters=10000) + if args.mode in ("fixed", "both"): fixed_res = bench_fixed(fixed_configs) diff --git a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp index 60dc4b34..f314723f 100644 --- a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp @@ -180,7 +180,7 @@ struct KdaChunkFwdIntraKernelSm100 { // === MMA -> CudaCore pipelines (UMMA) === typename PipelineQKDone::Params qk_done_pipe_params; - qk_done_pipe_params.producer_arv_count = NumMmaThreads; + qk_done_pipe_params.producer_arv_count = 1; qk_done_pipe_params.consumer_arv_count = NumCudaCoreThreads; if (role == WarpRole::Mma) { @@ -214,10 +214,7 @@ struct KdaChunkFwdIntraKernelSm100 { PipelineQKGInterReady qkg_inter_pipeline( shared_plan->pipe_qkg_inter_storage, qkg_inter_pipe_params, ClusterShape{}); - PipelineQKDone qk_done_pipeline( - shared_plan->pipe_qk_done_storage, - qk_done_pipe_params, - /*InitBarriers*/ cute::true_type{}); + PipelineQKDone qk_done_pipeline(shared_plan->pipe_qk_done_storage, qk_done_pipe_params, ClusterShape{}); PipelineKKInvReady kk_inv_pipeline( shared_plan->pipe_kk_inv_storage, diff --git a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp index 3aa2746b..849e910c 100644 --- a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp @@ -155,7 +155,7 @@ struct KdaChunkFwdIntraMainloopSm100 { using PipelineQKGInterReady = cutlass::PipelineUmmaConsumerAsync; - using PipelineQKDone = cutlass::PipelineAsync; + using PipelineQKDone = cutlass::PipelineUmmaAsync; using PipelineKKInvReady = cutlass::PipelineAsync; @@ -221,7 +221,6 @@ struct KdaChunkFwdIntraMainloopSm100 { alignas(16) typename PipelineKKInvReady::SharedStorage pipe_kk_inv_storage; - // TODO: support bf16 beta alignas(16) float beta_smem[StagesAcc][TileT]; array_aligned tmem_start_addr; }; diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp index 6bfa78cd..73cb4089 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp @@ -171,7 +171,7 @@ struct KdaChunkFwdRecompWUKernelSm100 { typename PipelineA::Params a_pipe_params; a_pipe_params.transaction_bytes = sizeof(bf16) * cosize_v; a_pipe_params.is_leader = lane_predicate && (role == WarpRole::Load); - a_pipe_params.num_consumers = cutlass::NumThreadsPerWarp; + a_pipe_params.num_consumers = 1; if (role == WarpRole::Load) { a_pipe_params.role = PipelineA::ThreadCategory::Producer; } else if (role == WarpRole::Mma) { @@ -236,11 +236,13 @@ struct KdaChunkFwdRecompWUKernelSm100 { // === Prologue → MMA pipelines === - // PipelinePrologueReady: Prologue+Epilogue(producer, 256 threads) → Mma(consumer, 32 threads) - // Unified pipeline for both K and V prologue ready (co-produced by Prologue and Epilogue) + // PipelinePrologueReady: Prologue+Epilogue(producer, 256 threads) → Mma(consumer, umma_arrive) + // Unified pipeline for both K and V prologue ready (co-produced by Prologue and Epilogue). + // Consumer side uses umma_arrive (tcgen05.commit::mbarrier::arrive), which internally + // elects exactly one thread, so consumer_arv_count must be 1. typename PipelinePrologueReady::Params prologue_ready_pipe_params; prologue_ready_pipe_params.producer_arv_count = NumPrologueThreads + NumEpilogueThreads; - prologue_ready_pipe_params.consumer_arv_count = NumMmaThreads; + prologue_ready_pipe_params.consumer_arv_count = 1; // umma_arrive elects one thread if (role == WarpRole::Prologue || role == WarpRole::Epilogue) { prologue_ready_pipe_params.role = PipelinePrologueReady::ThreadCategory::Producer; } else if (role == WarpRole::Mma) { @@ -276,7 +278,7 @@ struct KdaChunkFwdRecompWUKernelSm100 { /*InitBarriers*/ cute::true_type{}); PipelinePrologueReady prologue_ready_pipeline( - shared_plan->pipe_prologue_ready_storage, prologue_ready_pipe_params, /*InitBarriers*/ cute::true_type{}); + shared_plan->pipe_prologue_ready_storage, prologue_ready_pipe_params, ClusterShape{}); PipelineAccDone acc_done_pipeline(shared_plan->pipe_acc_done_storage, acc_done_pipe_params, ClusterShape{}); diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index 718e0753..0f6a66f8 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -95,8 +95,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { SM100_MMA_F16BF16_SS{})); // ===================== Pipeline Types ===================== + using ClusterShape = Shape<_1, _1, _1>; // TMA load -> MMA (Akk) - using PipelineA = cutlass::PipelineTmaAsync; + using PipelineA = cutlass::PipelineTmaUmmaAsync; // TMA load -> Compute (merged prologue+epilogue) using PipelineV = cutlass::PipelineTmaAsync; // TMA load -> Compute (K) @@ -109,7 +110,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { using PipelineBeta = cutlass::PipelineAsync; // Unified pipeline: Compute -> MMA (K/V prologue ready share one pipeline, used sequentially) - using PipelinePrologueReady = cutlass::PipelineAsync; + // NOTE: must be PipelineUmmaConsumerAsync so that the MMA warp's consumer_release + // issues tcgen05.commit::mbarrier::arrive (umma_arrive). + using PipelinePrologueReady = cutlass::PipelineUmmaConsumerAsync; // Unified pipeline: MMA -> Compute (W/U acc done share one pipeline, used sequentially) using PipelineAccDone = cutlass::PipelineUmmaAsync; diff --git a/cula/ops/chunk_delta_h.py b/cula/ops/chunk_delta_h.py index 3ce0cddd..fb81dd60 100644 --- a/cula/ops/chunk_delta_h.py +++ b/cula/ops/chunk_delta_h.py @@ -37,6 +37,8 @@ from cula.utils import USE_FAST_MATH, assert_blackwell +COMPILE_OPTIONS = "--enable-tvm-ffi --generate-line-info --ptxas-options '--verbose'" + # in FLA, cumsum returns int64 tensor by default @tensor_cache @@ -1958,7 +1960,7 @@ def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fas Int32(0), # store_final_state Int32(0), # save_v_new stream_fake, - options="--enable-tvm-ffi", + options=COMPILE_OPTIONS, ) return compiled_fn diff --git a/cula/ops/fwd_o.py b/cula/ops/fwd_o.py index e301f025..2c820aad 100644 --- a/cula/ops/fwd_o.py +++ b/cula/ops/fwd_o.py @@ -87,6 +87,8 @@ LN2 = 0.6931471805599453 RCP_LN2 = 1.4426950408889634 +COMPILE_OPTIONS = "--enable-tvm-ffi --generate-line-info --ptxas-options '--verbose'" + def make_thread_cooperative_group(size: int): return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) @@ -1721,7 +1723,7 @@ def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, us (Int32(1), Int32(1), Int32(H), Int32(K), Int32(V)), Int32(1), stream_fake, - options="--enable-tvm-ffi", + options=COMPILE_OPTIONS, ) return compiled_fn From d4fd474132635fcd9dbc29aef29341659a720353 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Thu, 14 May 2026 11:29:28 +0800 Subject: [PATCH 11/34] [Fix] delta_h race risk (#69) Co-authored-by: boyu.zbw --- cula/ops/chunk_delta_h.py | 72 +++++++++++++++++++++------------------ 1 file changed, 39 insertions(+), 33 deletions(-) diff --git a/cula/ops/chunk_delta_h.py b/cula/ops/chunk_delta_h.py index fb81dd60..b2a61de3 100644 --- a/cula/ops/chunk_delta_h.py +++ b/cula/ops/chunk_delta_h.py @@ -153,9 +153,10 @@ def __init__( barrier_id=3, num_threads=self.threads_per_warp * len(self.cuda_warp_ids), # 128 ) - # No CTA-wide barrier needed for WU scheduling: - # Load warp elect_one arrives on a lightweight mbarrier (count=1), - # other warps just wait on mbarrier phase (no arrive needed). + # WU scheduling uses full-thread mbarrier arrives (no elect_one) to guarantee + # all 32 lanes have completed their sWorkIdx read before the producer can overwrite it. + # sched_mbar: count=32 (all Load warp threads arrive). + # sched_consumed_mbar: count=32*7 (all threads of 7 consumer warps arrive). self.buffer_align_bytes = 1024 @staticmethod @@ -550,9 +551,9 @@ class SharedStorage: ] # Double-buffered work index for dynamic scheduling sWorkIdx: cute.struct.MemRange[Int32, 2] - # Double-buffered scheduling mbarriers (count=1 each, Load warp elect_one arrives) + # Double-buffered scheduling mbarriers (count=32, all Load warp threads arrive) sched_mbar: cute.struct.MemRange[Int64, 2] - # Double-buffered consumed mbarriers (count=4, one arrive per consumer warp group) + # Double-buffered consumed mbarriers (count=32*7, all threads of 7 consumer warps arrive) sched_consumed_mbar: cute.struct.MemRange[Int64, 2] self.shared_storage = SharedStorage @@ -749,14 +750,14 @@ def kernel( if cutlass.const_expr(self.is_varlen and self.persistent): sched_mbar_base = storage.sched_mbar.data_ptr() sched_consumed_mbar_base = storage.sched_consumed_mbar.data_ptr() - # Init 2 scheduling mbarriers with count=1: only Load warp elect_one arrives - # Init 2 consumed mbarriers with count=7: one arrive per non-load warp - # (MMA=1, CC=4, Store=1, Empty=1) + # Init 2 scheduling mbarriers with count=32: all 32 Load warp threads arrive. + # Init 2 consumed mbarriers with count=32*7: all 32 threads × 7 consumer warps arrive. + # (MMA=1, CC=4, Store=1, Empty=1; all full-thread arrives, no elect_one) if warp_idx == 0: - cute.arch.mbarrier_init(sched_mbar_base, 1) - cute.arch.mbarrier_init(sched_mbar_base + 1, 1) - cute.arch.mbarrier_init(sched_consumed_mbar_base, 7) - cute.arch.mbarrier_init(sched_consumed_mbar_base + 1, 7) + cute.arch.mbarrier_init(sched_mbar_base, 32) + cute.arch.mbarrier_init(sched_mbar_base + 1, 32) + cute.arch.mbarrier_init(sched_consumed_mbar_base, 32 * 7) + cute.arch.mbarrier_init(sched_consumed_mbar_base + 1, 32 * 7) cute.arch.mbarrier_init_fence() cute.arch.barrier(barrier_id=0, number_of_threads=self.threads_per_cta) @@ -854,8 +855,12 @@ def kernel( # ========================================================================= # DYNAMIC SCHEDULING: initial work_idx fetch (persistent varlen only) - # Load warp elect_one does atomicAdd → sWorkIdx[buf] → fence → arrive mbar[buf]. - # Other warps wait on mbar[buf] at the correct phase, then read sWorkIdx[buf]. + # Load warp: elect_one does atomicAdd → sWorkIdx[buf] → fence_acq_rel_cta; + # then ALL 32 threads arrive on sched_mbar[buf] (count=32) so the arrive + # is ordered after every lane's write, not just the elected lane's. + # Consumer warps: wait on sched_mbar[buf]; ALL 32 threads arrive on + # sched_consumed_mbar[buf] (count=32*7) so the Load warp's back-pressure + # wait is satisfied only after every consumer lane has read sWorkIdx. # Double-buffered to avoid phase racing. # ========================================================================= if cutlass.const_expr(self.is_varlen and self.persistent): @@ -864,14 +869,11 @@ def kernel( with cute.arch.elect_one(): first_work_idx = _atomic_add_global_i32(workspace_iter.toint().ir_value(), Int32(1).ir_value()) sWorkIdx[(0,)] = first_work_idx - cute.arch.fence_acq_rel_cta() - cute.arch.mbarrier_arrive(sched_mbar_base) + cute.arch.fence_acq_rel_cta() + cute.arch.mbarrier_arrive(sched_mbar_base) else: # All other warps: wait for Load warp's signal on mbar[0], phase=0 cute.arch.mbarrier_wait(sched_mbar_base, Int32(0)) - # Signal consumed_mbar[0]: buf 0 has been consumed - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(sched_consumed_mbar_base) # ========================================================================= # LOAD WARP @@ -1052,8 +1054,8 @@ def kernel( with cute.arch.elect_one(): next_idx = _atomic_add_global_i32(workspace_iter.toint().ir_value(), Int32(1).ir_value()) sWorkIdx[(sched_buf,)] = next_idx - cute.arch.fence_acq_rel_cta() - cute.arch.mbarrier_arrive(sched_mbar_base + sched_buf) + cute.arch.fence_acq_rel_cta() + cute.arch.mbarrier_arrive(sched_mbar_base + sched_buf) cute.arch.sync_warp() work_idx = sWorkIdx[(sched_buf,)] sched_buf = Int32(1) - sched_buf @@ -1071,6 +1073,8 @@ def kernel( wu_iter = Int32(0) if cutlass.const_expr(self.is_varlen and self.persistent): work_idx = sWorkIdx[(0,)] + # All 32 MMA warp threads arrive on consumed_mbar[0]: buf 0 has been read + cute.arch.mbarrier_arrive(sched_consumed_mbar_base) sched_buf = Int32(1) # next buffer to wait on sched_phase0 = Int32(1) # mbar[0]: init consumed phase=0, next=1 sched_phase1 = Int32(0) # mbar[1]: not yet used, next=0 @@ -1151,9 +1155,8 @@ def kernel( cute.arch.mbarrier_wait(sched_mbar_base + 1, sched_phase1) sched_phase1 = Int32(1) - sched_phase1 work_idx = sWorkIdx[(sched_buf,)] - # Signal consumed: Load warp can reuse this sched buffer - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) + # All 32 MMA warp threads arrive: Load warp can reuse this sched buffer + cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) sched_buf = Int32(1) - sched_buf should_continue = work_idx < total_work_units else: @@ -1261,6 +1264,8 @@ def kernel( wu_iter = Int32(0) if cutlass.const_expr(self.is_varlen and self.persistent): work_idx = sWorkIdx[(0,)] + # All 32 CC warp threads arrive on consumed_mbar[0]: buf 0 has been read + cute.arch.mbarrier_arrive(sched_consumed_mbar_base) sched_buf = Int32(1) sched_phase0 = Int32(1) sched_phase1 = Int32(0) @@ -1449,9 +1454,8 @@ def kernel( cute.arch.mbarrier_wait(sched_mbar_base + 1, sched_phase1) sched_phase1 = Int32(1) - sched_phase1 work_idx = sWorkIdx[(sched_buf,)] - # Signal consumed: CC has 4 warps, elect_one per warp → 4 arrives - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) + # Signal consumed: all 32 threads × 4 CC warps → 128 arrives + cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) sched_buf = Int32(1) - sched_buf should_continue = work_idx < total_work_units else: @@ -1474,6 +1478,8 @@ def kernel( wu_iter = Int32(0) if cutlass.const_expr(self.is_varlen and self.persistent): work_idx = sWorkIdx[(0,)] + # All 32 Store warp threads arrive on consumed_mbar[0]: buf 0 has been read + cute.arch.mbarrier_arrive(sched_consumed_mbar_base) sched_buf = Int32(1) sched_phase0 = Int32(1) sched_phase1 = Int32(0) @@ -1633,9 +1639,8 @@ def kernel( cute.arch.mbarrier_wait(sched_mbar_base + 1, sched_phase1) sched_phase1 = Int32(1) - sched_phase1 work_idx = sWorkIdx[(sched_buf,)] - # Signal consumed: Store warp → 1 arrive - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) + # Signal consumed: all 32 Store warp threads arrive + cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) sched_buf = Int32(1) - sched_buf should_continue = work_idx < total_work_units else: @@ -1650,6 +1655,8 @@ def kernel( # Dynamic scheduling: wait on double-buffered mbarriers for each WU if cutlass.const_expr(self.is_varlen and self.persistent): work_idx = sWorkIdx[(0,)] + # All 32 Empty warp threads arrive on consumed_mbar[0]: buf 0 has been read + cute.arch.mbarrier_arrive(sched_consumed_mbar_base) sched_buf = Int32(1) sched_phase0 = Int32(1) sched_phase1 = Int32(0) @@ -1661,9 +1668,8 @@ def kernel( cute.arch.mbarrier_wait(sched_mbar_base + 1, sched_phase1) sched_phase1 = Int32(1) - sched_phase1 work_idx = sWorkIdx[(sched_buf,)] - # Signal consumed: Empty warp → 1 arrive - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) + # Signal consumed: all 32 Empty warp threads arrive + cute.arch.mbarrier_arrive(sched_consumed_mbar_base + sched_buf) sched_buf = Int32(1) - sched_buf tmem.relinquish_alloc_permit() From 437dc21640ecc79bf001c6cbdf4ca641141ad301 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=97=A0=E8=A8=80=E7=8B=AC=E4=B8=8A=E6=9C=BA=E6=88=BF?= <88866917+sjmshsh@users.noreply.github.com> Date: Thu, 14 May 2026 11:47:31 +0800 Subject: [PATCH 12/34] [KDA] sm90 GVA enhance (#64) * sm90 * sm90 * sm90 * sm90 * sm90 * sm90 * sm90 * sm90 * sm90 * sm90 * sm90 * add sm90 benchmarks * benchmarks * benchmarks * fix * benchmark * benchmark * benchmark * benchmark * benchmark * benchmark * benchmark * benchmark * benchmark --------- Co-authored-by: sunnyxyli --- benchmarks/bench_kda_fused_fwd.py | 97 ++++++++++--- benchmarks/utils.py | 35 ++++- csrc/api/kda_sm90.cu | 42 ++++-- csrc/kda/sm90/collective/load_tma.hpp | 20 ++- csrc/kda/sm90/collective/mainloop_kda_fwd.hpp | 37 +++-- csrc/kda/sm90/collective/store_tma.hpp | 2 +- csrc/kda/sm90/kda_fwd_sm90.cu | 15 +- csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu | 4 + csrc/kda/sm90/kernel/kernel_kda_fwd.hpp | 5 +- csrc/kda/sm90/kernel/tile_scheduler.hpp | 46 ++++-- csrc/kda/sm90/prefill_kernel.hpp | 3 +- csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh | 23 +-- cula/kda/hopper_fused_fwd.py | 70 ++++++--- tests/test_kda_fused_fwd.py | 136 +++++++++++------- 14 files changed, 377 insertions(+), 158 deletions(-) diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index a87eed32..171c2bb9 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -29,8 +29,12 @@ - Fixed-length: various (B, T) configs - Varlen: sequences with 2-3x length variation +H (number of Q/K heads) is a module-level constant; HV (number of V heads) +defaults to H and can be overridden globally via --hv to run every config in +GVA (Grouped Value Attention) mode. HV must be a positive multiple of H. + Usage: - python bench_kda_fused_fwd.py [--mode fixed|varlen|both] [--ncu] + python bench_kda_fused_fwd.py [--mode fixed|varlen|both] [--hv HV] [--ncu] With --ncu, warmup=1 and iters=1 for ncu profiling: ncu --set full -o report python bench_kda_fused_fwd.py --mode varlen --ncu @@ -67,7 +71,11 @@ # ============================================================ # Constants # ============================================================ +# Default number of Q/K heads (H) and V heads (HV). When HV > H the run is in +# GVA mode (the kernel sees HV expanded q/k heads, prepared internally by +# prepare_safe_gate_inputs). HV is overridable globally via --hv. H, D = 64, 128 +HV = H WARMUP = 25 N_ITERS = 100 NCU_MODE = False @@ -159,7 +167,8 @@ def bench_fixed(configs): print("=" * 100) results = [] - for B, T in configs: + for cfg in configs: + B, T = cfg set_seed(SEED) device = torch.device("cuda") torch.cuda.empty_cache() @@ -167,7 +176,16 @@ def bench_fixed(configs): seq_lens = [T] * B cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=HAS_INIT_STATE) + inputs = prepare_safe_gate_inputs( + B, + T, + H, + D, + device, + cu_seqlens=cu_seqlens, + has_init_state=HAS_INIT_STATE, + num_v_heads=HV, + ) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -202,6 +220,8 @@ def bench_fixed(configs): { "B": B, "T": T, + "H": H, + "HV": HV, "rmse": rmse, "rel_max": rel_max, "mean_diff": mean_diff, @@ -226,7 +246,8 @@ def bench_varlen(configs): print("=" * 100) results = [] - for seq_lens, total_len, dist in configs: + for cfg in configs: + seq_lens, total_len, dist = cfg set_seed(SEED) device = torch.device("cuda") torch.cuda.empty_cache() @@ -234,7 +255,16 @@ def bench_varlen(configs): T = total_len cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=HAS_INIT_STATE) + inputs = prepare_safe_gate_inputs( + 1, + T, + H, + D, + device, + cu_seqlens=cu_seqlens, + has_init_state=HAS_INIT_STATE, + num_v_heads=HV, + ) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -276,6 +306,8 @@ def bench_varlen(configs): "dist": dist, "T_total": T, "n_seqs": n_seqs, + "H": H, + "HV": HV, "rmse": rmse, "rel_max": rel_max, "mean_diff": mean_diff, @@ -295,11 +327,13 @@ def bench_varlen(configs): # Report # ============================================================ def print_report(fixed_results, varlen_results): - sep = "=" * 110 + sep = "=" * 120 print(f"\n\n{sep}") print(" BENCHMARK REPORT: cula_kda_fused_fwd (fully-fused)") print(f" cuLA {_SM_TAG} fully-fused vs FLA Triton") - print(f" H={H} D={D} dtype=bf16 safe_gate=True has_init_state={HAS_INIT_STATE}") + print(f" D={D} dtype=bf16 safe_gate=True has_init_state={HAS_INIT_STATE}") + gva_note = f"GVA enabled (HV={HV} > H={H}, ratio={HV // H}x)" if HV > H else f"MHA (HV=H={H})" + print(f" {gva_note}") wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") @@ -308,35 +342,39 @@ def print_report(fixed_results, varlen_results): if fixed_results: print("\n [Fixed-Length]") - print(f" {'─' * 90}") + print(f" {'─' * 110}") print( - f" {'B':>3s} {'T':>6s} │ {'RMSE':>10s} {'rel_max':>10s} {'mean_diff':>10s}" - f" │ {'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" + f" {'B':>3s} {'T':>6s} {'H':>3s} {'HV':>3s} {'GVA':>4s} │ " + f"{'RMSE':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" ) - print(f" {'─' * 90}") + print(f" {'─' * 110}") for r in fixed_results: + gva_tag = f"{r['HV'] // r['H']}x" if r["HV"] > r["H"] else "no" print( - f" {r['B']:3d} {r['T']:6d} │ " + f" {r['B']:3d} {r['T']:6d} {r['H']:3d} {r['HV']:3d} {gva_tag:>4s} │ " f"{r['rmse']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" ) - print(f" {'─' * 90}") + print(f" {'─' * 110}") if varlen_results: print("\n [Varlen]") - print(f" {'─' * 105}") + print(f" {'─' * 120}") print( - f" {'Config':>45s} │ {'RMSE':>10s} {'rel_max':>10s} {'mean_diff':>10s}" - f" │ {'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" + f" {'Config':>45s} {'H':>3s} {'HV':>3s} {'GVA':>4s} │ " + f"{'RMSE':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" ) - print(f" {'─' * 105}") + print(f" {'─' * 120}") for r in varlen_results: + gva_tag = f"{r['HV'] // r['H']}x" if r["HV"] > r["H"] else "no" print( - f" {r['tag']:>45s} │ " + f" {r['tag']:>45s} {r['H']:3d} {r['HV']:3d} {gva_tag:>4s} │ " f"{r['rmse']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" ) - print(f" {'─' * 105}") + print(f" {'─' * 120}") print(f"\n{sep}\n") @@ -351,7 +389,7 @@ def main(): type=str, default="both", choices=["fixed", "varlen", "both"], - help="Which benchmark mode to run (default: both)", + help="Which benchmark mode to run (default: both).", ) parser.add_argument( "--ncu", @@ -368,9 +406,15 @@ def main(): action="store_true", help="Use non-zero initial state (default: False)", ) + parser.add_argument( + "--hv", + type=int, + default=None, + help=f"Override number of V heads (HV). Default: H ({H}, no GVA). Set HV > H to run all configs in GVA mode.", + ) args = parser.parse_args() - global NCU_MODE, SANITIZER_MODE, HAS_INIT_STATE + global NCU_MODE, SANITIZER_MODE, HAS_INIT_STATE, HV if args.ncu: NCU_MODE = True print("[NCU mode] warmup=1, iters=1") @@ -380,11 +424,21 @@ def main(): if args.init_state: HAS_INIT_STATE = True print("[init_state] using non-zero initial state") + if args.hv is not None: + if args.hv < H or args.hv % H != 0: + raise ValueError(f"--hv must be a positive multiple of H ({H}), got {args.hv}") + HV = args.hv + if HV > H: + print(f"[GVA] HV={HV} (H={H}, ratio={HV // H}x)") print( f"[Device] {torch.cuda.get_device_name(0)} compute capability {_SM_TAG} → using {cula_kda_fused_fwd.__module__}.{cula_kda_fused_fwd.__name__}" ) + # ------------------------------------------------------------------ + # Fixed-length configs — (B, T). Per-row H/HV defaults to global H/HV + # (HV overridable via --hv to switch all rows into GVA mode). + # ------------------------------------------------------------------ fixed_configs = [ # (B, T) (1, 512), @@ -399,6 +453,7 @@ def main(): (2, 16384), ] + # Varlen configs — same layout as fixed; HV is controlled globally via --hv. varlen_configs = build_varlen_configs( num_seqs_list=(10, 20), total_lens=(4096, 8192, 16384), diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 29bab04d..bfd0761e 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -256,25 +256,46 @@ def build_varlen_configs( def prepare_safe_gate_inputs( - batch_size, T, H, D, device, cu_seqlens=None, chunk_size=CHUNK_SIZE, seed=SEED, has_init_state=False + batch_size, + T, + H, + D, + device, + cu_seqlens=None, + chunk_size=CHUNK_SIZE, + seed=SEED, + has_init_state=False, + num_v_heads=None, ): """Prepare inputs for safe_gate benchmarks (use_gate_in_kernel=True, safe_gate=True). All tensors are flattened to (1, B*T, ...) for cu_seqlens compatibility. """ + HV = H if num_v_heads is None else num_v_heads + assert HV >= H and HV % H == 0, f"HV ({HV}) must be a positive multiple of H ({H}) with HV >= H." + dtype = torch.bfloat16 scale = D ** (-0.5) set_seed(seed) + # Allocate native GVA shapes: q = torch.randn(batch_size, T, H, D, dtype=dtype, device=device).requires_grad_(False) k = torch.randn(batch_size, T, H, D, dtype=dtype, device=device).requires_grad_(False) - v = torch.randn(batch_size, T, H, D, dtype=dtype, device=device).requires_grad_(False) - g = torch.randn(batch_size, T, H, D, dtype=dtype, device=device).requires_grad_(False) - beta = torch.randn(batch_size, T, H, dtype=torch.float, device=device).sigmoid().requires_grad_(False) + v = torch.randn(batch_size, T, HV, D, dtype=dtype, device=device).requires_grad_(False) + g = torch.randn(batch_size, T, HV, D, dtype=dtype, device=device).requires_grad_(False) + beta = torch.randn(batch_size, T, HV, dtype=torch.float, device=device).sigmoid().requires_grad_(False) + + # GVA expansion: bring q/k up to HV heads so all tensors share head dim. + group = HV // H + if group > 1: + q = q.repeat_interleave(group, dim=2).contiguous() + k = k.repeat_interleave(group, dim=2).contiguous() - A_log = torch.randn(H, dtype=torch.float, device=device).requires_grad_(False) - dt_bias = torch.randn(H * D, dtype=torch.float, device=device).requires_grad_(False) + # A_log / dt_bias must match the head count of `g` (HV), otherwise + # kda_gate_chunk_cumsum would index out of bounds for i_h >= H. + A_log = torch.randn(HV, dtype=torch.float, device=device).requires_grad_(False) + dt_bias = torch.randn(HV * D, dtype=torch.float, device=device).requires_grad_(False) # flatten to batch_size=1 for cu_seqlens compatibility if batch_size != 1: @@ -285,7 +306,7 @@ def prepare_safe_gate_inputs( init_state = None if has_init_state: num_seqs = cu_seqlens.shape[0] - 1 if cu_seqlens is not None else batch_size - init_state = torch.randn(num_seqs, H, D, D, dtype=torch.float, device=device).requires_grad_(False) + init_state = torch.randn(num_seqs, HV, D, D, dtype=torch.float, device=device).requires_grad_(False) return dict( q=q, diff --git a/csrc/api/kda_sm90.cu b/csrc/api/kda_sm90.cu index e8f3a545..9e016eb1 100644 --- a/csrc/api/kda_sm90.cu +++ b/csrc/api/kda_sm90.cu @@ -36,21 +36,32 @@ kda_fwd_prefill( float scale, bool output_final_state, bool safe_gate) { - // Q, K, V: [packed_seq, H, D] (already packed by Python layer) + // Q, K: [packed_seq, num_qk_heads, D] + // V/O/g: [packed_seq, num_v_heads, D] (GVA: num_v_heads is a positive integer multiple of num_qk_heads) auto packed_seq = q.size(0); - auto num_heads = q.size(1); + auto num_qk_heads = q.size(1); + auto num_v_heads = v.size(1); auto head_size = q.size(2); auto num_seqs = cu_seqlens.size(0) - 1; - // KDA constraint: all head counts must be the same - TORCH_CHECK(num_heads == k.size(1), "KDA requires num_q_heads == num_k_heads, got ", num_heads, " vs ", k.size(1)); - TORCH_CHECK(num_heads == v.size(1), "KDA requires num_q_heads == num_v_heads, got ", num_heads, " vs ", v.size(1)); + // GVA contract on the C++ side. Order matters: check positivity *before* the modulo to + // avoid % 0 / division-by-zero UB in case the Python layer passed a degenerate shape. + TORCH_CHECK(num_qk_heads > 0, "KDA requires num_qk_heads > 0, got ", num_qk_heads); + TORCH_CHECK(num_v_heads > 0, "KDA requires num_v_heads > 0, got ", num_v_heads); + TORCH_CHECK( + num_qk_heads == k.size(1), "KDA requires num_q_heads == num_k_heads, got ", num_qk_heads, " vs ", k.size(1)); + TORCH_CHECK( + num_v_heads % num_qk_heads == 0, + "KDA GVA requires num_v_heads to be a positive multiple of num_qk_heads, got num_v_heads=", + num_v_heads, + ", num_qk_heads=", + num_qk_heads); TORCH_CHECK(head_size == v.size(2), "KDA requires Q and V head dim to match, got ", head_size, " vs ", v.size(2)); // Allocate output if not provided torch::Tensor output = output_.has_value() ? output_.value() : torch::empty( - {packed_seq, num_heads, head_size}, + {packed_seq, num_v_heads, head_size}, torch::TensorOptions().dtype(q.dtype()).device(q.device())); // output_final_state controls the API side effect. If it is false, ignore @@ -61,7 +72,7 @@ kda_fwd_prefill( output_state = output_state_.has_value() ? output_state_.value() : torch::zeros( - {num_seqs, num_heads, head_size, head_size}, + {num_seqs, num_v_heads, head_size, head_size}, torch::TensorOptions().dtype(torch::kFloat32).device(q.device())); } @@ -92,8 +103,8 @@ kda_fwd_prefill( TORCH_CHECK(alpha.dtype() == torch::kFloat32, "alpha must be float32"); TORCH_CHECK(alpha.is_contiguous(), "alpha must be contiguous"); TORCH_CHECK( - alpha.size(0) == packed_seq && alpha.size(1) == num_heads && alpha.size(2) == head_size, - "alpha shape must be [packed_seq, num_heads, head_size]"); + alpha.size(0) == packed_seq && alpha.size(1) == num_v_heads && alpha.size(2) == head_size, + "alpha shape must be [packed_seq, num_v_heads, head_size]"); alpha_ptr = alpha.data_ptr(); } @@ -106,12 +117,17 @@ kda_fwd_prefill( beta.dtype()); TORCH_CHECK(beta.is_contiguous(), "beta must be contiguous"); TORCH_CHECK( - beta.size(0) == packed_seq && beta.size(1) == num_heads, "beta shape must be [packed_seq, num_heads]"); + beta.size(0) == packed_seq && beta.size(1) == num_v_heads, "beta shape must be [packed_seq, num_v_heads]"); } if (input_state_.has_value()) { auto& input_state = input_state_.value(); TORCH_CHECK(input_state.dtype() == torch::kFloat32, "input_state must be float32"); TORCH_CHECK(input_state.is_contiguous(), "input_state must be contiguous"); + // Defense in depth: also enforce shape on the C++ side (Python layer should already check). + TORCH_CHECK( + input_state.dim() == 4 && input_state.size(0) == num_seqs && input_state.size(1) == num_v_heads && + input_state.size(2) == head_size && input_state.size(3) == head_size, + "input_state shape must be [num_seqs, num_v_heads, head_size, head_size]"); input_state_ptr = input_state.data_ptr(); } @@ -142,7 +158,8 @@ kda_fwd_prefill( cu_seqlens.data_ptr(), workspace_buffer.data_ptr(), static_cast(num_seqs), - static_cast(num_heads), + static_cast(num_qk_heads), + static_cast(num_v_heads), static_cast(head_size), static_cast(packed_seq), scale, @@ -163,7 +180,8 @@ kda_fwd_prefill( cu_seqlens.data_ptr(), workspace_buffer.data_ptr(), static_cast(num_seqs), - static_cast(num_heads), + static_cast(num_qk_heads), + static_cast(num_v_heads), static_cast(head_size), static_cast(packed_seq), scale, diff --git a/csrc/kda/sm90/collective/load_tma.hpp b/csrc/kda/sm90/collective/load_tma.hpp index d7e4c8fc..1d427b0c 100644 --- a/csrc/kda/sm90/collective/load_tma.hpp +++ b/csrc/kda/sm90/collective/load_tma.hpp @@ -92,10 +92,11 @@ struct CollectiveLoadTma { work_desc.seq_idx, work_desc.q_head_idx(), work_desc.tok_offset); + // Q lives in the QK head space. Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( problem_size.total_seqlen, problem_size.head_size, - problem_size.num_heads)); // global view to the packed varlen sequence + problem_size.num_qk_heads)); // global view to the packed varlen sequence Tensor m_varlen = m_varlen_head(_, _, work_desc.q_head_idx()); // slice into current head_idx Tensor m_offset = domain_offset( make_coord(work_desc.tok_offset, _0{}), @@ -103,18 +104,19 @@ struct CollectiveLoadTma { Tensor g_full = local_tile(m_offset, make_tile(BlkSeqQ, HeadSize), make_coord(_, _0{})); // (blk, d, iter_blk) return g_full; - } else if constexpr (kind == LoadKind::kAlpha) { // same as Q currently + } else if constexpr (kind == LoadKind::kAlpha) { + // Alpha (gate) is per V/O head under GVA. DPRINTF0_W( "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", to_string(kind), work_desc.seq_idx, - work_desc.q_head_idx(), + work_desc.o_head_idx(), work_desc.tok_offset); Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( problem_size.total_seqlen, problem_size.head_size, - problem_size.num_heads)); // global view to the packed varlen sequence - Tensor m_varlen = m_varlen_head(_, _, work_desc.q_head_idx()); // slice into current head_idx + problem_size.num_v_heads)); // global view to the packed varlen sequence + Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx Tensor m_offset = domain_offset( make_coord(work_desc.tok_offset, _0{}), m_varlen); // offset to start of the current sequence @@ -122,7 +124,11 @@ struct CollectiveLoadTma { local_tile(m_offset, make_tile(BlkSeqQ, HeadSize), make_coord(_, _0{})); // (blk, d, iter_blk) return g_full; } else { - auto head_idx = (kind == LoadKind::kK ? work_desc.k_head_idx() : work_desc.v_head_idx()); + // K lives in the QK head space; V lives in the V head space. + // `kind` is a static constexpr LoadKind, so the head-count selection collapses at compile time. + constexpr bool kIsK = (kind == LoadKind::kK); + auto head_idx = kIsK ? work_desc.k_head_idx() : work_desc.v_head_idx(); + auto num_kv_heads = kIsK ? problem_size.num_qk_heads : problem_size.num_v_heads; DPRINTF0_W( "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", to_string(kind), @@ -132,7 +138,7 @@ struct CollectiveLoadTma { Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( problem_size.head_size, problem_size.total_seqlen, - problem_size.num_heads)); // global view to the packed varlen sequence + num_kv_heads)); // global view to the packed varlen sequence Tensor m_varlen = m_varlen_head(_, _, head_idx); // slice into current head_idx Tensor m_offset = domain_offset( make_coord(_0{}, work_desc.tok_offset), diff --git a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp index 1f781693..301dbfd4 100644 --- a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp +++ b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp @@ -471,7 +471,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { Element const* ptr_V; LayoutV dV; Element* ptr_O; LayoutO dO; float const* ptr_Alpha; LayoutAlpha dAlpha; - float* ptr_output_state; // layout fixed (kdim, vdim, num_heads, num_seqs):LayoutLeft{} + float* ptr_output_state; // layout fixed (kdim, vdim, num_v_heads, num_seqs):LayoutLeft{} float const* ptr_input_state; float scale; ElementBetaGmem const* beta_ptr; GmemStrideBeta beta_stride; @@ -506,15 +506,17 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { int64_t t = problem_size.total_seqlen; int32_t d = problem_size.head_size; + // GVA: Q/K are sized by num_qk_heads; V, alpha (gate), O, beta and the recurrent state + // are sized by num_v_heads. auto params_qk = CollectiveMmaQK::to_underlying_arguments( - make_shape(s, t, d, problem_size.num_heads), + make_shape(s, t, d, problem_size.num_qk_heads), typename CollectiveMmaQK::Arguments{ args.ptr_Q, args.dQ, args.ptr_K, args.dK, // never used, dummy }, /*workspace=*/nullptr); auto params_kv_k = CollectiveMmaKV_G2S::to_underlying_arguments( - make_shape(d, d, s, problem_size.num_heads), + make_shape(d, d, s, problem_size.num_qk_heads), typename CollectiveMmaKV_G2S::Arguments{ args.ptr_V, select<1, 0, 2>(args.dV), // not used @@ -523,7 +525,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }, /*workspace=*/nullptr); - auto alpha_shape = make_shape(s, d, problem_size.num_heads); + auto alpha_shape = make_shape(s, d, problem_size.num_v_heads); auto alpha_stride = make_stride( get<0>(args.dAlpha), // seqlen stride get<1>(args.dAlpha), // head_dim stride @@ -538,7 +540,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { size<0>(ClusterShape{})); auto params_kv_v = CollectiveMmaKV_G2S::to_underlying_arguments( - make_shape(d, d, s, problem_size.num_heads), + make_shape(d, d, s, problem_size.num_v_heads), typename CollectiveMmaKV_G2S::Arguments{ args.ptr_V, select<1, 0, 2>(args.dV), // used as G2S for V @@ -548,8 +550,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { /*workspace=*/nullptr); auto params_o = CollectiveStoreO::to_underlying_arguments( - make_shape(d, s, d, problem_size.num_heads), // in O1 - // make_shape(d, s, s, problem_size.num_heads), // in O2 + make_shape(d, s, d, problem_size.num_v_heads), // in O1 + // make_shape(d, s, s, problem_size.num_v_heads), // in O2 typename CollectiveStoreO::Arguments{args.ptr_O, select<1, 0, 2>(args.dO), workspace}, workspace); @@ -567,7 +569,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // TODO: refactor all name to varname_vartype .beta_ptr = args.beta_ptr, - .beta_layout = make_layout(make_shape(s, problem_size.num_heads), args.beta_stride), + .beta_layout = make_layout(make_shape(s, problem_size.num_v_heads), args.beta_stride), }; } @@ -899,7 +901,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto kv_load = [&](auto& tKVrKV) INLINE_LAMBDA { DPRINTF0_WG("[%d,%d,%d,%d]>> load tKVgKV -> tKVrKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); - int num_state_heads = problem_size.num_heads; + // GVA: state is stored per V/O head. + int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); auto gKV = make_tensor( make_gmem_ptr(params.ptr_input_state), @@ -914,8 +917,22 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }; auto kv_store = [&]() INLINE_LAMBDA { // tKVrKV is carried over whole mainloop + // Skip the final-state write-back entirely when the caller does not need it + // (i.e. output_final_state=False on the public API). The check is uniform + // across the launch, so the branch cost is negligible while it saves a full + // [N, num_v_heads, D, D] float32 store to GMEM. + if (params.ptr_output_state == nullptr) { + DPRINTF0_WG( + "[%d,%d,%d,%d]>> skip tKVrKV -> tKVgKV (output_final_state=false)\n", + seq_idx, + q_head_idx, + k_head_idx, + v_head_idx); + return; + } DPRINTF0_WG("[%d,%d,%d,%d]>> save tKVrKV -> tKVgKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); - int num_state_heads = problem_size.num_heads; + // GVA: state is stored per V/O head. + int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); auto gKV = make_tensor( make_gmem_ptr(params.ptr_output_state), diff --git a/csrc/kda/sm90/collective/store_tma.hpp b/csrc/kda/sm90/collective/store_tma.hpp index f7b986ab..0f7f7c1a 100644 --- a/csrc/kda/sm90/collective/store_tma.hpp +++ b/csrc/kda/sm90/collective/store_tma.hpp @@ -195,7 +195,7 @@ struct CollectiveStoreTma { Tensor m_varlen_head = tma_store_.get_tma_tensor(make_shape( problem_size.head_size, problem_size.total_seqlen, - problem_size.num_heads)); // global view to the packed varlen sequence + problem_size.num_v_heads)); // O lives in the V/O head space under GVA Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx Tensor m_offset = domain_offset( make_coord(_0{}, work_desc.tok_offset), diff --git a/csrc/kda/sm90/kda_fwd_sm90.cu b/csrc/kda/sm90/kda_fwd_sm90.cu index c13b1bd0..d668db9b 100644 --- a/csrc/kda/sm90/kda_fwd_sm90.cu +++ b/csrc/kda/sm90/kda_fwd_sm90.cu @@ -48,7 +48,8 @@ launch_kda_fwd_prefill_kernel_gbai( int32_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs, - int32_t num_heads, + int32_t num_qk_heads, + int32_t num_v_heads, int32_t head_size, int64_t total_seqlen, float scale, @@ -74,7 +75,8 @@ launch_kda_fwd_prefill_kernel( int32_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs, - int32_t num_heads, + int32_t num_qk_heads, + int32_t num_v_heads, int32_t head_size, int64_t total_seqlen, float scale, @@ -98,7 +100,8 @@ launch_kda_fwd_prefill_kernel( cu_seqlens, \ workspace_buffer, \ num_seqs, \ - num_heads, \ + num_qk_heads, \ + num_v_heads, \ head_size, \ total_seqlen, \ scale, \ @@ -137,7 +140,8 @@ launch_kda_fwd_prefill_kernel( int32_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs, - int32_t num_heads, + int32_t num_qk_heads, + int32_t num_v_heads, int32_t head_size, int64_t total_seqlen, float scale, @@ -159,7 +163,8 @@ launch_kda_fwd_prefill_kernel( int32_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs, - int32_t num_heads, + int32_t num_qk_heads, + int32_t num_v_heads, int32_t head_size, int64_t total_seqlen, float scale, diff --git a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu index 93fe8693..309cefa0 100644 --- a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu +++ b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu @@ -40,6 +40,7 @@ launch_kda_fwd_prefill_kernel_gbai CUTE_DEVICE bool @@ -42,11 +43,11 @@ struct WorkDesc { CUTE_DEVICE int32_t q_head_idx() const { - return head_idx; + return qk_head_idx; } CUTE_DEVICE int32_t k_head_idx() const { - return head_idx; + return qk_head_idx; } CUTE_DEVICE int32_t v_head_idx() const { @@ -64,11 +65,15 @@ struct WorkDesc { } }; +// Each block handles a single (seq, v_head) work item; CTAs do not cooperate. +// GVA optimization: heads_per_group is precomputed on the host and stored in +// Params, so the device side does not redo the integer division per CTA. struct IndividualTileScheduler { struct Params { dim3 grid; int32_t num_seqs; - int32_t num_heads; + int32_t num_v_heads; + int32_t heads_per_group; // = num_v_heads / num_qk_heads, precomputed on host }; bool scheduled = false; // a once flag @@ -84,19 +89,26 @@ struct IndividualTileScheduler { cutlass::KernelHardwareInfo const& hw_info, ClusterShape const& cluster_shape, TileShape const& tile_shape) { + // Compute heads_per_group once on the host so every CTA does not have to redo + // the integer division. + int32_t const heads_per_group = problem_size.num_v_heads / problem_size.num_qk_heads; dim3 grid(0, 1, 1); - grid.x = problem_size.num_seqs * problem_size.num_heads; + grid.x = problem_size.num_seqs * problem_size.num_v_heads; DPRINTF( - "to_underlying_arguments: grid:{.x:%d, .y:%d, .z:%d}, num_seqs:%d, num_heads:%d\n", + "to_underlying_arguments: grid:{.x:%d, .y:%d, .z:%d}, num_seqs:%d, num_qk_heads:%d, num_v_heads:%d, " + "heads_per_group:%d\n", grid.x, grid.y, grid.z, problem_size.num_seqs, - problem_size.num_heads); + problem_size.num_qk_heads, + problem_size.num_v_heads, + heads_per_group); return { .grid = grid, .num_seqs = problem_size.num_seqs, - .num_heads = problem_size.num_heads, + .num_v_heads = problem_size.num_v_heads, + .heads_per_group = heads_per_group, }; } @@ -108,8 +120,10 @@ struct IndividualTileScheduler { template CUTE_DEVICE WorkDesc get_next_work(Params params, ProblemSize const& problem_size) { - int32_t seq_idx = blockIdx.x / params.num_heads; - int32_t head_idx = blockIdx.x % params.num_heads; + int32_t seq_idx = blockIdx.x / params.num_v_heads; + int32_t head_idx = blockIdx.x % params.num_v_heads; + // GVA: use the host-precomputed heads_per_group to avoid device-side division. + int32_t qk_head_idx = head_idx / params.heads_per_group; int32_t s = problem_size.cu_seqlens[seq_idx]; int32_t e = problem_size.cu_seqlens[seq_idx + 1]; @@ -120,8 +134,9 @@ struct IndividualTileScheduler { } else { scheduled = true; DPRINTF0_W( - "get_next_work: this_work={seq_idx:%d head_idx:%d tok_offset:%lld seq_len:%lld}\n", + "get_next_work: this_work={seq_idx:%d qk_head_idx:%d head_idx:%d tok_offset:%lld seq_len:%lld}\n", seq_idx, + qk_head_idx, head_idx, s, seq_len); @@ -129,6 +144,7 @@ struct IndividualTileScheduler { return { .seq_idx = seq_idx, + .qk_head_idx = qk_head_idx, .head_idx = head_idx, .tok_offset = s, .seq_len = seq_len, diff --git a/csrc/kda/sm90/prefill_kernel.hpp b/csrc/kda/sm90/prefill_kernel.hpp index 00cef2df..d56fafae 100644 --- a/csrc/kda/sm90/prefill_kernel.hpp +++ b/csrc/kda/sm90/prefill_kernel.hpp @@ -40,7 +40,8 @@ launch_kda_fwd_prefill_kernel( int32_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs, - int32_t num_heads, + int32_t num_qk_heads, + int32_t num_v_heads, int32_t head_size, int64_t total_seqlen, float scale, diff --git a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh index 2fb9bda9..72f13a6f 100644 --- a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh +++ b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh @@ -53,7 +53,8 @@ launch_kda_fwd_prefill_kernel_gbai( int32_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs, - int32_t num_heads, + int32_t num_qk_heads, + int32_t num_v_heads, int32_t head_size, int64_t total_seqlen, float scale, @@ -105,8 +106,11 @@ launch_kda_fwd_prefill_kernel_gbai( using Arguments = typename Operation::Arguments; // NOTE: LayoutQ/K/V in (seq, head_size, (b,h)) coordinate semantics + // GVA: Q/K rows are packed as [packed_seq, num_qk_heads, head_size]; + // V/O/g/beta rows are packed as [packed_seq, num_v_heads, head_size]. - int32_t tok_stride = num_heads * head_size; + int32_t qk_tok_stride = num_qk_heads * head_size; + int32_t v_tok_stride = num_v_heads * head_size; int32_t head_stride = head_size; Operation op; @@ -116,21 +120,22 @@ launch_kda_fwd_prefill_kernel_gbai( .cu_seqlens = cu_seqlens, .total_seqlen = total_seqlen, .num_seqs = num_seqs, - .num_heads = num_heads, + .num_qk_heads = num_qk_heads, + .num_v_heads = num_v_heads, .head_size = head_size, }, .mainloop = { // clang-format off - .ptr_Q = (T*)q, .dQ = {tok_stride, _1{}, head_stride}, - .ptr_K = (T*)k, .dK = {tok_stride, _1{}, head_stride}, - .ptr_V = (T*)v, .dV = {tok_stride, _1{}, head_stride}, - .ptr_O = (T*)output, .dO = {tok_stride, _1{}, head_stride}, - .ptr_Alpha = alpha, .dAlpha = {tok_stride, _1{}, head_stride}, + .ptr_Q = (T*)q, .dQ = {qk_tok_stride, _1{}, head_stride}, + .ptr_K = (T*)k, .dK = {qk_tok_stride, _1{}, head_stride}, + .ptr_V = (T*)v, .dV = {v_tok_stride, _1{}, head_stride}, + .ptr_O = (T*)output, .dO = {v_tok_stride, _1{}, head_stride}, + .ptr_Alpha = alpha, .dAlpha = {v_tok_stride, _1{}, head_stride}, .ptr_output_state = (float*)output_state, .ptr_input_state = (float*)input_state, .scale = scale, - .beta_ptr = beta, .beta_stride = {num_heads, 1}, + .beta_ptr = beta, .beta_stride = {num_v_heads, 1}, }, // clang-format on .hw_info = hw_info}; diff --git a/cula/kda/hopper_fused_fwd.py b/cula/kda/hopper_fused_fwd.py index cc42827a..c0399bbe 100644 --- a/cula/kda/hopper_fused_fwd.py +++ b/cula/kda/hopper_fused_fwd.py @@ -49,9 +49,18 @@ def forward( chunk_indices: torch.IntTensor | None = None, ): chunk_size = 64 - assert q.shape[-2] == v.shape[-2] == k.shape[-2], "Number of heads must be the same for q, k, v." + # GVA: q/k share num_qk_heads; v/g/beta share num_v_heads. + # num_v_heads must be a positive multiple of num_qk_heads (heads_per_group = HV / H). + assert q.shape == k.shape, "q and k must have the same shape." + assert q.shape[:2] == v.shape[:2] == g.shape[:2], "q, k, v, g must share batch and sequence dimensions." - batch_size, seq_len, num_heads, head_dim = q.shape + batch_size, seq_len, num_qk_heads, head_dim = q.shape + num_v_heads = v.shape[-2] + assert num_qk_heads > 0, f"num_qk_heads must be positive, got {num_qk_heads}." + assert num_v_heads > 0, f"num_v_heads must be positive, got {num_v_heads}." + assert num_v_heads % num_qk_heads == 0, ( + f"num_v_heads ({num_v_heads}) must be a positive multiple of num_qk_heads ({num_qk_heads})." + ) if cu_seqlens is None: cu_seqlens = prepare_uniform_cu_seqlens(batch_size, seq_len, q.device, torch.int32) @@ -88,13 +97,13 @@ def forward( q, q_rstd = l2norm_fwd(q) k, k_rstd = l2norm_fwd(k) - # reshape to packed [T, H, K] for the C++ kernel + # reshape q/k to packed [T, H, K] and v/g to [T, HV, K], beta to [T, HV] for the C++ kernel packed_seq = batch_size * seq_len - q = q.reshape(packed_seq, num_heads, head_dim).contiguous() - k = k.reshape(packed_seq, num_heads, head_dim).contiguous() - v = v.reshape(packed_seq, num_heads, head_dim).contiguous() - g = g.reshape(packed_seq, num_heads, head_dim).contiguous() - beta = beta.reshape(packed_seq, num_heads).contiguous() + q = q.reshape(packed_seq, num_qk_heads, head_dim).contiguous() + k = k.reshape(packed_seq, num_qk_heads, head_dim).contiguous() + v = v.reshape(packed_seq, num_v_heads, head_dim).contiguous() + g = g.reshape(packed_seq, num_v_heads, head_dim).contiguous() + beta = beta.reshape(packed_seq, num_v_heads).contiguous() # workspace buffer for TMA Store O tensormap sm_count = get_device_sm_count(q.device) @@ -159,19 +168,19 @@ def cula_kda_prefill( k (torch.Tensor): keys of shape `[B, T, H, K]`. v (torch.Tensor): - values of shape `[B, T, H, V]`. + values of shape `[B, T, HV, K]`. g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, H, K]`. + (forget) gating tensor (in log space!) of shape `[B, T, HV, K]`. beta (torch.Tensor): - betas of shape `[B, T, H]`. + betas of shape `[B, T, HV]`. scale (Optional[float]): Scale factor for the KDA attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + If not provided, it will default to `1 / sqrt(D)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. + Initial state of shape `[N, HV, K, K]` for `N` input sequences. Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + Whether to output the final state of shape `[N, HV, K, K]`. Default: `False`. use_qk_l2norm_in_kernel (bool): Whether to apply L2norm to the q,k tensor internally. Default: `False`. use_gate_in_kernel (bool): @@ -189,9 +198,9 @@ def cula_kda_prefill( Returns: o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. + Outputs of shape `[B, T, HV, K]`. final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + Final state of shape `[N, HV, K, K]` if `output_final_state=True` else `None`. """ assert_hopper() assert safe_gate, "Only support safe_gate=True." @@ -219,9 +228,32 @@ def cula_kda_prefill( if not (-5 <= lower_bound < 0): raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." - assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." - assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." + assert q.shape == k.shape, "q and k must have the same shape." + assert q.shape[:2] == v.shape[:2] == g.shape[:2], "q, k, v, g must share batch and sequence dimensions." + + batch_size, seq_len, num_qk_heads, head_dim = q.shape + num_v_heads = v.shape[-2] + # Order matters here: positivity *first*, modulo second, to avoid ZeroDivisionError on bad inputs. + assert num_qk_heads > 0, f"num_qk_heads must be positive, got {num_qk_heads}." + assert num_v_heads > 0, f"num_v_heads must be positive, got {num_v_heads}." + assert num_v_heads % num_qk_heads == 0, ( + f"num_v_heads ({num_v_heads}) must be a positive multiple of num_qk_heads ({num_qk_heads})." + ) + assert g.shape == (batch_size, seq_len, num_v_heads, head_dim), ( + f"g must have shape (B, T, HV, D)=({batch_size}, {seq_len}, {num_v_heads}, {head_dim}), got {tuple(g.shape)}." + ) + assert v.shape == (batch_size, seq_len, num_v_heads, head_dim), ( + f"v must have shape (B, T, HV, D)=({batch_size}, {seq_len}, {num_v_heads}, {head_dim}), got {tuple(v.shape)}." + ) + assert beta.shape == (batch_size, seq_len, num_v_heads), ( + f"beta must have shape (B, T, HV)=({batch_size}, {seq_len}, {num_v_heads}), got {tuple(beta.shape)}." + ) + if initial_state is not None: + expected_num_states = (len(cu_seqlens) - 1) if cu_seqlens is not None else batch_size + assert initial_state.shape == (expected_num_states, num_v_heads, head_dim, head_dim), ( + f"initial_state must have shape (N, HV, D, D)=" + f"({expected_num_states}, {num_v_heads}, {head_dim}, {head_dim}), got {tuple(initial_state.shape)}." + ) assert q.dtype == k.dtype == v.dtype == torch.bfloat16, "q, k, v must be in bfloat16." assert beta.dtype == torch.bfloat16 or beta.dtype == torch.float32, "beta must be in bfloat16 or float32." assert q.shape[-1] == k.shape[-1] == v.shape[-1] == 128, "Currently we only support head dim of 128 for KDA" diff --git a/tests/test_kda_fused_fwd.py b/tests/test_kda_fused_fwd.py index 9c325524..354f0c9c 100644 --- a/tests/test_kda_fused_fwd.py +++ b/tests/test_kda_fused_fwd.py @@ -36,28 +36,34 @@ "B", "T", "H", + "HV", "D", "gate_logit_normalizer", "mask_p", "use_qk_l2norm_in_kernel", "use_gate_in_kernel", "safe_gate", + "use_initial_state", "dtype", ), [ pytest.param( *test, - id="B{}-T{}-H{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), + id=("B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-init{}-{}").format(*test), ) for test in [ - (1, 63, 1, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 500, 3, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 1000, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), - (3, 1024, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 128, 1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (4, 2048, 8, 128, 1, 0, False, True, True, torch.bfloat16), + (1, 63, 1, 1, 128, 1, 0, False, False, True, True, torch.bfloat16), + (2, 500, 3, 3, 128, 1, 0, False, False, True, True, torch.bfloat16), + (2, 1000, 3, 3, 128, 1, 0.5, False, False, True, True, torch.bfloat16), + (3, 1024, 4, 4, 128, 0.1, 0, False, False, True, True, torch.bfloat16), + (4, 1024, 4, 4, 128, 1, 0, False, False, True, True, torch.bfloat16), + (4, 1024, 4, 4, 128, 1, 0, True, False, True, True, torch.bfloat16), + (2, 1500, 4, 4, 128, 10, 0, False, True, True, True, torch.bfloat16), + (4, 2048, 8, 8, 128, 1, 0, False, True, True, True, torch.bfloat16), + (2, 512, 2, 4, 128, 1, 0, False, False, True, True, torch.bfloat16), + (2, 1024, 2, 8, 128, 1, 0, False, True, True, True, torch.bfloat16), + (1, 64, 1, 2, 128, 1, 0, False, False, True, True, torch.bfloat16), + (1, 65, 1, 4, 128, 1, 0, False, False, True, False, torch.bfloat16), ] ], ) @@ -65,12 +71,14 @@ def test_safe_gate_chunk( B: int, T: int, H: int, + HV: int, D: int, gate_logit_normalizer: float, mask_p: float, use_qk_l2norm_in_kernel: bool, use_gate_in_kernel: bool, safe_gate: bool, + use_initial_state: bool, dtype: torch.dtype, beta_dtype: torch.dtype, ): @@ -81,11 +89,11 @@ def test_safe_gate_chunk( torch.manual_seed(42) q = torch.rand(B, T, H, D, dtype=dtype) k = torch.rand(B, T, H, D, dtype=dtype) - v = torch.rand(B, T, H, D, dtype=dtype) - g = torch.randn(B, T, H, D, dtype=torch.float if not use_gate_in_kernel else dtype) + v = torch.rand(B, T, HV, D, dtype=dtype) + g = torch.randn(B, T, HV, D, dtype=torch.float if not use_gate_in_kernel else dtype) if use_gate_in_kernel: - A_log = torch.randn(H, dtype=torch.float) - dt_bias = torch.randn(H * D, dtype=torch.float) + A_log = torch.randn(HV, dtype=torch.float) + dt_bias = torch.randn(HV * D, dtype=torch.float) else: g = F.logsigmoid(g) / gate_logit_normalizer g = g * (torch.rand_like(g) > mask_p) @@ -98,33 +106,39 @@ def test_safe_gate_chunk( lower_bound = None naive_kda_gate_fn = naive_kda_gate - beta = torch.randn(B, T, H, dtype=torch.float32).sigmoid().to(beta_dtype) - h0 = torch.randn(B, H, D, D, dtype=torch.float32) + beta = torch.randn(B, T, HV, dtype=torch.float32).sigmoid().to(beta_dtype) + h0 = torch.randn(B, HV, D, D, dtype=torch.float32) # NOTE: for inference scenarios, we only use transposed state layout for better decoding performance h0_vk = h0.transpose(-1, -2).contiguous() if use_gate_in_kernel: A_log, dt_bias = map(lambda x: x.to(device).requires_grad_(False), (A_log, dt_bias)) q, k, v, g, beta, h0, h0_vk = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, g, beta, h0, h0_vk)) + initial_state = h0.clone() if use_initial_state else None + initial_state_vk = h0_vk.clone() if use_initial_state else None + + heads_per_group = HV // H + q_ref = q.repeat_interleave(heads_per_group, dim=2) + k_ref = k.repeat_interleave(heads_per_group, dim=2) ref, ref_ht = naive_recurrent_kda( - q=F.normalize(q.clone(), p=2, dim=-1), - k=F.normalize(k.clone(), p=2, dim=-1), + q=F.normalize(q_ref.clone(), p=2, dim=-1), + k=F.normalize(k_ref.clone(), p=2, dim=-1), v=v.clone(), g=(naive_kda_gate_fn(g, A_log, dt_bias) if use_gate_in_kernel else g.clone()), beta=beta.clone(), - initial_state=h0.clone(), + initial_state=initial_state, output_final_state=True, ) ref_fla, ref_ht_fla = fla_chunk_kda( - q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), - k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + q=F.normalize(q_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q_ref.clone(), + k=F.normalize(k_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k_ref.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), A_log=(A_log.clone() if use_gate_in_kernel else None), dt_bias=(dt_bias.clone() if use_gate_in_kernel else None), - initial_state=h0.clone(), + initial_state=initial_state, output_final_state=True, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_gate_in_kernel=use_gate_in_kernel, @@ -133,14 +147,14 @@ def test_safe_gate_chunk( ) ref_fla_trans, ref_ht_fla_trans = fla_chunk_kda( - q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), - k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + q=F.normalize(q_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q_ref.clone(), + k=F.normalize(k_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k_ref.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), A_log=(A_log.clone() if use_gate_in_kernel else None), dt_bias=(dt_bias.clone() if use_gate_in_kernel else None), - initial_state=h0_vk.clone(), + initial_state=initial_state_vk, output_final_state=True, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_gate_in_kernel=use_gate_in_kernel, @@ -157,7 +171,7 @@ def test_safe_gate_chunk( beta=beta.clone(), A_log=(A_log.clone() if use_gate_in_kernel else None), dt_bias=(dt_bias.clone() if use_gate_in_kernel else None), - initial_state=h0_vk.clone(), + initial_state=initial_state_vk, output_final_state=True, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_gate_in_kernel=use_gate_in_kernel, @@ -166,10 +180,10 @@ def test_safe_gate_chunk( ) assert_close("o", ref, tri, 0.005) - assert_close("ht", ref_ht, tri_ht.transpose(-1, -2), 0.005) assert_close("o", ref_fla, tri, 0.005) - assert_close("ht", ref_ht_fla, tri_ht.transpose(-1, -2), 0.005) assert_close("o", ref_fla_trans, tri, 0.005) + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2), 0.005) + assert_close("ht", ref_ht_fla, tri_ht.transpose(-1, -2), 0.005) assert_close("ht", ref_ht_fla_trans, tri_ht, 0.005) @@ -222,58 +236,75 @@ def test_safe_gate_chunk_no_final_state(): @pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) @pytest.mark.parametrize( - ("H", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), + ("H", "HV", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate", "use_initial_state"), [ - pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) + pytest.param( + *test, + id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}-init{}".format(*test), + ) for test in [ - (4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), - (4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), + (4, 4, 128, 0.1, [0, 15], torch.bfloat16, True, True), + (4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True, True), + (4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True, True), + (4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True, True), + (4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True, True), + (2, 4, 128, 0, [0, 63, 130], torch.bfloat16, True, True), + (1, 2, 128, 0, [0, 1], torch.bfloat16, True, True), + (1, 2, 128, 0, [0, 63, 64, 65], torch.bfloat16, True, True), + (2, 4, 128, 0, [0, 17, 64, 65, 130], torch.bfloat16, True, False), # ======Varlen test with simulated trace======= ( + 32, 32, 128, 0, [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096, 4824, 5101, 5919, 6426, 7137, 7392, 7800, 8192], torch.bfloat16, True, + True, ), ( + 32, 32, 128, 0, [0, 652, 1255, 1600, 2083, 2345, 2756, 3172, 3767, 4096, 4891, 5236, 5543, 6255, 6480, 6947, 7616, 8192], torch.bfloat16, True, + True, ), ( + 32, 32, 128, 0, [0, 315, 973, 1283, 2162, 2459, 2678, 2998, 3781, 4096, 4503, 5459, 6318, 6669, 6979, 7583, 8192], torch.bfloat16, True, + True, ), ( + 32, 32, 128, 0, [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], torch.bfloat16, True, + True, ), ] ], ) def test_safe_gate_chunk_varlen( H: int, + HV: int, D: int, mask_p: float, cu_seqlens: list[int], dtype: torch.dtype, safe_gate: bool, + use_initial_state: bool, beta_dtype: torch.dtype, ): cula_kda_fused_fwd = get_kda_fused_fwd(device) @@ -286,19 +317,24 @@ def test_safe_gate_chunk_varlen( q = torch.randn((1, T, H, D), dtype=dtype) k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) - v = torch.randn((1, T, H, D), dtype=dtype) - g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)) + v = torch.randn((1, T, HV, D), dtype=dtype) + g = F.logsigmoid(torch.randn(1, T, HV, D, dtype=torch.float)) mask = torch.rand_like(g) > mask_p g = g * mask + (~mask) * (-1000) if safe_gate: g = g.clamp(-5, 0) - beta = torch.randn(1, T, H, dtype=torch.float32).sigmoid().to(beta_dtype) - h0 = torch.randn((N, H, D, D), dtype=torch.float32) + beta = torch.randn(1, T, HV, dtype=torch.float32).sigmoid().to(beta_dtype) + h0 = torch.randn((N, HV, D, D), dtype=torch.float32) # NOTE: for inference scenarios, we only use transposed state layout for better decoding performance h0_vk = h0.transpose(-1, -2).contiguous() q, k, v, g, beta, h0, h0_vk = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, g, beta, h0, h0_vk)) + initial_state = h0.clone() if use_initial_state else None + initial_state_vk = h0_vk.clone() if use_initial_state else None + heads_per_group = HV // H + q_ref = q.repeat_interleave(heads_per_group, dim=2) + k_ref = k.repeat_interleave(heads_per_group, dim=2) tri, tri_ht = cula_kda_fused_fwd( q=F.normalize(q.clone(), p=2, dim=-1), @@ -306,7 +342,7 @@ def test_safe_gate_chunk_varlen( v=v.clone(), g=g.clone(), beta=beta.clone(), - initial_state=h0_vk.clone(), + initial_state=initial_state_vk, output_final_state=True, cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, @@ -315,12 +351,12 @@ def test_safe_gate_chunk_varlen( ) ref_fla, ref_ht_fla = fla_chunk_kda( - q=F.normalize(q.clone(), p=2, dim=-1), - k=k.clone(), + q=F.normalize(q_ref.clone(), p=2, dim=-1), + k=k_ref.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), - initial_state=h0.clone(), + initial_state=initial_state, output_final_state=True, cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, @@ -329,12 +365,12 @@ def test_safe_gate_chunk_varlen( ) ref_fla_trans, ref_ht_fla_trans = fla_chunk_kda( - q=F.normalize(q.clone(), p=2, dim=-1), - k=k.clone(), + q=F.normalize(q_ref.clone(), p=2, dim=-1), + k=k_ref.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), - initial_state=h0_vk.clone(), + initial_state=initial_state_vk, output_final_state=True, cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, @@ -347,12 +383,12 @@ def test_safe_gate_chunk_varlen( ref_ht = [] for i in range(N): ref_i, ref_ht_i = naive_recurrent_kda( - q=F.normalize(q[:, cu_seqlens[i] : cu_seqlens[i + 1]], p=2, dim=-1), - k=k[:, cu_seqlens[i] : cu_seqlens[i + 1]], + q=F.normalize(q_ref[:, cu_seqlens[i] : cu_seqlens[i + 1]], p=2, dim=-1), + k=k_ref[:, cu_seqlens[i] : cu_seqlens[i + 1]], v=v[:, cu_seqlens[i] : cu_seqlens[i + 1]], beta=beta[:, cu_seqlens[i] : cu_seqlens[i + 1]], g=g[:, cu_seqlens[i] : cu_seqlens[i + 1]], - initial_state=h0[i], + initial_state=h0[i] if use_initial_state else None, output_final_state=True, ) ref.append(ref_i) @@ -361,8 +397,8 @@ def test_safe_gate_chunk_varlen( ref_ht = torch.cat(ref_ht, 0) assert_close("o", ref, tri, 0.005) - assert_close("ht", ref_ht, tri_ht.transpose(-1, -2), 0.005) assert_close("o", ref_fla, tri, 0.005) - assert_close("ht", ref_ht_fla, tri_ht.transpose(-1, -2), 0.005) assert_close("o", ref_fla_trans, tri, 0.005) + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2), 0.005) + assert_close("ht", ref_ht_fla, tri_ht.transpose(-1, -2), 0.005) assert_close("ht", ref_ht_fla_trans, tri_ht, 0.005) From e4ea536706bab97a5ef343ac2a1d66ba26ae2d13 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Fri, 22 May 2026 10:33:10 +0800 Subject: [PATCH 13/34] [Feat] upgrade FLA to v0.5.0 (#72) * upgrade fla and update b200 bench, update readme and fix lightning test param * update h200 bench result with fla bug fixed * update b200 bench * update b200 bench * fix readme * remove useless repeat_interleave for fla * fix readme --------- Co-authored-by: boyu.zbw --- BENCHMARK_GB200_CUDA_130.md | 184 ++++++++++----------- BENCHMARK_H200.md | 58 +++---- README.md | 18 +- benchmarks/bench_kda_fused_fwd.py | 10 +- benchmarks/generate_benchmark_hopper_md.py | 2 +- benchmarks/generate_benchmark_md.py | 2 +- benchmarks/utils.py | 6 - tests/test_kda_fused_fwd.py | 16 +- tests/test_lightning_attn.py | 3 +- third_party/flash-linear-attention | 2 +- 10 files changed, 150 insertions(+), 151 deletions(-) diff --git a/BENCHMARK_GB200_CUDA_130.md b/BENCHMARK_GB200_CUDA_130.md index b14d8f9b..38bb9bd6 100644 --- a/BENCHMARK_GB200_CUDA_130.md +++ b/BENCHMARK_GB200_CUDA_130.md @@ -1,10 +1,10 @@ # Benchmark Results -> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-05-12. +> Auto-generated by `benchmarks/generate_benchmark_md.py` on 2026-05-19. > **GPU:** NVIDIA GB200 | **CUDA:** 13.0 | **PyTorch:** 2.9.1+cu130 -> FLA baseline: [flash-linear-attention v0.4.2](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.4.2) +> FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0) @@ -14,44 +14,44 @@ | B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------|---------| -| 1 | 512 | 0.582 | 0.483 | **1.21x** | -| 1 | 1024 | 0.579 | 0.493 | **1.17x** | -| 1 | 4096 | 0.749 | 0.541 | **1.38x** | -| 1 | 8192 | 1.393 | 1.009 | **1.38x** | -| 1 | 16384 | 2.706 | 1.931 | **1.40x** | -| 2 | 512 | 0.595 | 0.510 | **1.17x** | -| 2 | 1024 | 0.619 | 0.498 | **1.24x** | -| 2 | 4096 | 1.394 | 1.016 | **1.37x** | -| 2 | 8192 | 2.701 | 1.949 | **1.39x** | -| 2 | 16384 | 5.297 | 3.875 | **1.37x** | +| 1 | 512 | 0.838 | 0.604 | **1.39x** | +| 1 | 1024 | 0.694 | 0.571 | **1.22x** | +| 1 | 4096 | 0.759 | 0.564 | **1.35x** | +| 1 | 8192 | 1.406 | 1.026 | **1.37x** | +| 1 | 16384 | 2.734 | 1.965 | **1.39x** | +| 2 | 512 | 0.665 | 0.555 | **1.20x** | +| 2 | 1024 | 0.695 | 0.562 | **1.24x** | +| 2 | 4096 | 1.408 | 1.034 | **1.36x** | +| 2 | 8192 | 2.733 | 1.978 | **1.38x** | +| 2 | 16384 | 5.354 | 3.877 | **1.38x** | -Summary (10 configs): **avg=1.31x**, min=1.17x, max=1.40x. +Summary (10 configs): **avg=1.33x**, min=1.20x, max=1.39x. ### Variable-Length (H=64, D=128, bf16) | Config | FLA Triton (ms) | cuLA (ms) | Speedup | |--------|-----------------|-----------|---------| -| uniform 10seqs T=4096 [409..415] avg=409 | 0.783 | 0.585 | **1.34x** | -| random 10seqs T=4096 [24..1201] avg=409 | 0.777 | 0.579 | **1.34x** | -| skewed 10seqs T=4096 [227..2053] avg=409 | 0.776 | 0.578 | **1.34x** | -| uniform 20seqs T=4096 [204..220] avg=204 | 0.855 | 0.633 | **1.35x** | -| random 20seqs T=4096 [5..787] avg=204 | 0.828 | 0.619 | **1.34x** | -| skewed 20seqs T=4096 [107..2063] avg=204 | 0.811 | 0.597 | **1.36x** | -| uniform 10seqs T=8192 [819..821] avg=819 | 1.386 | 1.028 | **1.35x** | -| random 10seqs T=8192 [48..2401] avg=819 | 1.414 | 1.048 | **1.35x** | -| skewed 10seqs T=8192 [455..4097] avg=819 | 1.441 | 1.049 | **1.37x** | -| uniform 20seqs T=8192 [409..421] avg=409 | 1.476 | 1.074 | **1.37x** | -| random 20seqs T=8192 [9..1574] avg=409 | 1.475 | 1.079 | **1.37x** | -| skewed 20seqs T=8192 [215..4107] avg=409 | 1.482 | 1.081 | **1.37x** | -| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.671 | 1.963 | **1.36x** | -| random 10seqs T=16384 [95..4802] avg=1638 | 2.684 | 1.965 | **1.37x** | -| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.688 | 1.972 | **1.36x** | -| uniform 20seqs T=16384 [819..823] avg=819 | 2.680 | 1.966 | **1.36x** | -| random 20seqs T=16384 [19..3147] avg=819 | 2.712 | 1.990 | **1.36x** | -| skewed 20seqs T=16384 [431..8195] avg=819 | 2.691 | 1.970 | **1.37x** | - -Summary (18 configs): **avg=1.36x**, min=1.34x, max=1.37x. +| uniform 10seqs T=4096 [409..415] avg=409 | 0.796 | 0.600 | **1.33x** | +| random 10seqs T=4096 [24..1201] avg=409 | 0.789 | 0.587 | **1.34x** | +| skewed 10seqs T=4096 [227..2053] avg=409 | 0.790 | 0.590 | **1.34x** | +| uniform 20seqs T=4096 [204..220] avg=204 | 0.871 | 0.649 | **1.34x** | +| random 20seqs T=4096 [5..787] avg=204 | 0.843 | 0.634 | **1.33x** | +| skewed 20seqs T=4096 [107..2063] avg=204 | 0.822 | 0.608 | **1.35x** | +| uniform 10seqs T=8192 [819..821] avg=819 | 1.405 | 1.045 | **1.34x** | +| random 10seqs T=8192 [48..2401] avg=819 | 1.433 | 1.070 | **1.34x** | +| skewed 10seqs T=8192 [455..4097] avg=819 | 1.458 | 1.068 | **1.37x** | +| uniform 20seqs T=8192 [409..421] avg=409 | 1.494 | 1.095 | **1.36x** | +| random 20seqs T=8192 [9..1574] avg=409 | 1.494 | 1.097 | **1.36x** | +| skewed 20seqs T=8192 [215..4107] avg=409 | 1.499 | 1.101 | **1.36x** | +| uniform 10seqs T=16384 [1638..1642] avg=1638 | 2.696 | 1.988 | **1.36x** | +| random 10seqs T=16384 [95..4802] avg=1638 | 2.704 | 1.990 | **1.36x** | +| skewed 10seqs T=16384 [910..8194] avg=1638 | 2.715 | 2.000 | **1.36x** | +| uniform 20seqs T=16384 [819..823] avg=819 | 2.718 | 1.998 | **1.36x** | +| random 20seqs T=16384 [19..3147] avg=819 | 2.742 | 2.023 | **1.36x** | +| skewed 20seqs T=16384 [431..8195] avg=819 | 2.723 | 2.001 | **1.36x** | + +Summary (18 configs): **avg=1.35x**, min=1.33x, max=1.37x. To reproduce: @@ -66,14 +66,14 @@ python benchmarks/bench_kda.py --mode both | B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------|---------| -| 1 | 1024 | 0.087 | 0.070 | **1.24x** | +| 1 | 1024 | 0.112 | 0.073 | **1.53x** | | 1 | 4096 | 0.175 | 0.157 | **1.11x** | -| 1 | 8192 | 0.330 | 0.292 | **1.13x** | -| 1 | 16384 | 0.628 | 0.563 | **1.12x** | -| 2 | 1024 | 0.099 | 0.064 | **1.53x** | -| 2 | 4096 | 0.327 | 0.175 | **1.87x** | +| 1 | 8192 | 0.329 | 0.292 | **1.13x** | +| 1 | 16384 | 0.629 | 0.563 | **1.12x** | +| 2 | 1024 | 0.099 | 0.068 | **1.45x** | +| 2 | 4096 | 0.327 | 0.176 | **1.86x** | | 2 | 8192 | 0.631 | 0.327 | **1.93x** | -| 2 | 16384 | 1.249 | 0.632 | **1.98x** | +| 2 | 16384 | 1.257 | 0.632 | **1.99x** | ### Variable-Length (H=64, D=128, bf16) @@ -81,50 +81,50 @@ Persistent CuTe DSL kernel vs FLA Triton varlen. | N (seqs) | T | cuLA (ms) | FLA Triton (ms) | Speedup | |----------|---|-----------|-----------------|---------| -| 5 | 1020 | 0.089 | 0.171 | **1.91x** | -| 5 | 2045 | 0.111 | 0.189 | **1.71x** | -| 5 | 4095 | 0.163 | 0.249 | **1.53x** | -| 5 | 8190 | 0.264 | 0.399 | **1.51x** | -| 5 | 16380 | 0.463 | 0.702 | **1.52x** | -| 5 | 32765 | 0.858 | 1.283 | **1.49x** | -| 8 | 1024 | 0.086 | 0.156 | **1.82x** | -| 8 | 2048 | 0.111 | 0.183 | **1.65x** | -| 8 | 4096 | 0.157 | 0.250 | **1.59x** | -| 8 | 8192 | 0.243 | 0.402 | **1.66x** | -| 8 | 16384 | 0.413 | 0.688 | **1.67x** | -| 8 | 32768 | 0.756 | 1.252 | **1.66x** | -| 10 | 1020 | 0.104 | 0.162 | **1.56x** | -| 10 | 2040 | 0.133 | 0.200 | **1.51x** | -| 10 | 4090 | 0.179 | 0.269 | **1.50x** | -| 10 | 8190 | 0.267 | 0.414 | **1.55x** | -| 10 | 16380 | 0.439 | 0.693 | **1.58x** | -| 10 | 32760 | 0.788 | 1.260 | **1.60x** | -| 12 | 1020 | 0.119 | 0.175 | **1.47x** | -| 12 | 2040 | 0.143 | 0.197 | **1.38x** | -| 12 | 4092 | 0.189 | 0.265 | **1.40x** | -| 12 | 8184 | 0.281 | 0.405 | **1.44x** | -| 12 | 16380 | 0.452 | 0.703 | **1.55x** | -| 12 | 32760 | 0.793 | 1.259 | **1.59x** | -| 16 | 1024 | 0.121 | 0.157 | **1.30x** | -| 16 | 2048 | 0.149 | 0.183 | **1.23x** | -| 16 | 4096 | 0.187 | 0.256 | **1.37x** | +| 5 | 1020 | 0.095 | 0.199 | **2.08x** | +| 5 | 2045 | 0.112 | 0.219 | **1.96x** | +| 5 | 4095 | 0.164 | 0.262 | **1.60x** | +| 5 | 8190 | 0.266 | 0.410 | **1.54x** | +| 5 | 16380 | 0.464 | 0.698 | **1.50x** | +| 5 | 32765 | 0.860 | 1.289 | **1.50x** | +| 8 | 1024 | 0.096 | 0.165 | **1.72x** | +| 8 | 2048 | 0.111 | 0.197 | **1.78x** | +| 8 | 4096 | 0.157 | 0.248 | **1.58x** | +| 8 | 8192 | 0.241 | 0.389 | **1.61x** | +| 8 | 16384 | 0.412 | 0.680 | **1.65x** | +| 8 | 32768 | 0.757 | 1.250 | **1.65x** | +| 10 | 1020 | 0.105 | 0.159 | **1.52x** | +| 10 | 2040 | 0.133 | 0.199 | **1.50x** | +| 10 | 4090 | 0.180 | 0.261 | **1.45x** | +| 10 | 8190 | 0.266 | 0.403 | **1.51x** | +| 10 | 16380 | 0.440 | 0.688 | **1.56x** | +| 10 | 32760 | 0.789 | 1.264 | **1.60x** | +| 12 | 1020 | 0.118 | 0.164 | **1.39x** | +| 12 | 2040 | 0.142 | 0.190 | **1.35x** | +| 12 | 4092 | 0.189 | 0.260 | **1.37x** | +| 12 | 8184 | 0.280 | 0.401 | **1.43x** | +| 12 | 16380 | 0.454 | 0.697 | **1.54x** | +| 12 | 32760 | 0.795 | 1.250 | **1.57x** | +| 16 | 1024 | 0.121 | 0.162 | **1.35x** | +| 16 | 2048 | 0.149 | 0.186 | **1.24x** | +| 16 | 4096 | 0.188 | 0.254 | **1.35x** | | 16 | 8192 | 0.267 | 0.398 | **1.49x** | -| 16 | 16384 | 0.424 | 0.686 | **1.62x** | -| 16 | 32768 | 0.740 | 1.247 | **1.68x** | -| 20 | 1020 | 0.162 | 0.174 | **1.07x** | -| 20 | 2040 | 0.191 | 0.207 | **1.08x** | -| 20 | 4080 | 0.233 | 0.288 | **1.24x** | -| 20 | 8180 | 0.319 | 0.424 | **1.33x** | -| 20 | 16380 | 0.478 | 0.703 | **1.47x** | -| 20 | 32760 | 0.800 | 1.261 | **1.58x** | -| 25 | 1000 | 0.193 | 0.176 | **0.91x** | -| 25 | 2025 | 0.221 | 0.227 | **1.03x** | -| 25 | 4075 | 0.258 | 0.286 | **1.11x** | -| 25 | 8175 | 0.347 | 0.445 | **1.28x** | -| 25 | 16375 | 0.517 | 0.720 | **1.39x** | -| 25 | 32750 | 0.831 | 1.270 | **1.53x** | - -Summary (126 configs across uniform/skewed/random): **avg=1.48x**, min=0.91x, max=2.01x. +| 16 | 16384 | 0.424 | 0.688 | **1.62x** | +| 16 | 32768 | 0.742 | 1.242 | **1.67x** | +| 20 | 1020 | 0.162 | 0.173 | **1.07x** | +| 20 | 2040 | 0.191 | 0.203 | **1.06x** | +| 20 | 4080 | 0.235 | 0.283 | **1.20x** | +| 20 | 8180 | 0.319 | 0.415 | **1.30x** | +| 20 | 16380 | 0.481 | 0.691 | **1.44x** | +| 20 | 32760 | 0.804 | 1.262 | **1.57x** | +| 25 | 1000 | 0.193 | 0.184 | **0.95x** | +| 25 | 2025 | 0.223 | 0.225 | **1.01x** | +| 25 | 4075 | 0.260 | 0.288 | **1.11x** | +| 25 | 8175 | 0.349 | 0.450 | **1.29x** | +| 25 | 16375 | 0.520 | 0.718 | **1.38x** | +| 25 | 32750 | 0.834 | 1.275 | **1.53x** | + +Summary (126 configs across uniform/skewed/random): **avg=1.47x**, min=0.92x, max=2.16x. To reproduce: @@ -140,21 +140,21 @@ Single-token decode: la_decode (CuTe DSL) vs fla fused_recurrent (Triton). | B | FLA Triton (ms) | cuLA (ms) | Speedup | |---|-----------------|-----------|---------| -| 1 | 0.0740 | 0.0134 | **5.53x** | -| 4 | 0.0698 | 0.0130 | **5.39x** | -| 16 | 0.0731 | 0.0209 | **3.50x** | -| 64 | 0.0996 | 0.0843 | **1.18x** | -| 256 | 0.3501 | 0.3126 | **1.12x** | +| 1 | 0.0728 | 0.0149 | **4.88x** | +| 4 | 0.0722 | 0.0147 | **4.92x** | +| 16 | 0.0763 | 0.0209 | **3.66x** | +| 64 | 0.0997 | 0.0843 | **1.18x** | +| 256 | 0.3494 | 0.3123 | **1.12x** | #### Wrapper (Full Call Path) | B | FLA Triton (ms) | cuLA (ms) | Speedup | |---|-----------------|-----------|---------| -| 1 | 0.0958 | 0.0189 | **5.08x** | -| 4 | 0.0920 | 0.0186 | **4.95x** | -| 16 | 0.0934 | 0.0211 | **4.43x** | -| 64 | 0.0990 | 0.0850 | **1.17x** | -| 256 | 0.3492 | 0.3133 | **1.11x** | +| 1 | 0.0953 | 0.0194 | **4.91x** | +| 4 | 0.0924 | 0.0193 | **4.80x** | +| 16 | 0.0977 | 0.0233 | **4.20x** | +| 64 | 0.1029 | 0.0846 | **1.22x** | +| 256 | 0.3490 | 0.3133 | **1.11x** | To reproduce: diff --git a/BENCHMARK_H200.md b/BENCHMARK_H200.md index 181e397d..ce0d46f7 100644 --- a/BENCHMARK_H200.md +++ b/BENCHMARK_H200.md @@ -1,10 +1,10 @@ # Benchmark Results — Hopper (SM90) -> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on 2026-04-05. +> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on 2026-05-19. > **GPU:** NVIDIA H200 | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 -> FLA baseline: [flash-linear-attention v0.4.2](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.4.2) +> FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0) @@ -16,39 +16,39 @@ Fully-fused KDA forward prefill kernel (sm90). | B | T | FLA Triton (ms) | cuLA Fused (ms) | Speedup | |---|---|-----------------|-----------------|---------| -| 1 | 512 | 0.576 | 0.230 | **2.51x** | -| 1 | 1024 | 0.572 | 0.248 | **2.31x** | -| 1 | 4096 | 0.936 | 0.899 | **1.04x** | -| 1 | 8192 | 1.819 | 1.758 | **1.03x** | -| 1 | 16384 | 3.599 | 3.521 | **1.02x** | -| 2 | 512 | 0.569 | 0.228 | **2.49x** | -| 2 | 1024 | 0.572 | 0.306 | **1.87x** | -| 2 | 4096 | 1.818 | 1.108 | **1.64x** | -| 2 | 8192 | 3.605 | 2.210 | **1.63x** | -| 2 | 16384 | 7.173 | 4.485 | **1.60x** | +| 1 | 512 | 0.556 | 0.224 | **2.48x** | +| 1 | 1024 | 0.581 | 0.248 | **2.34x** | +| 1 | 4096 | 0.936 | 0.896 | **1.04x** | +| 1 | 8192 | 1.810 | 1.754 | **1.03x** | +| 1 | 16384 | 3.576 | 3.492 | **1.02x** | +| 2 | 512 | 0.567 | 0.226 | **2.51x** | +| 2 | 1024 | 0.585 | 0.315 | **1.86x** | +| 2 | 4096 | 1.815 | 1.170 | **1.55x** | +| 2 | 8192 | 3.576 | 2.283 | **1.57x** | +| 2 | 16384 | 7.115 | 4.408 | **1.61x** | ### Variable-Length (H=64, D=128, bf16) | Config | FLA Triton (ms) | cuLA Fused (ms) | Speedup | |--------|-----------------|-----------------|---------| -| uniform 10seqs T=4096 [409..415] avg=409 | 1.016 | 0.707 | **1.44x** | -| random 10seqs T=4096 [24..1201] avg=409 | 1.008 | 0.660 | **1.53x** | -| skewed 10seqs T=4096 [227..2053] avg=409 | 1.005 | 0.668 | **1.50x** | -| uniform 20seqs T=4096 [204..220] avg=204 | 1.087 | 0.919 | **1.18x** | -| random 20seqs T=4096 [5..787] avg=204 | 1.066 | 0.736 | **1.45x** | -| skewed 20seqs T=4096 [107..2063] avg=204 | 1.038 | 0.724 | **1.43x** | -| uniform 10seqs T=8192 [819..821] avg=819 | 1.855 | 1.179 | **1.57x** | -| random 10seqs T=8192 [48..2401] avg=819 | 1.893 | 1.215 | **1.56x** | -| skewed 10seqs T=8192 [455..4097] avg=819 | 1.906 | 1.209 | **1.58x** | -| uniform 20seqs T=8192 [409..421] avg=409 | 1.961 | 1.406 | **1.39x** | -| random 20seqs T=8192 [9..1574] avg=409 | 1.954 | 1.283 | **1.52x** | +| uniform 10seqs T=4096 [409..415] avg=409 | 1.019 | 0.707 | **1.44x** | +| random 10seqs T=4096 [24..1201] avg=409 | 1.013 | 0.669 | **1.51x** | +| skewed 10seqs T=4096 [227..2053] avg=409 | 1.010 | 0.681 | **1.48x** | +| uniform 20seqs T=4096 [204..220] avg=204 | 1.098 | 0.932 | **1.18x** | +| random 20seqs T=4096 [5..787] avg=204 | 1.074 | 0.748 | **1.44x** | +| skewed 20seqs T=4096 [107..2063] avg=204 | 1.048 | 0.732 | **1.43x** | +| uniform 10seqs T=8192 [819..821] avg=819 | 1.851 | 1.174 | **1.58x** | +| random 10seqs T=8192 [48..2401] avg=819 | 1.890 | 1.217 | **1.55x** | +| skewed 10seqs T=8192 [455..4097] avg=819 | 1.905 | 1.225 | **1.55x** | +| uniform 20seqs T=8192 [409..421] avg=409 | 1.960 | 1.406 | **1.39x** | +| random 20seqs T=8192 [9..1574] avg=409 | 1.953 | 1.290 | **1.51x** | | skewed 20seqs T=8192 [215..4107] avg=409 | 1.957 | 1.300 | **1.51x** | -| uniform 10seqs T=16384 [1638..1642] avg=1638 | 3.646 | 2.188 | **1.67x** | -| random 10seqs T=16384 [95..4802] avg=1638 | 3.646 | 2.306 | **1.58x** | -| skewed 10seqs T=16384 [910..8194] avg=1638 | 3.656 | 2.335 | **1.57x** | -| uniform 20seqs T=16384 [819..823] avg=819 | 3.679 | 2.355 | **1.56x** | -| random 20seqs T=16384 [19..3147] avg=819 | 3.713 | 2.323 | **1.60x** | -| skewed 20seqs T=16384 [431..8195] avg=819 | 3.670 | 2.384 | **1.54x** | +| uniform 10seqs T=16384 [1638..1642] avg=1638 | 3.642 | 2.162 | **1.68x** | +| random 10seqs T=16384 [95..4802] avg=1638 | 3.609 | 2.279 | **1.58x** | +| skewed 10seqs T=16384 [910..8194] avg=1638 | 3.625 | 2.354 | **1.54x** | +| uniform 20seqs T=16384 [819..823] avg=819 | 3.644 | 2.320 | **1.57x** | +| random 20seqs T=16384 [19..3147] avg=819 | 3.681 | 2.293 | **1.61x** | +| skewed 20seqs T=16384 [431..8195] avg=819 | 3.634 | 2.371 | **1.53x** | Summary (28 configs): **avg=1.58x**, min=1.02x, max=2.51x. diff --git a/README.md b/README.md index d5764052..b894cd32 100644 --- a/README.md +++ b/README.md @@ -101,25 +101,23 @@ See [USAGE.md](USAGE.md) for detailed usage examples and notes. ## Benchmarks -Benchmarks run on a single **NVIDIA GB300/GB200/H200** GPU with **CUDA Toolkit 12.9**, **PyTorch 2.9.1**, **Triton 3.5.1**. +Benchmarks run on a single **NVIDIA GB200/H200** GPU with **PyTorch 2.9.1**, **Triton 3.5.1**. -FLA baseline: [flash-linear-attention v0.4.2](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.4.2). +FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0). **Blackwell (SM10X)** -See [BENCHMARK_GB300.md](BENCHMARK_GB300.md) for detailed results. - -See [BENCHMARK_GB200.md](BENCHMARK_GB200.md) for detailed results. +See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA 13.0 for detailed results. **Hopper (SM90)** -See [BENCHMARK_H200.md](BENCHMARK_H200.md) for detailed results. +See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 12.9 for detailed results. **Highlights:** -- **KDA Modular Forward (Blackwell):** **avg 1.45x** speedup on fixed-length, **avg 1.32x** on variable-length (18 configs, uniform/skewed/random). -- **Lightning Attention Prefill (Blackwell):** up to **1.86x** speedup (B=2). -- **Lightning Attention Varlen (Blackwell):** **avg 1.54x** speedup across 126 configs (uniform/skewed/random). -- **KDA Fused Forward (Hopper):** **avg 1.52x** speedup across fixed-length and variable-length sequences. +- **KDA Modular Forward (Blackwell):** **avg 1.33x** speedup on fixed-length, **avg 1.35x** on variable-length (18 configs, uniform/skewed/random). +- **Lightning Attention Prefill (Blackwell):** up to **2.08x** speedup (B=2). +- **Lightning Attention Varlen (Blackwell):** **avg 1.47x** speedup across 126 configs (uniform/skewed/random). +- **KDA Fused Forward (Hopper):** **avg 1.58x** speedup across fixed-length and variable-length sequences. To regenerate benchmarks: diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index 171c2bb9..0b2dd538 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -34,7 +34,7 @@ GVA (Grouped Value Attention) mode. HV must be a positive multiple of H. Usage: - python bench_kda_fused_fwd.py [--mode fixed|varlen|both] [--hv HV] [--ncu] + python bench_kda_fused_fwd.py [--mode fixed|varlen|both] [--heads H] [--hv HV] [--ncu] With --ncu, warmup=1 and iters=1 for ncu profiling: ncu --set full -o report python bench_kda_fused_fwd.py --mode varlen --ncu @@ -406,6 +406,13 @@ def main(): action="store_true", help="Use non-zero initial state (default: False)", ) + global H + parser.add_argument( + "--heads", + type=int, + default=H, + help=f"Number of Q/K heads (H). Default: {H}", + ) parser.add_argument( "--hv", type=int, @@ -415,6 +422,7 @@ def main(): args = parser.parse_args() global NCU_MODE, SANITIZER_MODE, HAS_INIT_STATE, HV + H = args.heads if args.ncu: NCU_MODE = True print("[NCU mode] warmup=1, iters=1") diff --git a/benchmarks/generate_benchmark_hopper_md.py b/benchmarks/generate_benchmark_hopper_md.py index 177ac241..13ffd1d8 100644 --- a/benchmarks/generate_benchmark_hopper_md.py +++ b/benchmarks/generate_benchmark_hopper_md.py @@ -82,7 +82,7 @@ def format_benchmark_md(env, kda_fused_fixed, kda_fused_varlen, has_init_state: w(f"> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on {datetime.now().strftime('%Y-%m-%d')}.\n") w(f"> **GPU:** {env['gpu']} | **CUDA:** {env['cuda']} | **PyTorch:** {env['torch']}\n") w( - "> FLA baseline: [flash-linear-attention v0.4.2](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.4.2)\n" + "> FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0)\n" ) w("") diff --git a/benchmarks/generate_benchmark_md.py b/benchmarks/generate_benchmark_md.py index 1d047557..96a09091 100644 --- a/benchmarks/generate_benchmark_md.py +++ b/benchmarks/generate_benchmark_md.py @@ -150,7 +150,7 @@ def format_benchmark_md(env, kda_fixed, kda_varlen, la_standard, la_varlen, la_d w(f"> Auto-generated by `benchmarks/generate_benchmark_md.py` on {datetime.now().strftime('%Y-%m-%d')}.\n") w(f"> **GPU:** {env['gpu']} | **CUDA:** {env['cuda']} | **PyTorch:** {env['torch']}\n") w( - "> FLA baseline: [flash-linear-attention v0.4.2](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.4.2)\n" + "> FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0)\n" ) w("") diff --git a/benchmarks/utils.py b/benchmarks/utils.py index bfd0761e..75d7ef52 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -286,12 +286,6 @@ def prepare_safe_gate_inputs( g = torch.randn(batch_size, T, HV, D, dtype=dtype, device=device).requires_grad_(False) beta = torch.randn(batch_size, T, HV, dtype=torch.float, device=device).sigmoid().requires_grad_(False) - # GVA expansion: bring q/k up to HV heads so all tensors share head dim. - group = HV // H - if group > 1: - q = q.repeat_interleave(group, dim=2).contiguous() - k = k.repeat_interleave(group, dim=2).contiguous() - # A_log / dt_bias must match the head count of `g` (HV), otherwise # kda_gate_chunk_cumsum would index out of bounds for i_h >= H. A_log = torch.randn(HV, dtype=torch.float, device=device).requires_grad_(False) diff --git a/tests/test_kda_fused_fwd.py b/tests/test_kda_fused_fwd.py index 354f0c9c..05121327 100644 --- a/tests/test_kda_fused_fwd.py +++ b/tests/test_kda_fused_fwd.py @@ -131,8 +131,8 @@ def test_safe_gate_chunk( ) ref_fla, ref_ht_fla = fla_chunk_kda( - q=F.normalize(q_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q_ref.clone(), - k=F.normalize(k_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k_ref.clone(), + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), @@ -147,8 +147,8 @@ def test_safe_gate_chunk( ) ref_fla_trans, ref_ht_fla_trans = fla_chunk_kda( - q=F.normalize(q_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q_ref.clone(), - k=F.normalize(k_ref.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k_ref.clone(), + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), @@ -351,8 +351,8 @@ def test_safe_gate_chunk_varlen( ) ref_fla, ref_ht_fla = fla_chunk_kda( - q=F.normalize(q_ref.clone(), p=2, dim=-1), - k=k_ref.clone(), + q=F.normalize(q.clone(), p=2, dim=-1), + k=k.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), @@ -365,8 +365,8 @@ def test_safe_gate_chunk_varlen( ) ref_fla_trans, ref_ht_fla_trans = fla_chunk_kda( - q=F.normalize(q_ref.clone(), p=2, dim=-1), - k=k_ref.clone(), + q=F.normalize(q.clone(), p=2, dim=-1), + k=k.clone(), v=v.clone(), g=g.clone(), beta=beta.clone(), diff --git a/tests/test_lightning_attn.py b/tests/test_lightning_attn.py index 8958f54c..26fcc16e 100644 --- a/tests/test_lightning_attn.py +++ b/tests/test_lightning_attn.py @@ -363,7 +363,7 @@ def test_against_fla(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rto g_gamma = -decay # FLA reference (scale=1.0 to match our kernel) - O_fla, _ = chunk_simple_gla(Q, K, V, g_gamma=g_gamma, scale=1.0, head_first=False) + O_fla, _ = chunk_simple_gla(Q, K, V, g_gamma=g_gamma, scale=1.0) # Our kernel O_cute, _ = run_cute_kernel(Q, K, V, decay, scale=1.0, chunk_size=C) @@ -411,7 +411,6 @@ def test_against_fla_with_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, ato scale=1.0, initial_state=h0.clone(), output_final_state=True, - head_first=False, ) # Ours (expects BHVK state) diff --git a/third_party/flash-linear-attention b/third_party/flash-linear-attention index ca910f88..3a9ce1c8 160000 --- a/third_party/flash-linear-attention +++ b/third_party/flash-linear-attention @@ -1 +1 @@ -Subproject commit ca910f88529565b28b6e16465258f2e239a02dc7 +Subproject commit 3a9ce1c83a13994d824dbb3421e2989d330bb38b From f17c36d505e0c9e073d5bdb5295587118160b2df Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Fri, 22 May 2026 14:28:10 +0800 Subject: [PATCH 14/34] [Fix] add cross-proxy fence for recomp_wu kernel (#77) * fix q pipeline race condition * add fence * remove redundant tma->cuda core fence * update e2e bench * add fence for cuda core->tma * fix fence postition --- benchmarks/bench_kda_fwd_bwd_e2e.py | 14 ++--- benchmarks/bench_recompute_wu.py | 42 +++++++++++++ .../kda_fwd_recomp_w_u_mainloop_sm100.hpp | 59 +++++++++++-------- 3 files changed, 82 insertions(+), 33 deletions(-) diff --git a/benchmarks/bench_kda_fwd_bwd_e2e.py b/benchmarks/bench_kda_fwd_bwd_e2e.py index c6b4117b..24c67213 100644 --- a/benchmarks/bench_kda_fwd_bwd_e2e.py +++ b/benchmarks/bench_kda_fwd_bwd_e2e.py @@ -222,6 +222,8 @@ def check_determinism(num_seqs=5, T=512, iters=20): for i in range(iters): out = run_kda_e2e_with_grads(**common, fn=cula_chunk_kda) for name in ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"): + assert torch.isnan(out[name]).sum() == 0, f"[determinism] cuLA {name} has NaNs at iter {i}" + assert torch.isfinite(out[name]).all(), f"[determinism] cuLA {name} has infs at iter {i}" assert torch.equal(out[name], ref[name]), f"[determinism] cuLA {name} mismatch at iter {i}" return True @@ -559,11 +561,6 @@ def main(): action="store_true", help="Disable recompute in both FLA and cuLA (pre-compute QG)", ) - parser.add_argument( - "--check_determinism", - action="store_true", - help="Run determinism check: verify cuLA produces identical outputs across repeated runs", - ) args = parser.parse_args() global NCU_MODE, SANITIZER_MODE, DISABLE_RECOMPUTE, PHASE @@ -578,14 +575,13 @@ def main(): print("[Disable recompute] pre-compute QG in forward") PHASE = args.phase - if args.check_determinism: + if not (args.ncu or args.sanitizer): det_configs = [(5, 1024), (10, 4096), (10, 8192), (10, 16384)] print("\n[Determinism Check] cuLA chunk_kda E2E ...") for num_seqs, T in det_configs: - result = check_determinism(num_seqs=num_seqs, T=T, iters=20) - print(f" num_seqs={num_seqs} T={T:5d} iters=20 {'PASS' if result else 'FAIL'}") + result = check_determinism(num_seqs=num_seqs, T=T, iters=1000) + print(f" num_seqs={num_seqs} T={T:5d} {'PASS' if result else 'FAIL'}") print("[Determinism Check] All passed.\n") - return fixed_configs = [ # (B, T) diff --git a/benchmarks/bench_recompute_wu.py b/benchmarks/bench_recompute_wu.py index c7d0873a..0beca3b5 100644 --- a/benchmarks/bench_recompute_wu.py +++ b/benchmarks/bench_recompute_wu.py @@ -25,6 +25,7 @@ from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra from fla.ops.kda.wy_fast import recompute_w_u_fwd as fla_recompute_w_u_fwd +from fla.utils import get_abs_err, get_err_ratio import cula.cudac as cula_cuda from benchmarks.utils import SEED, exclusive_cumsum, generate_random_seq_lens, prepare_intra_inputs @@ -241,6 +242,45 @@ def benchmark_recompute_wu_varlen(): print("─" * 100) +def check_determinism(num_seqs=NUM_SEQS, T=2001, H=H, iters=1000): + """Run the recompute_w_u kernel multiple times and check for deterministic outputs.""" + device = torch.device("cuda") + chunk_size = BT + + seq_lens = generate_random_seq_lens(num_seqs, T, MIN_SEQ_LEN, VARIANCE, SEED) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + + q, k, v, g, beta, Akk, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs( + B=1, T=T, H=H, D=D, device=device, cu_seqlens=cu_seqlens, chunk_size=chunk_size + ) + + ref_w, ref_u, ref_qg, ref_kg = run_cula_recompute_wu( + k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, chunk_size, disable_recompute=DISABLE_RECOMPUTE + ) + + for i in range(iters): + w, u, qg, kg = run_cula_recompute_wu( + k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, chunk_size, disable_recompute=DISABLE_RECOMPUTE + ) + + if not torch.equal(w, ref_w): + print(f"Iteration {i}: w mismatch") + print(f"{get_abs_err(ref_w, w):.6f} absolute error, {get_err_ratio(ref_w, w):.6f} relative error") + raise AssertionError("Non-deterministic output detected in w") + if not torch.equal(u, ref_u): + print(f"Iteration {i}: u mismatch") + print(f"{get_abs_err(ref_u, u):.6f} absolute error, {get_err_ratio(ref_u, u):.6f} relative error") + raise AssertionError("Non-deterministic output detected in u") + if kg is not None and not torch.equal(kg, ref_kg): + print(f"Iteration {i}: kg mismatch") + print(f"{get_abs_err(ref_kg, kg):.6f} absolute error, {get_err_ratio(ref_kg, kg):.6f} relative error") + raise AssertionError("Non-deterministic output detected in kg") + if qg is not None and not torch.equal(qg, ref_qg): + print(f"Iteration {i}: qg mismatch") + print(f"{get_abs_err(ref_qg, qg):.6f} absolute error, {get_err_ratio(ref_qg, qg):.6f} relative error") + raise AssertionError("Non-deterministic output detected in qg") + + if __name__ == "__main__": parser = argparse.ArgumentParser(description="bench_recompute_wu: cuLA vs FLA Triton for recompute_w_u") parser.add_argument( @@ -254,5 +294,7 @@ def benchmark_recompute_wu_varlen(): DISABLE_RECOMPUTE = True print("[Disable recompute] pre-compute QG in forward") + check_determinism(iters=100000) + benchmark_recompute_wu_uniform() benchmark_recompute_wu_varlen() diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index 0f6a66f8..8b51b65b 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -318,9 +318,6 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } } - // Release K SMEM back to Load warp (done reading K) - k_pipeline.consumer_release(k_pipe_state_read); - ++k_pipe_state_read; g_pipeline.consumer_wait(g_pipe_state_read); Tensor sG = @@ -393,11 +390,14 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } } - fence_view_async_shared(); prologue_ready_pipeline.producer_commit(prologue_ready_pipe_state_write); ++prologue_ready_pipe_state_write; + // Release K SMEM back to Load warp (done reading K) + k_pipeline.consumer_release(k_pipe_state_read); + ++k_pipe_state_read; + // ---- Step 3: Compute KG = K * exp2(g_last - G) → write to out (K-major) → store to GMEM ---- // Load g_last from SMEM into registers (only need last valid row) float4 g_last_reg[TileK / 64][2]; @@ -408,9 +408,6 @@ struct KdaChunkFwdRecompWUMainloopSm100 { g_last_reg[k_yi][1] = *reinterpret_cast(&sG(sub_seq_len - 1, y + 4)); } - g_pipeline.consumer_release(g_pipe_state_read); - ++g_pipe_state_read; - // Need NamedBarrier to ensure all 128 prologue threads finish previous sKG_out writes cutlass::arch::NamedBarrier::arrive_and_wait( NumPrologueThreads, KdaChunkFwdRecompWUSm100NamedBarriers::PrologueCudaCore); @@ -464,6 +461,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } + g_pipeline.consumer_release(g_pipe_state_read); + ++g_pipe_state_read; + // Ensure all 128 prologue threads have finished writing sKG_out cutlass::arch::NamedBarrier::arrive_and_wait( NumPrologueThreads, KdaChunkFwdRecompWUSm100NamedBarriers::PrologueCudaCore); @@ -517,22 +517,12 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } } - // Release Q SMEM back to Load warp - q_pipeline.consumer_release(q_pipe_state_read); - ++q_pipe_state_read; - - // Need NamedBarrier to ensure all 128 prologue threads finish previous sKG_out writes (from KG - // store) - cutlass::arch::NamedBarrier::arrive_and_wait( - NumPrologueThreads, KdaChunkFwdRecompWUSm100NamedBarriers::PrologueCudaCore); - - // Compute QG = Q * exp2(G) and write to sKG_out (reuse output buffer) + // Compute QG = Q * exp2(G) #pragma unroll for (int ti = 0; ti < TileT / 16; ++ti) { int t = x_local + ti * 16; #pragma unroll for (int k_yi = 0; k_yi < TileK / 64; ++k_yi) { - int y = k_y_base + k_yi * 64; if (t < sub_seq_len) { // lo half (cols y..y+3) float2 qf_01 = __bfloat1622float2(q_reg[ti][k_yi].a01); @@ -554,18 +544,39 @@ struct KdaChunkFwdRecompWUMainloopSm100 { out.a23 = __float22bfloat162_rn(res_23); out.a45 = __float22bfloat162_rn(res_45); out.a67 = __float22bfloat162_rn(res_67); - *reinterpret_cast(&sKG_out(t, y)) = out; + q_reg[ti][k_yi] = out; } else { bf16x8 zero; zero.a01 = __float2bfloat162_rn(0.0f); zero.a23 = __float2bfloat162_rn(0.0f); zero.a45 = __float2bfloat162_rn(0.0f); zero.a67 = __float2bfloat162_rn(0.0f); - *reinterpret_cast(&sKG_out(t, y)) = zero; + q_reg[ti][k_yi] = zero; } } } + // Need NamedBarrier to ensure all 128 prologue threads finish previous sKG_out writes + cutlass::arch::NamedBarrier::arrive_and_wait( + NumPrologueThreads, KdaChunkFwdRecompWUSm100NamedBarriers::PrologueCudaCore); + + // write to sKG_out +#pragma unroll + for (int ti = 0; ti < TileT / 16; ++ti) { + int t = x_local + ti * 16; +#pragma unroll + for (int k_yi = 0; k_yi < TileK / 64; ++k_yi) { + int y = k_y_base + k_yi * 64; + *reinterpret_cast(&sKG_out(t, y)) = q_reg[ti][k_yi]; + } + } + + // NOTE: must make smem visible from CUDA Core (general proxy) to TMA (async proxy) + fence_view_async_shared(); + // Release Q SMEM back to Load warp + q_pipeline.consumer_release(q_pipe_state_read); + ++q_pipe_state_read; + // Ensure all 128 prologue threads have finished writing QG to sKG_out cutlass::arch::NamedBarrier::arrive_and_wait( NumPrologueThreads, KdaChunkFwdRecompWUSm100NamedBarriers::PrologueCudaCore); @@ -711,15 +722,15 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } - // Release V SMEM back to Load warp (done reading sV, all in sV_dst now) - v_pipeline.consumer_release(v_pipe_state_read); - ++v_pipe_state_read; - // Co-commit prologue_ready with Prologue → MMA can now consume k_mma + v_mma fence_view_async_shared(); prologue_ready_pipeline.producer_commit(prologue_ready_pipe_state_write); ++prologue_ready_pipe_state_write; + // Release V SMEM back to Load warp (done reading sV, all in sV_dst now) + v_pipeline.consumer_release(v_pipe_state_read); + ++v_pipe_state_read; + // ---- w/u output: wait K-GEMM & V-GEMM acc → T2R → bf16 → R2G ---- // Split into 2 iterations of TileK/2 to reduce register pressure (avoid spill) acc_done_pipeline.consumer_wait(acc_done_pipe_state_read); From bb39cf93ab0e236c25b3635a89ae83eb40d28b81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=97=A0=E8=A8=80=E7=8B=AC=E4=B8=8A=E6=9C=BA=E6=88=BF?= <88866917+sjmshsh@users.noreply.github.com> Date: Sat, 23 May 2026 16:11:16 +0800 Subject: [PATCH 15/34] [KDA] sm100 GVA enhance for chunk_intra and recomp_wu (#65) * feat(kda/sm100): support GVA (HV > HQK) in fwd intra/recomp kernels Follow the GVA pattern used in the SM90 KDA (and in gated_delta_rule GVA) so that the SM100 KDA forward pass can handle num_v_heads > num_qk_heads. C++ changes: - tile_scheduler: Params now carries heads_per_group; decode_tile_coord enumerates tiles in v-head space and returns both v_head_idx and qk_head_idx (= v_head_idx / heads_per_group). When HV == HQK this degenerates to the previous behaviour. - kda_config: KDA_fwd_intra_params / KDA_fwd_recomp_w_u_params split h into h_qk and h_v and cache heads_per_group; Akk and w/u/kg/qg layouts now live in v-head space. - intra kernel/mainloop: Q/K TMA descriptors use shape_QK (total, d, h_qk); g TMA uses shape_VG (total, d, h_v). Load warp slices Q/K with qk_head_idx and g with v_head_idx; Aqk row stride and beta stride now use params.h_v. - recomp_w_u kernel/mainloop: K/Q TMA descriptors use shape_QK; V/g TMA use shape_VG; Akk TMA uses shape_Akk (total, BT, h_v). Load warp slices K/Q with qk_head_idx and V/g/Akk with v_head_idx; w/u/kg/qg write stride and beta stride now use params.h_v. API / Python: - kda_sm100.cu: derive h_qk from Q/K and h_v from V/g; validate HV % HQK == 0 and beta/qg_out shapes. - cula/kda/chunk_intra.py: infer HQK from k.shape[2] and HV from v.shape[2]; allocate Aqk, Akk, w, kg, qg in v-head space; add shape assertions. Backward compatible: when HV == HQK, heads_per_group == 1 and qk_head_idx == v_head_idx, and all shapes/strides reduce to the pre-GVA layout. * add sm100 test * benchmark * benchmark and test * benchmark and test * benchmark and test * benchmark and test * benchmark and test * benchmark and test * benchmark and test * benchmark and test * benchmark * benchmark * benchmark * benchmark --------- Co-authored-by: sunnyxyli --- benchmarks/bench_kda_chunk_intra.py | 212 ++++------ benchmarks/bench_recompute_wu.py | 126 +++--- benchmarks/utils.py | 30 +- csrc/api/kda_sm100.cu | 88 +++- csrc/kda/sm100/kda_config.hpp | 50 ++- csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp | 29 +- .../sm100/kda_fwd_intra_mainloop_sm100.hpp | 51 ++- .../sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp | 33 +- .../kda_fwd_recomp_w_u_mainloop_sm100.hpp | 57 ++- csrc/kda/sm100/tile_scheduler.hpp | 29 +- cula/kda/chunk_intra.py | 21 +- tests/test_kda_gva_intra_sm100.py | 386 ++++++++++++++++++ 12 files changed, 828 insertions(+), 284 deletions(-) create mode 100644 tests/test_kda_gva_intra_sm100.py diff --git a/benchmarks/bench_kda_chunk_intra.py b/benchmarks/bench_kda_chunk_intra.py index 29aa5ab7..4c2718d5 100644 --- a/benchmarks/bench_kda_chunk_intra.py +++ b/benchmarks/bench_kda_chunk_intra.py @@ -12,6 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +bench_kda_chunk_intra.py — Benchmark: cuLA vs FLA Triton for chunk_kda_fwd_intra + +Supports both standard (HV=H) and GVA (HV > H) modes. +In GVA mode both FLA (v0.5.0+) and cuLA accept compact q/k in HQK space natively. + +Usage: + python bench_kda_chunk_intra.py [--heads H] [--hv HV] [--disable_recompute] +""" + import argparse import os import pathlib @@ -30,6 +40,7 @@ # Constant params B, H, D = 2, 64, 128 +HV = H # overridable via --hv; HV > H enables GVA mode BT = 64 # chunk size # Varlen benchmark params @@ -54,196 +65,117 @@ def accuracy_stats(a, b): # ============================================================================== -# Uniform seqlen benchmark +# Unified uniform seqlen benchmark (handles both standard and GVA) # ============================================================================== def benchmark_chunk_intra_uniform(): device = torch.device("cuda") chunk_size = BT + HQK = H + gva_mode = HV > HQK + group_size = HV // HQK T_vals = [512, 1024, 4096, 8192, 16384, 32768] - print("=" * 90) + gva_note = f"HQK={HQK} HV={HV} (group_size={group_size})" if gva_mode else f"H={HQK}" + print("=" * 100) print( - f" Uniform-Length ChunkIntra Benchmark: cuLA vs FLA Triton B={B} H={H} D={D} disable_recompute={DISABLE_RECOMPUTE}" + f" Uniform-Length ChunkIntra Benchmark: cuLA vs FLA Triton " + f"B={B} {gva_note} D={D} disable_recompute={DISABLE_RECOMPUTE}" ) - print("=" * 90) + print("=" * 100) print( f"{'B':>4} {'T':>7} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) - print("─" * 90) + print("─" * 100) for T in T_vals: seq_lens = [T] * B cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - q, k, v, g, beta, scale, cu_seqlens, chunk_indices = prepare_intra_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens) - - # Accuracy: run once and compare - out_fla = fla_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, + q, k, v, g, beta, scale, cu_seqlens, chunk_indices = prepare_intra_inputs( + B, T, HQK, D, device, cu_seqlens=cu_seqlens, num_v_heads=HV ) - out_cula = cula_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, + + common = dict( + q=q, k=k, v=v, gk=g, beta=beta, scale=scale, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=True, disable_recompute=DISABLE_RECOMPUTE, ) - # Compare the first output tensor (o) + + # Accuracy: run once and compare + out_fla = fla_chunk_kda_fwd_intra(**common) + out_cula = cula_chunk_kda_fwd_intra(**common) o_fla = out_fla[0] if isinstance(out_fla, (tuple, list)) else out_fla o_cula = out_cula[0] if isinstance(out_cula, (tuple, list)) else out_cula rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) # Performance - ms_fla = triton.testing.do_bench( - lambda: fla_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, - ), - ) - ms_cula = triton.testing.do_bench( - lambda: cula_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, - ), - ) + ms_fla = triton.testing.do_bench(lambda: fla_chunk_kda_fwd_intra(**common)) + ms_cula = triton.testing.do_bench(lambda: cula_chunk_kda_fwd_intra(**common)) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( f"{B:>4} {T:>7} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) - print("─" * 90) + print("─" * 100) # ============================================================================== -# Varlen benchmark +# Unified varlen benchmark (handles both standard and GVA) # ============================================================================== def benchmark_chunk_intra_varlen(): device = torch.device("cuda") chunk_size = BT + HQK = H + gva_mode = HV > HQK + group_size = HV // HQK total_len_vals = [8192, 16384, 32768, 65536] + gva_note = f"HQK={HQK} HV={HV} (group_size={group_size})" if gva_mode else f"H={HQK}" print() - print("=" * 100) + print("=" * 110) print( - f" Varlen ChunkIntra Benchmark: cuLA vs FLA Triton NUM_SEQS={NUM_SEQS} H={H} D={D} disable_recompute={DISABLE_RECOMPUTE}" + f" Varlen ChunkIntra Benchmark: cuLA vs FLA Triton " + f"NUM_SEQS={NUM_SEQS} {gva_note} D={D} disable_recompute={DISABLE_RECOMPUTE}" ) - print("=" * 100) + print("=" * 110) print( f"{'total_len':>10} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) - print("─" * 100) + print("─" * 110) for total_len in total_len_vals: seq_lens = generate_random_seq_lens(NUM_SEQS, total_len, MIN_SEQ_LEN, VARIANCE, SEED) T = total_len cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - q, k, v, g, beta, scale, cu_seqlens, chunk_indices = prepare_intra_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens) - - # Accuracy - out_fla = fla_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, + q, k, v, g, beta, scale, cu_seqlens, chunk_indices = prepare_intra_inputs( + 1, T, HQK, D, device, cu_seqlens=cu_seqlens, num_v_heads=HV ) - out_cula = cula_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, + + common = dict( + q=q, k=k, v=v, gk=g, beta=beta, scale=scale, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=True, disable_recompute=DISABLE_RECOMPUTE, ) + + # Accuracy + out_fla = fla_chunk_kda_fwd_intra(**common) + out_cula = cula_chunk_kda_fwd_intra(**common) o_fla = out_fla[0] if isinstance(out_fla, (tuple, list)) else out_fla o_cula = out_cula[0] if isinstance(out_cula, (tuple, list)) else out_cula rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) # Performance - ms_fla = triton.testing.do_bench( - lambda: fla_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, - ), - ) - ms_cula = triton.testing.do_bench( - lambda: cula_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=DISABLE_RECOMPUTE, - ), - ) + ms_fla = triton.testing.do_bench(lambda: fla_chunk_kda_fwd_intra(**common)) + ms_cula = triton.testing.do_bench(lambda: cula_chunk_kda_fwd_intra(**common)) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( f"{total_len:>10} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) - print("─" * 100) + print("─" * 110) if __name__ == "__main__": @@ -253,11 +185,37 @@ def benchmark_chunk_intra_varlen(): action="store_true", help="Disable recompute in both FLA and cuLA (pre-compute QG)", ) + parser.add_argument( + "--heads", + type=int, + default=None, + help=f"Override number of Q/K heads (H). Default: {H}.", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help=f"Override number of V heads (HV). Default: H (no GVA). Set HV > H to enable GVA mode.", + ) args = parser.parse_args() if args.disable_recompute: DISABLE_RECOMPUTE = True print("[Disable recompute] pre-compute QG in forward") + if args.heads is not None: + if args.heads <= 0: + raise ValueError(f"--heads must be a positive integer, got {args.heads}") + H = args.heads + HV = H # reset HV to new H before --hv override + + if args.hv is not None: + if args.hv < H or args.hv % H != 0: + raise ValueError(f"--hv must be a positive multiple of H ({H}), got {args.hv}") + HV = args.hv + + if HV > H: + print(f"[GVA] HV={HV} (H={H}, group_size={HV // H}x)") + benchmark_chunk_intra_uniform() benchmark_chunk_intra_varlen() diff --git a/benchmarks/bench_recompute_wu.py b/benchmarks/bench_recompute_wu.py index 0beca3b5..c40dac31 100644 --- a/benchmarks/bench_recompute_wu.py +++ b/benchmarks/bench_recompute_wu.py @@ -12,6 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +bench_recompute_wu.py — Benchmark: cuLA vs FLA Triton for recompute_w_u + +Supports both standard (HV=H) and GVA (HV > H) modes. +In GVA mode both FLA (v0.5.0+) and cuLA accept compact q/k in HQK space natively. + +Usage: + python bench_recompute_wu.py [--heads H] [--hv HV] [--disable_recompute] +""" + import argparse import os import pathlib @@ -29,9 +39,11 @@ import cula.cudac as cula_cuda from benchmarks.utils import SEED, exclusive_cumsum, generate_random_seq_lens, prepare_intra_inputs +from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra # Constant params B, H, D = 2, 64, 128 +HV = H # overridable via --hv; HV > H enables GVA mode BT = 64 # chunk size # Varlen benchmark params @@ -55,36 +67,27 @@ def accuracy_stats(a, b): return rmse, rel_max, mean_diff -def prepare_recompute_wu_inputs(B, T, H, D, device, cu_seqlens=None, chunk_size=BT): - """Prepare inputs for recompute_w_u benchmarking. +def prepare_recompute_wu_inputs(B, T, device, cu_seqlens=None, chunk_size=BT): + """Prepare inputs for recompute_w_u benchmarking (handles both MHA and GVA). - Runs chunk_kda_fwd_intra (FLA) to produce Akk, then returns - all tensors needed for recompute_w_u_fwd / recompute_w_u_cuda. + Uses cuLA's GVA-aware chunk_kda_fwd_intra to produce Akk in HV head space, + which is valid for both MHA (HV=H) and GVA (HV>H) layouts. """ q, k, v, g, beta, scale, cu_seqlens, chunk_indices = prepare_intra_inputs( - B, T, H, D, device, cu_seqlens=cu_seqlens, chunk_size=chunk_size + B, T, H, D, device, cu_seqlens=cu_seqlens, chunk_size=chunk_size, num_v_heads=HV ) - # Run FLA chunk_kda_fwd_intra to get Akk (shared input for both impls) - _, _, _, _, Aqk, Akk = fla_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=False, + _, _, _, _, _, Akk = cula_chunk_kda_fwd_intra( + q=q, k=k, v=v, gk=g, beta=beta, scale=scale, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=True, disable_recompute=False, ) return q, k, v, g, beta, Akk, cu_seqlens, chunk_indices def run_fla_recompute_wu(k, v, beta, Akk, q, gk, cu_seqlens, chunk_indices, disable_recompute): - """Run FLA recompute_w_u_fwd.""" + """FLA recompute_w_u reference (handles both MHA and GVA natively as of v0.5.0).""" return fla_recompute_w_u_fwd( k=k, v=v, @@ -98,11 +101,12 @@ def run_fla_recompute_wu(k, v, beta, Akk, q, gk, cu_seqlens, chunk_indices, disa def run_cula_recompute_wu(k, v, beta, Akk, q, gk, cu_seqlens, chunk_indices, chunk_size, disable_recompute): - """Run cuLA recompute_w_u_cuda.""" - w = torch.empty_like(k) + """cuLA recompute_w_u (handles both MHA and GVA; w/u/qg/kg allocated in HV head space).""" + B_flat, T, HV_out, Dv = v.shape + w = torch.empty_like(v) u = torch.empty_like(v) - qg = torch.empty_like(q) if disable_recompute else None - kg = torch.empty_like(k) if gk is not None else None + qg = torch.empty(B_flat, T, HV_out, Dv, device=q.device, dtype=q.dtype) if disable_recompute else None + kg = torch.empty_like(v) if gk is not None else None cula_cuda.recompute_w_u_cuda( k, v, beta, Akk, gk, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q if disable_recompute else None, qg @@ -111,29 +115,32 @@ def run_cula_recompute_wu(k, v, beta, Akk, q, gk, cu_seqlens, chunk_indices, chu # ============================================================================== -# Uniform seqlen benchmark +# Unified uniform seqlen benchmark (handles both standard and GVA) # ============================================================================== def benchmark_recompute_wu_uniform(): device = torch.device("cuda") chunk_size = BT + gva_mode = HV > H + gva_note = f"HQK={H} HV={HV} (group_size={HV // H})" if gva_mode else f"H={H}" T_vals = [512, 1024, 4096, 8192, 16384, 32768] - print("=" * 90) + print("=" * 100) print( - f" Uniform-Length RecomputeWU Benchmark: cuLA vs FLA Triton B={B} H={H} D={D} disable_recompute={DISABLE_RECOMPUTE}" + f" Uniform-Length RecomputeWU Benchmark: cuLA vs FLA Triton " + f"B={B} {gva_note} D={D} disable_recompute={DISABLE_RECOMPUTE}" ) - print("=" * 90) + print("=" * 100) print( f"{'B':>4} {'T':>7} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) - print("─" * 90) + print("─" * 100) for T in T_vals: seq_lens = [T] * B cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) q, k, v, g, beta, Akk, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs( - B, T, H, D, device, cu_seqlens=cu_seqlens, chunk_size=chunk_size + B, T, device, cu_seqlens=cu_seqlens, chunk_size=chunk_size ) # Accuracy: run once and compare @@ -144,17 +151,12 @@ def benchmark_recompute_wu_uniform(): k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, chunk_size, DISABLE_RECOMPUTE ) - # Compare w, u, qg, kg stats = {} for name, t_fla, t_cula in [ - ("w", w_fla, w_cula), - ("u", u_fla, u_cula), - ("qg", qg_fla, qg_cula), - ("kg", kg_fla, kg_cula), + ("w", w_fla, w_cula), ("u", u_fla, u_cula), ("qg", qg_fla, qg_cula), ("kg", kg_fla, kg_cula), ]: if t_fla is not None and t_cula is not None: stats[name] = accuracy_stats(t_fla, t_cula) - # Use max across all outputs for display rmse = max(s[0] for s in stats.values()) rel_max = max(s[1] for s in stats.values()) mean_diff = max(s[2] for s in stats.values()) @@ -172,27 +174,30 @@ def benchmark_recompute_wu_uniform(): f"{B:>4} {T:>7} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) - print("─" * 90) + print("─" * 100) # ============================================================================== -# Varlen benchmark +# Unified varlen benchmark (handles both standard and GVA) # ============================================================================== def benchmark_recompute_wu_varlen(): device = torch.device("cuda") chunk_size = BT + gva_mode = HV > H + gva_note = f"HQK={H} HV={HV} (group_size={HV // H})" if gva_mode else f"H={H}" total_len_vals = [8192, 16384, 32768, 65536] print() - print("=" * 100) + print("=" * 110) print( - f" Varlen RecomputeWU Benchmark: cuLA vs FLA Triton NUM_SEQS={NUM_SEQS} H={H} D={D} disable_recompute={DISABLE_RECOMPUTE}" + f" Varlen RecomputeWU Benchmark: cuLA vs FLA Triton " + f"NUM_SEQS={NUM_SEQS} {gva_note} D={D} disable_recompute={DISABLE_RECOMPUTE}" ) - print("=" * 100) + print("=" * 110) print( f"{'total_len':>10} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) - print("─" * 100) + print("─" * 110) for total_len in total_len_vals: seq_lens = generate_random_seq_lens(NUM_SEQS, total_len, MIN_SEQ_LEN, VARIANCE, SEED) @@ -200,7 +205,7 @@ def benchmark_recompute_wu_varlen(): cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) q, k, v, g, beta, Akk, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs( - 1, T, H, D, device, cu_seqlens=cu_seqlens, chunk_size=chunk_size + 1, T, device, cu_seqlens=cu_seqlens, chunk_size=chunk_size ) # Accuracy @@ -211,17 +216,12 @@ def benchmark_recompute_wu_varlen(): k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, chunk_size, DISABLE_RECOMPUTE ) - # Compare w, u, qg, kg stats = {} for name, t_fla, t_cula in [ - ("w", w_fla, w_cula), - ("u", u_fla, u_cula), - ("qg", qg_fla, qg_cula), - ("kg", kg_fla, kg_cula), + ("w", w_fla, w_cula), ("u", u_fla, u_cula), ("qg", qg_fla, qg_cula), ("kg", kg_fla, kg_cula), ]: if t_fla is not None and t_cula is not None: stats[name] = accuracy_stats(t_fla, t_cula) - # Use max across all outputs for display rmse = max(s[0] for s in stats.values()) rel_max = max(s[1] for s in stats.values()) mean_diff = max(s[2] for s in stats.values()) @@ -239,7 +239,7 @@ def benchmark_recompute_wu_varlen(): f"{total_len:>10} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) - print("─" * 100) + print("─" * 110) def check_determinism(num_seqs=NUM_SEQS, T=2001, H=H, iters=1000): @@ -251,7 +251,7 @@ def check_determinism(num_seqs=NUM_SEQS, T=2001, H=H, iters=1000): cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) q, k, v, g, beta, Akk, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs( - B=1, T=T, H=H, D=D, device=device, cu_seqlens=cu_seqlens, chunk_size=chunk_size + B=1, T=T, device=device, cu_seqlens=cu_seqlens, chunk_size=chunk_size ) ref_w, ref_u, ref_qg, ref_kg = run_cula_recompute_wu( @@ -288,12 +288,38 @@ def check_determinism(num_seqs=NUM_SEQS, T=2001, H=H, iters=1000): action="store_true", help="Disable recompute in both FLA and cuLA (pre-compute QG)", ) + parser.add_argument( + "--heads", + type=int, + default=None, + help=f"Override number of Q/K heads (H). Default: {H}.", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help=f"Override number of V heads (HV). Default: H (no GVA). Set HV > H to enable GVA mode.", + ) args = parser.parse_args() if args.disable_recompute: DISABLE_RECOMPUTE = True print("[Disable recompute] pre-compute QG in forward") + if args.heads is not None: + if args.heads <= 0: + raise ValueError(f"--heads must be a positive integer, got {args.heads}") + H = args.heads + HV = H # reset HV to new H before --hv override + + if args.hv is not None: + if args.hv < H or args.hv % H != 0: + raise ValueError(f"--hv must be a positive multiple of H ({H}), got {args.hv}") + HV = args.hv + + if HV > H: + print(f"[GVA] HV={HV} (H={H}, group_size={HV // H}x)") + check_determinism(iters=100000) benchmark_recompute_wu_uniform() diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 75d7ef52..64591662 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -318,11 +318,25 @@ def prepare_safe_gate_inputs( ) -def prepare_intra_inputs(batch_size, T, H, D, device, cu_seqlens=None, chunk_size=CHUNK_SIZE, seed=SEED): +def prepare_intra_inputs( + batch_size, T, H, D, device, cu_seqlens=None, chunk_size=CHUNK_SIZE, seed=SEED, num_v_heads=None +): """Prepare preprocessed inputs ready for chunk_kda_fwd_intra. - All tensors are flattened to (1, B*T, ...) for cu_seqlens compatibility. + Supports both standard (HV=H) and GVA (HV > H) layouts via ``num_v_heads``: + + q, k : (batch_size_flat, T, H, D) — Q/K head space (always compact) + v : (batch_size_flat, T, HV, D) — V head space + g : (batch_size_flat, T, HV, D) — gate in V head space (after cumsum) + beta : (batch_size_flat, T, HV) — beta in V head space + + When ``num_v_heads`` is None or equal to H this matches the original non-GVA + behaviour exactly. All tensors are flattened to batch_size=1 for cu_seqlens + compatibility. """ + HV = H if num_v_heads is None else num_v_heads + assert HV >= H and HV % H == 0, f"num_v_heads ({HV}) must be a positive multiple of H ({H})" + dtype = torch.bfloat16 scale = D ** (-0.5) @@ -330,9 +344,9 @@ def prepare_intra_inputs(batch_size, T, H, D, device, cu_seqlens=None, chunk_siz q = torch.randn(batch_size, T, H, D, dtype=dtype, device=device) k = torch.randn(batch_size, T, H, D, dtype=dtype, device=device) - v = torch.randn(batch_size, T, H, D, dtype=dtype, device=device) - g_raw = torch.randn(batch_size, T, H, D, dtype=dtype, device=device) - beta = torch.randn(batch_size, T, H, dtype=torch.float, device=device).sigmoid() + v = torch.randn(batch_size, T, HV, D, dtype=dtype, device=device) + g_raw = torch.randn(batch_size, T, HV, D, dtype=dtype, device=device) + beta = torch.randn(batch_size, T, HV, dtype=torch.float, device=device).sigmoid() # l2norm q, k q, _ = l2norm_fwd(q) @@ -342,9 +356,9 @@ def prepare_intra_inputs(batch_size, T, H, D, device, cu_seqlens=None, chunk_siz if batch_size != 1: q, k, v, g_raw, beta = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (q, k, v, g_raw, beta)) - # gate preprocessing - A_log = torch.randn(H, dtype=torch.float, device=device) - dt_bias = torch.randn(H * D, dtype=torch.float, device=device) + # gate preprocessing — A_log / dt_bias live in HV head space + A_log = torch.randn(HV, dtype=torch.float, device=device) + dt_bias = torch.randn(HV * D, dtype=torch.float, device=device) chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None diff --git a/csrc/api/kda_sm100.cu b/csrc/api/kda_sm100.cu index ac324113..ff89887a 100644 --- a/csrc/api/kda_sm100.cu +++ b/csrc/api/kda_sm100.cu @@ -37,7 +37,33 @@ ChunkKDAFwdIntra( KDA_fwd_intra_params params; params.total_q_len = q.size(0) * q.size(1); params.b = cu_seqlens.size(0) - 1; - params.h = q.size(2); + // GVA: Q/K are in h_qk head space (from q.size(2)); g/beta/Aqk/Akk are in h_v head + // space (from g.size(2)). When HV == HQK, heads_per_group == 1 and behaviour matches + // the pre-GVA path. + params.h_qk = q.size(2); + params.h_v = g.size(2); + TORCH_CHECK( + k.size(2) == params.h_qk, + "ChunkKDAFwdIntra: k.size(2) (", + k.size(2), + ") must match q.size(2) (", + params.h_qk, + ") under GVA (Q/K share h_qk)."); + TORCH_CHECK( + beta.size(-1) == params.h_v, + "ChunkKDAFwdIntra: beta.size(-1) (", + beta.size(-1), + ") must equal h_v (", + params.h_v, + ")."); + TORCH_CHECK( + params.h_qk > 0 && params.h_v > 0 && params.h_v % params.h_qk == 0, + "ChunkKDAFwdIntra: h_v (", + params.h_v, + ") must be a positive multiple of h_qk (", + params.h_qk, + ")."); + params.heads_per_group = params.h_v / params.h_qk; params.d = q.size(3); params.chunk_size = chunk_size; params.scale = scale; @@ -56,13 +82,15 @@ ChunkKDAFwdIntra( params.chunk_indices_ptr = chunk_indices.data_ptr(); params.Aqk_out_ptr = Aqk_out.data_ptr(); params.Akk_out_ptr = Akk_out.data_ptr(); - params.shape_Akk = cute::make_shape(params.total_q_len, params.chunk_size, params.h); - params.stride_Akk = cute::make_stride(params.chunk_size * params.h, cute::_1{}, params.chunk_size); + // Akk is laid out per v-head: (total_len, chunk_size, h_v). + params.shape_Akk = cute::make_shape(params.total_q_len, params.chunk_size, params.h_v); + params.stride_Akk = cute::make_stride(params.chunk_size * params.h_v, cute::_1{}, params.chunk_size); int tile_num = chunk_indices.size(0); auto device_prop = at::cuda::getCurrentDeviceProperties(); params.num_sm = device_prop->multiProcessorCount; - params.tile_scheduler_params = - StaticPersistentTileScheduler::Params{tile_num, params.h, params.num_sm, (int*)tile_counter.data_ptr()}; + // Tiles are enumerated in v-head space. + params.tile_scheduler_params = StaticPersistentTileScheduler::Params{ + tile_num, params.h_v, params.heads_per_group, params.num_sm, (int*)tile_counter.data_ptr()}; kda::sm100::run_kda_fwd_intra_sm100(params, at::cuda::getCurrentCUDAStream()); } @@ -85,7 +113,31 @@ ChunkKDAFwdRecompWU( KDA_fwd_recomp_w_u_params params; params.total_len = k.size(0) * k.size(1); params.b = cu_seqlens.size(0) - 1; - params.h = k.size(2); + // GVA: K (and optional Q) live in h_qk space; V/G/beta/A/w/u/kg/qg live in h_v space. + params.h_qk = k.size(2); + params.h_v = v.size(2); + TORCH_CHECK( + g.size(2) == params.h_v, + "ChunkKDAFwdRecompWU: g.size(2) (", + g.size(2), + ") must equal v.size(2) (", + params.h_v, + ")."); + TORCH_CHECK( + beta.size(-1) == params.h_v, + "ChunkKDAFwdRecompWU: beta.size(-1) (", + beta.size(-1), + ") must equal h_v (", + params.h_v, + ")."); + TORCH_CHECK( + params.h_qk > 0 && params.h_v > 0 && params.h_v % params.h_qk == 0, + "ChunkKDAFwdRecompWU: h_v (", + params.h_v, + ") must be a positive multiple of h_qk (", + params.h_qk, + ")."); + params.heads_per_group = params.h_v / params.h_qk; params.d = k.size(3); params.chunk_size = chunk_size; TORCH_CHECK( @@ -108,14 +160,32 @@ ChunkKDAFwdRecompWU( TORCH_CHECK( has_q == has_qg_out, "ChunkKDAFwdRecompWU: q and qg_out must either both be provided or both be omitted."); params.store_qg = has_q && has_qg_out; + if (params.store_qg) { + TORCH_CHECK( + q->size(2) == params.h_qk, + "ChunkKDAFwdRecompWU: q.size(2) (", + q->size(2), + ") must equal h_qk (", + params.h_qk, + ")."); + TORCH_CHECK( + qg_out->size(2) == params.h_v, + "ChunkKDAFwdRecompWU: qg_out.size(2) (", + qg_out->size(2), + ") must equal h_v (", + params.h_v, + ")."); + } params.q_ptr = params.store_qg ? q->data_ptr() : nullptr; params.qg_out_ptr = params.store_qg ? qg_out->data_ptr() : nullptr; - params.shape_wukg = cute::make_shape(params.total_len, params.d, params.h); - params.stride_wukg = cute::make_stride(params.d * params.h, cute::_1{}, params.d); + // w/u/kg/qg are per v-head: (total_len, d, h_v). + params.shape_wukg = cute::make_shape(params.total_len, params.d, params.h_v); + params.stride_wukg = cute::make_stride(params.d * params.h_v, cute::_1{}, params.d); int tile_num = chunk_indices.size(0); auto device_prop = at::cuda::getCurrentDeviceProperties(); params.num_sm = device_prop->multiProcessorCount; - params.tile_scheduler_params = StaticPersistentTileScheduler::Params{tile_num, params.h, params.num_sm, nullptr}; + params.tile_scheduler_params = StaticPersistentTileScheduler::Params{ + tile_num, params.h_v, params.heads_per_group, params.num_sm, nullptr}; kda::sm100::run_kda_fwd_recomp_w_u_sm100(params, at::cuda::getCurrentCUDAStream()); } \ No newline at end of file diff --git a/csrc/kda/sm100/kda_config.hpp b/csrc/kda/sm100/kda_config.hpp index 6f96529a..67b496a9 100644 --- a/csrc/kda/sm100/kda_config.hpp +++ b/csrc/kda/sm100/kda_config.hpp @@ -17,12 +17,18 @@ #include "kda/sm100/tile_scheduler.hpp" struct KDA_fwd_intra_params { - using GmemShapeAkk = cute::Shape; // (seqlen_kv, seqlen_kv, h) + // Akk shape is (total_seqlen, chunk_size, num_v_heads). Under GVA (num_v_heads > num_qk_heads), + // Aqk and Akk are produced per v-head because g/beta/Akk scaling all live in v-head space. + using GmemShapeAkk = cute::Shape; // (seqlen_kv, chunk_size, h_v) using GmemStrideAkk = cute::Stride; int total_q_len; int b; - int h; + // GVA: Q/K are sized by num_qk_heads; V, g, beta are sized by num_v_heads; Aqk/Akk are per v-head. + // When num_v_heads == num_qk_heads, heads_per_group == 1 and behaviour matches the pre-GVA path. + int h_qk; + int h_v; + int heads_per_group; // = h_v / h_qk, precomputed on host int d; int chunk_size; float scale; @@ -30,12 +36,12 @@ struct KDA_fwd_intra_params { bool unified_gref; bool is_beta_bf16; - void* __restrict__ q_ptr; //[b, t, h, d] - void* __restrict__ k_ptr; //[b, t, h, d] - void* __restrict__ g_ptr; //[b, t, h, d] - void* __restrict__ beta_ptr; //[b, t, h] - void* __restrict__ Aqk_out_ptr; //[b, t, h, BT] - void* __restrict__ Akk_out_ptr; //[b, t, h, BT] + void* __restrict__ q_ptr; //[b, t, h_qk, d] + void* __restrict__ k_ptr; //[b, t, h_qk, d] + void* __restrict__ g_ptr; //[b, t, h_v, d] + void* __restrict__ beta_ptr; //[b, t, h_v] + void* __restrict__ Aqk_out_ptr; //[b, t, h_v, BT] + void* __restrict__ Akk_out_ptr; //[b, t, h_v, BT] void* __restrict__ cu_seqlens_ptr; //[b + 1] void* __restrict__ chunk_indices_ptr; //[(b * t) / chunk_size, 2] @@ -48,28 +54,32 @@ struct KDA_fwd_intra_params { }; struct KDA_fwd_recomp_w_u_params { - using GmemShapeWUKg = cute::Shape; // (seqlen_kv, seqlen_kv, h) + // w/u/kg/qg all have shape (total_seqlen, d, num_v_heads) under GVA. + using GmemShapeWUKg = cute::Shape; // (seqlen_kv, d, h_v) using GmemStrideWUKg = cute::Stride; int total_len; int b; - int h; + // GVA: K and (optional) Q are sized by num_qk_heads; V/G/beta/Akk/w/u/kg/qg are per v-head. + int h_qk; + int h_v; + int heads_per_group; // = h_v / h_qk, precomputed on host int d; int chunk_size; bool is_beta_bf16; - void* __restrict__ k_ptr; //[b, t, h, d] - void* __restrict__ v_ptr; //[b, t, h, d] - void* __restrict__ q_ptr; //[b, t, h, d] (optional, for StoreQG) - void* __restrict__ beta_ptr; //[b, t, h] - void* __restrict__ A_ptr; //[b. t, h, BT] - void* __restrict__ g_ptr; //[b, t, h, d] + void* __restrict__ k_ptr; //[b, t, h_qk, d] + void* __restrict__ v_ptr; //[b, t, h_v, d] + void* __restrict__ q_ptr; //[b, t, h_qk, d] (optional, for StoreQG) + void* __restrict__ beta_ptr; //[b, t, h_v] + void* __restrict__ A_ptr; //[b, t, h_v, BT] + void* __restrict__ g_ptr; //[b, t, h_v, d] void* __restrict__ cu_seqlens_ptr; //[b + 1] void* __restrict__ chunk_indices_ptr; //[(b * t) / chunk_size, 2] - void* __restrict__ w_out_ptr; //[b, t, h, d] - void* __restrict__ u_out_ptr; //[b, t, h, d] - void* __restrict__ kg_out_ptr; //[b, t, h, d] - void* __restrict__ qg_out_ptr; //[b, t, h, d] (optional, for StoreQG) + void* __restrict__ w_out_ptr; //[b, t, h_v, d] + void* __restrict__ u_out_ptr; //[b, t, h_v, d] + void* __restrict__ kg_out_ptr; //[b, t, h_v, d] + void* __restrict__ qg_out_ptr; //[b, t, h_v, d] (optional, for StoreQG) bool store_qg; diff --git a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp index f314723f..f928616d 100644 --- a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp @@ -53,8 +53,8 @@ struct KdaChunkFwdIntraKernelSm100 { using SmemLayoutInputFP32 = typename Mainloop::SmemLayoutInputFP32; // TMA params (for host launcher) - template - using TmaParams = typename Mainloop::template TmaParams; + template + using TmaParams = typename Mainloop::template TmaParams; // Pipeline types (for construction in operator()) using PipelineQKG = typename Mainloop::PipelineQKG; @@ -318,29 +318,40 @@ __launch_bounds__(512, 1, 1) kda_fwd_intra_sm100_kernel_entry( template inline void run_kda_fwd_intra_sm100_impl_dispatch(KDA_fwd_intra_params& params, cudaStream_t stream) { - auto shape_QKG = make_shape(params.total_q_len, params.d, params.h); - auto stride_QKG = make_stride(params.h * params.d, _1{}, params.d); + // GVA: Q/K are sized by `h_qk`; G is sized by `h_v`. When HV == HQK + // (heads_per_group == 1), shape_QK and shape_VG coincide with the + // pre-GVA shape_QKG and behaviour is unchanged. + auto shape_QK = make_shape(params.total_q_len, params.d, params.h_qk); + auto stride_QK = make_stride(params.h_qk * params.d, _1{}, params.d); + auto shape_VG = make_shape(params.total_q_len, params.d, params.h_v); + auto stride_VG = make_stride(params.h_v * params.d, _1{}, params.d); // --- Build TMA descriptors --- auto tma_Q = cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((ku::bf16*)params.q_ptr), make_layout(shape_QKG, stride_QKG)), + make_tensor(make_gmem_ptr((ku::bf16*)params.q_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); auto tma_K = cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((ku::bf16*)params.k_ptr), make_layout(shape_QKG, stride_QKG)), + make_tensor(make_gmem_ptr((ku::bf16*)params.k_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); auto tma_G = cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_QKG, stride_QKG)), + make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), typename Kernel::SmemLayoutInputFP32{}); // --- Pack TMA params --- - typename Kernel::template TmaParams + typename Kernel::template TmaParams< + decltype(shape_QK), + decltype(shape_VG), + decltype(tma_Q), + decltype(tma_K), + decltype(tma_G)> tma_params = { - shape_QKG, + shape_QK, + shape_VG, tma_Q, tma_K, tma_G, diff --git a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp index 849e910c..7b624ebb 100644 --- a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp @@ -226,9 +226,13 @@ struct KdaChunkFwdIntraMainloopSm100 { }; // ===================== TMA Params ===================== - template + // GVA: Q/K live in h_qk head space (shape_qk), while G lives in h_v + // head space (shape_vg). When h_v == h_qk both shapes coincide and the + // TMA descriptors degrade to the pre-GVA behaviour. + template struct TmaParams { - ShapeQKG shape_qkg; + ShapeQK shape_qk; + ShapeVG shape_vg; TMA_Q tma_q; TMA_K tma_k; TMA_G tma_g; @@ -317,7 +321,10 @@ struct KdaChunkFwdIntraMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // head_idx here is the v-head index (Aqk/Akk/beta/g live in v-head space). + // qk_head_idx is only consumed by the TMA load warp for Q/K slicing. + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -501,7 +508,8 @@ struct KdaChunkFwdIntraMainloopSm100 { int token_offset = cu_seqlens_ptr[batch_idx]; int row = idx_in_warpgroup % 64; int BT = TileT; - int H = params.h; + // Aqk is laid out per v-head: row-stride is h_v * BT, head slot offset is head_idx * BT. + int H = params.h_v; __nv_bfloat16* Aqk_base = reinterpret_cast<__nv_bfloat16*>(params.Aqk_out_ptr); __nv_bfloat16* qk_out_row = Aqk_base + static_cast(token_offset + tile_idx * TileT + row) * H * BT + head_idx * BT; @@ -567,7 +575,10 @@ struct KdaChunkFwdIntraMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // MMA loop does not actually consume head_idx, but we decode to advance the + // same tile space as the other warps (num_blocks * num_v_heads). + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -702,21 +713,24 @@ struct KdaChunkFwdIntraMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - // Decode tile coordinates - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // Decode tile coordinates. head_idx is the v-head index (used for G), + // and qk_head_idx is the companion Q/K head (computed from heads_per_group). + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); - int head_idx = get<1>(blk_coord); + int head_idx = get<1>(blk_coord); // v-head index int tile_idx = get<2>(blk_coord); + int qk_head_idx = get<3>(blk_coord); // == head_idx / heads_per_group int token_offset = cu_seqlens_ptr[batch_idx]; int seq_len = cu_seqlens_ptr[batch_idx + 1] - cu_seqlens_ptr[batch_idx]; int sub_seq_len = min(TileT, seq_len - tile_idx * TileT); Tensor mQ = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_q.get_tma_tensor(tma_params.shape_qkg)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_q.get_tma_tensor(tma_params.shape_qk)); Tensor mK = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qkg)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qk)); Tensor mG = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_qkg)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); // TMA load body (Q, K, G — unified pipeline, single barrier per stage) CUTE_NO_UNROLL @@ -726,12 +740,13 @@ struct KdaChunkFwdIntraMainloopSm100 { Tensor sK = make_tensor(make_smem_ptr(shared_plan->k[buf_idx].data()), SmemLayoutInputBF16{}); Tensor sG = make_tensor(make_smem_ptr(shared_plan->g[buf_idx].data()), SmemLayoutInputFP32{}); + // GVA: K and Q are sliced by qk_head_idx; G is sliced by head_idx (v-head). Tensor gK = local_tile( - mK(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); + mK(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); Tensor gG = local_tile( mG(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); Tensor gQ = local_tile( - mQ(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); + mQ(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); // Single acquire for all three TMA copies qkg_load_pipeline.producer_acquire(qkg_load_pipe_state_write); @@ -768,7 +783,9 @@ struct KdaChunkFwdIntraMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // Akk is laid out per v-head (params.shape_Akk uses h_v), so we index by head_idx. + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -881,7 +898,9 @@ struct KdaChunkFwdIntraMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // beta is per v-head: layout (total_seqlen, h_v), row stride = h_v. + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -895,7 +914,7 @@ struct KdaChunkFwdIntraMainloopSm100 { shared_plan->beta_smem[beta_pipe_state_write.index()][thread_idx] = (thread_idx < sub_seq_len) ? float(reinterpret_cast( - params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h + head_idx]) + params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h_v + head_idx]) : float(0); } fence_view_async_shared(); diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp index 73cb4089..2cae6c04 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp @@ -41,14 +41,16 @@ struct KdaChunkFwdRecompWUKernelSm100 { // TMA params (for host launcher) template < - typename ShapeKVG, + typename ShapeQK, + typename ShapeVG, typename ShapeAkk, typename TMA_V, typename TMA_K, typename TMA_G, typename TMA_Akk, typename TMA_Q = int> - using TmaParams = typename Mainloop::template TmaParams; + using TmaParams = + typename Mainloop::template TmaParams; // Pipeline types (for construction in operator()) using PipelineA = typename Mainloop::PipelineA; @@ -431,25 +433,29 @@ __launch_bounds__(384, 1, 1) kda_fwd_recomp_w_u_sm100_kernel_entry( template inline void run_kda_fwd_recomp_w_u_sm100_impl_dispatch(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { - auto shape_KVG = make_shape(params.total_len, params.d, params.h); - auto stride_KVG = make_stride(params.h * params.d, _1{}, params.d); - auto shape_Akk = make_shape(params.total_len, params.chunk_size, params.h); - auto stride_Akk = make_stride(params.h * params.chunk_size, _1{}, params.chunk_size); + // GVA: K and (optional) Q are sized by h_qk; V and G are sized by h_v. + // Akk lives in v-head space (BT x BT per v-head). + auto shape_QK = make_shape(params.total_len, params.d, params.h_qk); + auto stride_QK = make_stride(params.h_qk * params.d, _1{}, params.d); + auto shape_VG = make_shape(params.total_len, params.d, params.h_v); + auto stride_VG = make_stride(params.h_v * params.d, _1{}, params.d); + auto shape_Akk = make_shape(params.total_len, params.chunk_size, params.h_v); + auto stride_Akk = make_stride(params.h_v * params.chunk_size, _1{}, params.chunk_size); // --- Build TMA descriptors --- auto tma_V = cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((bf16*)params.v_ptr), make_layout(shape_KVG, stride_KVG)), + make_tensor(make_gmem_ptr((bf16*)params.v_ptr), make_layout(shape_VG, stride_VG)), typename Kernel::SmemLayoutInputBF16{}); auto tma_K = cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((bf16*)params.k_ptr), make_layout(shape_KVG, stride_KVG)), + make_tensor(make_gmem_ptr((bf16*)params.k_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); auto tma_G = cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_KVG, stride_KVG)), + make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), typename Kernel::SmemLayoutInputFP32{}); auto tma_Akk = cute::make_tma_copy( @@ -457,12 +463,12 @@ run_kda_fwd_recomp_w_u_sm100_impl_dispatch(KDA_fwd_recomp_w_u_params& params, cu make_tensor(make_gmem_ptr((bf16*)params.A_ptr), make_layout(shape_Akk, stride_Akk)), typename Kernel::SmemLayoutInputAkkBF16{}); - // Q TMA descriptor (only meaningful when StoreQG=true) + // Q TMA descriptor (only meaningful when StoreQG=true). Q lives in h_qk head space. auto tma_Q = [&]() { if constexpr (Kernel::StoreQG) { return cute::make_tma_copy( SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((bf16*)params.q_ptr), make_layout(shape_KVG, stride_KVG)), + make_tensor(make_gmem_ptr((bf16*)params.q_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); } else { return 0; // placeholder, not used @@ -471,14 +477,15 @@ run_kda_fwd_recomp_w_u_sm100_impl_dispatch(KDA_fwd_recomp_w_u_params& params, cu // --- Pack TMA params --- typename Kernel::template TmaParams< - decltype(shape_KVG), + decltype(shape_QK), + decltype(shape_VG), decltype(shape_Akk), decltype(tma_V), decltype(tma_K), decltype(tma_G), decltype(tma_Akk), decltype(tma_Q)> - tma_params = {shape_KVG, shape_Akk, tma_V, tma_K, tma_G, tma_Akk, tma_Q}; + tma_params = {shape_QK, shape_VG, shape_Akk, tma_V, tma_K, tma_G, tma_Akk, tma_Q}; // --- Launch config --- auto kernel_fn = &kda_fwd_recomp_w_u_sm100_kernel_entry; diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index 8b51b65b..d5ff9a9e 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -190,8 +190,11 @@ struct KdaChunkFwdRecompWUMainloopSm100 { }; // ===================== TMA Params ===================== + // GVA: K and (optional) Q live in h_qk head space (shape_qk), while V + // and G live in h_v head space (shape_vg). Akk is per v-head. template < - typename ShapeKVG, + typename ShapeQK, + typename ShapeVG, typename ShapeAkk, typename TMA_V, typename TMA_K, @@ -199,7 +202,8 @@ struct KdaChunkFwdRecompWUMainloopSm100 { typename TMA_Akk, typename TMA_Q = int> struct TmaParams { - ShapeKVG shape_kvg; + ShapeQK shape_qk; + ShapeVG shape_vg; ShapeAkk shape_Akk; TMA_V tma_v; TMA_K tma_k; @@ -255,7 +259,10 @@ struct KdaChunkFwdRecompWUMainloopSm100 { CUTE_NO_UNROLL for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // Prologue touches K (h_qk) and G (h_v) + beta (h_v) + optional Q (h_qk). + // head_idx is the v-head index; qk_head_idx is derived via heads_per_group. + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -646,7 +653,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { CUTE_NO_UNROLL for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // Epilogue consumes V/beta (both h_v) and writes w/u/kg/qg (all h_v). + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -746,9 +755,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // each thread processes one row of W/U (TileK columns) int row = (idx_in_wg / 32) * 16 + (idx_in_wg % 16); - // GMEM output address: layout [total_len, d, h], stride [d*h, 1, d] + // GMEM output address: layout [total_len, d, h_v], stride [d*h_v, 1, d] __nv_bfloat16* out_row_base = - out_ptr_base + (token_offset_cur + row) * params.d * params.h + head_idx * params.d; + out_ptr_base + (token_offset_cur + row) * params.d * params.h_v + head_idx * params.d; constexpr int QuarK = TileK / 4; @@ -810,7 +819,8 @@ struct KdaChunkFwdRecompWUMainloopSm100 { CUTE_NO_UNROLL for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { // int tid = tile_scheduler.get_current_tile_id(); - // auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h_v, params.heads_per_group, + // chunk_indices_ptr, cu_seqlens_ptr); // ============================================================ // Once per WU: Wait for Akk in SMEM (from Load warp) @@ -890,31 +900,36 @@ struct KdaChunkFwdRecompWUMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - // Decode tile coordinates - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // Decode tile coordinates. head_idx is the v-head (used for V/G/Akk + // TMA loads); qk_head_idx (= head_idx / heads_per_group) is used for + // K/Q TMA loads under GVA. + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); - int head_idx = get<1>(blk_coord); + int head_idx = get<1>(blk_coord); // v-head int tile_idx = get<2>(blk_coord); + int qk_head_idx = get<3>(blk_coord); // qk-head int token_offset = cu_seqlens_ptr[batch_idx]; int seq_len = cu_seqlens_ptr[batch_idx + 1] - cu_seqlens_ptr[batch_idx]; int sub_seq_len = min(TileT, seq_len - tile_idx * TileT); // Build GMEM tensor views (with domain offset for batch) + // K and Q live in h_qk head space (shape_qk); V, G and Akk live in h_v space. Tensor mK = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_kvg)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qk)); Tensor mV = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_v.get_tma_tensor(tma_params.shape_kvg)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_v.get_tma_tensor(tma_params.shape_vg)); Tensor mG = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_kvg)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); Tensor mA = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_akk.get_tma_tensor(tma_params.shape_Akk)); - // Q GMEM tensor (only used when StoreQG=true) + // Q GMEM tensor (only used when StoreQG=true). Q lives in h_qk space. [[maybe_unused]] auto mQ = [&]() { if constexpr (StoreQG) { return domain_offset( make_coord(token_offset, _0{}, _0{}), - tma_params.tma_q.get_tma_tensor(tma_params.shape_kvg)); + tma_params.tma_q.get_tma_tensor(tma_params.shape_qk)); } else { return 0; // unused placeholder } @@ -947,8 +962,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { Tensor sG = make_tensor( make_smem_ptr(shared_plan->g[g_pipe_state_write.index()].data()), SmemLayoutInputFP32{}); + // GVA slicing: K uses qk_head_idx; V and G use the v-head index. Tensor gK = local_tile( - mK(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); + mK(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); Tensor gV = local_tile( mV(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); Tensor gG = local_tile( @@ -974,8 +990,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { Tensor sQ = make_tensor( make_smem_ptr(shared_plan->q_buf.q[q_pipe_state_write.index()].data()), SmemLayoutInputBF16{}); + // Q (StoreQG) lives in h_qk space → slice by qk_head_idx. Tensor gQ = local_tile( - mQ(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); + mQ(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); q_pipeline.producer_acquire(q_pipe_state_write); ku::launch_tma_copy( tma_params.tma_q, gQ, sQ, *q_pipeline.producer_get_barrier(q_pipe_state_write)); @@ -1008,7 +1025,9 @@ struct KdaChunkFwdRecompWUMainloopSm100 { for (; tile_scheduler.is_valid(); tile_scheduler.advance()) { int tid = tile_scheduler.get_current_tile_id(); - auto blk_coord = TileScheduler::decode_tile_coord(tid, params.h, chunk_indices_ptr, cu_seqlens_ptr); + // LoadAux: beta is per v-head (row stride = h_v). + auto blk_coord = TileScheduler::decode_tile_coord( + tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); int head_idx = get<1>(blk_coord); int tile_idx = get<2>(blk_coord); @@ -1024,7 +1043,7 @@ struct KdaChunkFwdRecompWUMainloopSm100 { float beta_val = (thread_idx < sub_seq_len) ? float(reinterpret_cast( - params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h + head_idx]) + params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h_v + head_idx]) : float(0); shared_plan->beta_smem[beta_pipe_state_write.index()][thread_idx] = beta_val; } diff --git a/csrc/kda/sm100/tile_scheduler.hpp b/csrc/kda/sm100/tile_scheduler.hpp index 47044aa5..695bb26c 100644 --- a/csrc/kda/sm100/tile_scheduler.hpp +++ b/csrc/kda/sm100/tile_scheduler.hpp @@ -26,11 +26,20 @@ // No smem synchronization needed — every CTA processes tiles starting // at blockIdx.x and striding by gridDim.x. All warps within a CTA // independently maintain the same tile_id, so no tile pipeline is needed. +// +// GVA (Grouped V-head Attention) support: +// Q/K are sized by `num_qk_heads`; V, g, beta, O and state tensors are +// sized by `num_v_heads`. We enumerate tiles by `num_v_heads` so that +// each v-head is scheduled independently, and derive the companion +// `qk_head_idx = v_head_idx / heads_per_group` on the device side. +// `heads_per_group = num_v_heads / num_qk_heads` is precomputed on the +// host to avoid a per-tile integer division. // =================================================================== struct StaticPersistentTileScheduler { struct Params { - int num_blocks; // number of sequence chunks (from chunk_indices) - int num_heads; + int num_blocks; // number of sequence chunks (from chunk_indices) + int num_heads; // == num_v_heads; tiles are enumerated by v-head + int heads_per_group; // == num_v_heads / num_qk_heads, precomputed on host int num_sm; int* tile_counter; // unused @@ -77,14 +86,22 @@ struct StaticPersistentTileScheduler { return current_tile_id < total_tiles(); } + // Decode tile_id -> (batch_idx, v_head_idx, seq_idx, qk_head_idx). + // `num_v_heads` is the number of V/O/g/beta heads; tile enumeration is + // done in v-head space. `heads_per_group` (= num_v_heads/num_qk_heads) + // is used to derive the companion Q/K head index for GVA. + // For backward compatibility, when HV == HQK, `heads_per_group == 1` + // and `qk_head_idx == v_head_idx`. CUTLASS_DEVICE static auto - decode_tile_coord(int tile_id, int num_heads, int* chunk_indices_ptr, int* cu_seqlens_ptr) { + decode_tile_coord( + int tile_id, int num_v_heads, int heads_per_group, int* chunk_indices_ptr, int* /*cu_seqlens_ptr*/) { using namespace cute; - int tile_idx_raw = tile_id / num_heads; - int head_idx = tile_id % num_heads; + int tile_idx_raw = tile_id / num_v_heads; + int v_head_idx = tile_id % num_v_heads; + int qk_head_idx = v_head_idx / heads_per_group; int batch_idx = chunk_indices_ptr[tile_idx_raw * 2]; int seq_idx = chunk_indices_ptr[tile_idx_raw * 2 + 1]; - return make_coord(batch_idx, head_idx, seq_idx, 0); + return make_coord(batch_idx, v_head_idx, seq_idx, qk_head_idx); } }; \ No newline at end of file diff --git a/cula/kda/chunk_intra.py b/cula/kda/chunk_intra.py index 07036383..aeb063fa 100644 --- a/cula/kda/chunk_intra.py +++ b/cula/kda/chunk_intra.py @@ -759,7 +759,12 @@ def chunk_kda_fwd_intra( unified_gref: bool = False, # Set True for ~5% extra perf (slightly lower precision) ): assert safe_gate, "Only safe_gate=True is supported in chunk_kda_fwd_intra for now" - B, T, H, K = k.shape + B, T, H_QK, K = k.shape + # GVA: g/beta/v live in h_v head space; q/k live in h_qk head space. + H_V = v.size(2) + assert H_QK > 0 and H_V > 0 and H_V % H_QK == 0, ( + f"HV ({H_V}) must be a positive multiple of HQK ({H_QK})" + ) BT = chunk_size if cu_seqlens is None: @@ -773,18 +778,20 @@ def chunk_kda_fwd_intra( "cu_seqlens and chunk_indices must be int32 for cuda impl" ) - Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) - Akk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) + # Aqk and Akk are produced per v-head by the intra kernel. + Aqk = torch.empty(B, T, H_V, BT, device=k.device, dtype=k.dtype) + Akk = torch.empty(B, T, H_V, BT, device=k.device, dtype=k.dtype) tile_counter = torch.zeros(1, dtype=torch.int32, device=q.device) cula_cuda.chunk_kda_fwd_intra_cuda( q, k, gk, beta, cu_seqlens, chunk_indices, Aqk, Akk, tile_counter, scale, chunk_size, use_tf32_inverse, unified_gref ) - w = torch.empty_like(k) + # w, u, kg, qg all live in h_v head space. + w = torch.empty_like(v) u = torch.empty_like(v) - qg = torch.empty_like(q) if disable_recompute else None - kg = torch.empty_like(k) if gk is not None else None + qg = torch.empty(B, T, H_V, K, device=q.device, dtype=q.dtype) if disable_recompute else None + kg = torch.empty(B, T, H_V, K, device=k.device, dtype=k.dtype) if gk is not None else None cula_cuda.recompute_w_u_cuda( k, v, beta, Akk, gk, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q if disable_recompute else None, qg @@ -857,4 +864,4 @@ def chunk_kda_bwd_intra( db = db2.sum(0).add_(db) dg = dg2 - return dq, dk, db, dg + return dq, dk, db, dg \ No newline at end of file diff --git a/tests/test_kda_gva_intra_sm100.py b/tests/test_kda_gva_intra_sm100.py new file mode 100644 index 00000000..8082946e --- /dev/null +++ b/tests/test_kda_gva_intra_sm100.py @@ -0,0 +1,386 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for SM100 KDA GVA (HV > HQK) support in chunk_kda_fwd_intra. + +The SM100 kernels (kda_fwd_intra / kda_fwd_recomp_w_u) now accept: + * q, k with head-dim ``HQK`` + * v, g, beta with head-dim ``HV`` where ``HV = group_size * HQK`` (group_size >= 1) + +This file verifies that the cuLA GVA path produces numerically matching results +compared to the FLA Triton reference, where the FLA reference does not natively +support GVA and therefore receives ``k`` replicated along the head axis to +``HV`` heads. Both uniform-length and varlen layouts are covered, and an +additional degeneracy test asserts that ``HV == HQK`` (group_size == 1) keeps +the non-GVA behaviour untouched. +""" + +from __future__ import annotations + +import pytest +import torch +from einops import rearrange +from fla.modules.l2norm import l2norm_fwd +from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import assert_close, device + +from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra +from cula.utils import prepare_uniform_cu_seqlens + +pytestmark = pytest.mark.sm100_only + + +# ========================================================================= +# Helpers +# ========================================================================= + +def _repeat_head(x: torch.Tensor, group_size: int, head_dim: int = 2) -> torch.Tensor: + """Replicate ``x`` along the head axis by ``group_size``. + + Mirrors GVA's broadcasting semantics: each QK head is paired with + ``group_size`` consecutive V heads, so ``k[..., h_qk, :]`` is used by + ``v[..., h_qk * group_size : (h_qk + 1) * group_size, :]``. + """ + return x.repeat_interleave(group_size, dim=head_dim).contiguous() + + +def _make_gva_inputs( + B: int, + T: int, + HQK: int, + HV: int, + D: int, + chunk_size: int, + cu_seqlens: torch.Tensor | None = None, + dtype: torch.dtype = torch.bfloat16, + seed: int = 42, +): + """Construct inputs for chunk_kda_fwd_intra in GVA layout. + + Returns: + q, k : (B, T, HQK, D) dtype + v : (B, T, HV, D) dtype + g : (B, T, HV, D) float32, after kda_gate_chunk_cumsum + beta : (B, T, HV) float32 in (0, 1) + scale : float + cu_seqlens : (N+1,) int32 or None + chunk_indices: (NT, 2) int32 or None + """ + assert HV % HQK == 0 and HV >= HQK, f"invalid HV/HQK: {HV}/{HQK}" + + torch.manual_seed(seed) + scale = D ** (-0.5) + + # QK are in HQK head space; V / gates / beta live in HV space. + q = torch.randn(B, T, HQK, D, dtype=dtype, device=device) + k = torch.randn(B, T, HQK, D, dtype=dtype, device=device) + v = torch.randn(B, T, HV, D, dtype=dtype, device=device) + g_raw = torch.randn(B, T, HV, D, dtype=dtype, device=device) + beta = torch.randn(B, T, HV, dtype=torch.float, device=device).sigmoid() + + # l2-normalise q/k so that scale/gate ranges match production use. + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + + # FLA gate cumsum only supports packed batch (B=1) when cu_seqlens is set. + if B != 1: + q, k, v, g_raw, beta = map( + lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), + (q, k, v, g_raw, beta), + ) + + # Per-HV gate preprocessing (cumsum inside chunks). + A_log = torch.randn(HV, dtype=torch.float, device=device) + dt_bias = torch.randn(HV * D, dtype=torch.float, device=device) + + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + ) + g = kda_gate_chunk_cumsum( + g=g_raw, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=-5.0, + ) + return q, k, v, g, beta, scale, cu_seqlens, chunk_indices + + +def _run_fla_ref(q, k_hqk, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, group_size, disable_recompute): + """Reference: replicate k along head axis to HV, then call FLA intra. + + FLA's chunk_kda_fwd_intra assumes H == HQK == HV (no GVA), so we construct + the HV-head view of k and q before invoking it. + """ + k_hv = _repeat_head(k_hqk, group_size) + q_hv = _repeat_head(q, group_size) + return fla_chunk_kda_fwd_intra( + q=q_hv, + k=k_hv, + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=True, + disable_recompute=disable_recompute, + ) + + +def _run_cula_gva(q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute): + return cula_chunk_kda_fwd_intra( + q=q, + k=k, + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=True, + disable_recompute=disable_recompute, + ) + + +def _assert_intra_outputs_match(ref, tri, disable_recompute: bool) -> None: + """Compare cuLA vs FLA on user-visible intra outputs. + + We intentionally skip ``Aqk``: the cuLA SM100 fused kernel does not + materialise every off-diagonal slot that FLA's multi-kernel path writes, + and the FLA reference can contain NaNs in unused ``Aqk`` entries. The + downstream tensors ``w`` / ``u`` / ``kg`` (and ``Akk``) are the meaningful + correctness signals and match the benchmark's comparison strategy. + """ + w_r, u_r, qg_r, kg_r, _Aqk_r, Akk_r = ref + w_c, u_c, qg_c, kg_c, _Aqk_c, Akk_c = tri + + assert Akk_c.shape == Akk_r.shape, (Akk_c.shape, Akk_r.shape) + assert w_c.shape == w_r.shape, (w_c.shape, w_r.shape) + assert u_c.shape == u_r.shape, (u_c.shape, u_r.shape) + assert kg_c.shape == kg_r.shape, (kg_c.shape, kg_r.shape) + + assert_close("Akk", Akk_r, Akk_c, 0.008) + assert_close("w", w_r, w_c, 0.008) + assert_close("u", u_r, u_c, 0.008) + assert_close("kg", kg_r, kg_c, 0.005) + + if disable_recompute: + assert qg_c is not None and qg_r is not None + assert qg_c.shape == qg_r.shape, (qg_c.shape, qg_r.shape) + assert_close("qg", qg_r, qg_c, 0.005) + else: + assert qg_c is None, "cuLA must not materialise qg when disable_recompute=False" + + +# ========================================================================= +# Uniform-length tests +# ========================================================================= + +@pytest.mark.parametrize("disable_recompute", [False, True], ids=["recomp", "no_recomp"]) +@pytest.mark.parametrize( + ("B", "T", "HQK", "group_size", "D"), + [ + pytest.param(*cfg, id="B{}-T{}-HQK{}-gs{}-D{}".format(*cfg)) + for cfg in [ + # group_size == 2: classic GVA 2:1 + (1, 256, 2, 2, 128), + (2, 512, 4, 2, 128), + # group_size == 4: wider grouping + (1, 1024, 2, 4, 128), + (2, 1024, 4, 4, 128), + # Non-multiple-of-BT sequence length to stress boundary handling. + (1, 500, 2, 2, 128), + (1, 1000, 4, 2, 128), + ] + ], +) +def test_gva_intra_uniform(B, T, HQK, group_size, D, disable_recompute): + """cuLA GVA path must match FLA(k-replicated-to-HV) for uniform seqlens.""" + HV = HQK * group_size + chunk_size = 64 + + cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) + q, k, v, g, beta, scale, cu_seqlens, chunk_indices = _make_gva_inputs( + B=B, T=T, HQK=HQK, HV=HV, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens, + ) + + # cuLA GVA path (k in HQK head space). + w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c = _run_cula_gva( + q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute, + ) + + # FLA reference (k replicated to HV). + w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r = _run_fla_ref( + q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, group_size, disable_recompute, + ) + + _assert_intra_outputs_match( + (w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r), + (w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c), + disable_recompute, + ) + + +# ========================================================================= +# Varlen tests +# ========================================================================= + +@pytest.mark.parametrize("disable_recompute", [False, True], ids=["recomp", "no_recomp"]) +@pytest.mark.parametrize( + ("HQK", "group_size", "D", "cu_seqlens"), + [ + pytest.param(*cfg, id="HQK{}-gs{}-D{}-ns{}".format(cfg[0], cfg[1], cfg[2], len(cfg[3]) - 1)) + for cfg in [ + (2, 2, 128, [0, 256, 500, 1000]), + (4, 2, 128, [0, 100, 300, 1200, 2000]), + (2, 4, 128, [0, 15, 100, 300, 1200, 2048]), + # Simulated realistic trace. + ( + 4, 2, 128, + [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096], + ), + ] + ], +) +def test_gva_intra_varlen(HQK, group_size, D, cu_seqlens, disable_recompute): + """GVA correctness under variable-length (packed) inputs.""" + HV = HQK * group_size + chunk_size = 64 + + cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + T = int(cu_seqlens_t[-1].item()) + # Packed layout uses B=1 and a flat time axis. + q, k, v, g, beta, scale, cu_seqlens_t, chunk_indices = _make_gva_inputs( + B=1, T=T, HQK=HQK, HV=HV, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens_t, + ) + + w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c = _run_cula_gva( + q, k, v, g, beta, scale, cu_seqlens_t, chunk_indices, chunk_size, disable_recompute, + ) + w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r = _run_fla_ref( + q, k, v, g, beta, scale, cu_seqlens_t, chunk_indices, chunk_size, group_size, disable_recompute, + ) + + _assert_intra_outputs_match( + (w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r), + (w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c), + disable_recompute, + ) + + +# ========================================================================= +# Degeneracy: HV == HQK must match the non-GVA (same-shape) reference +# ========================================================================= + +@pytest.mark.parametrize("disable_recompute", [False, True], ids=["recomp", "no_recomp"]) +@pytest.mark.parametrize( + ("B", "T", "H", "D"), + [ + pytest.param(*cfg, id="B{}-T{}-H{}-D{}".format(*cfg)) + for cfg in [ + (1, 512, 4, 128), + (2, 1024, 4, 128), + ] + ], +) +def test_gva_intra_degenerate_equals_non_gva(B, T, H, D, disable_recompute): + """When HV == HQK, the GVA code path must be byte-for-byte equivalent + to the non-GVA path that existed before this change. + + We do not have a separate "non-GVA" entrypoint, but we can assert the + cuLA path matches FLA with *no* head replication (group_size=1), which + exercises the ``HV == HQK`` fast-path inside the new kernels. + """ + chunk_size = 64 + cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) + q, k, v, g, beta, scale, cu_seqlens, chunk_indices = _make_gva_inputs( + B=B, T=T, HQK=H, HV=H, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens, + ) + + w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c = _run_cula_gva( + q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute, + ) + # group_size=1 → no replication; identical input shape to cuLA. + w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r = fla_chunk_kda_fwd_intra( + q=q, k=k, v=v, gk=g, beta=beta, scale=scale, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=True, disable_recompute=disable_recompute, + ) + + _assert_intra_outputs_match( + (w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r), + (w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c), + disable_recompute, + ) + + +# ========================================================================= +# Shape / contract sanity checks (run even without a reference) +# ========================================================================= + +@pytest.mark.parametrize("group_size", [1, 2, 4]) +def test_gva_intra_output_shapes(group_size): + """All outputs of chunk_kda_fwd_intra must live in HV-head space.""" + B, T, HQK, D = 1, 256, 2, 128 + HV = HQK * group_size + chunk_size = 64 + cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) + q, k, v, g, beta, scale, cu_seqlens, chunk_indices = _make_gva_inputs( + B=B, T=T, HQK=HQK, HV=HV, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens, + ) + w, u, qg, kg, Aqk, Akk = _run_cula_gva( + q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute=True, + ) + + assert Aqk.shape == (B, T, HV, chunk_size), Aqk.shape + assert Akk.shape == (B, T, HV, chunk_size), Akk.shape + assert w.shape == (B, T, HV, D), w.shape + assert u.shape == (B, T, HV, D), u.shape + assert kg.shape == (B, T, HV, D), kg.shape + assert qg is not None and qg.shape == (B, T, HV, D), (None if qg is None else qg.shape) + + +# ========================================================================= +# Negative / assertion tests +# ========================================================================= + +def test_gva_intra_rejects_non_multiple_ratio(): + """HV must be a positive integer multiple of HQK.""" + B, T, HQK, HV, D = 1, 128, 3, 5, 128 # 5 % 3 != 0 + chunk_size = 64 + cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) + # We intentionally do not use _make_gva_inputs because the assert fires + # before kernel launch on the python side. + dtype = torch.bfloat16 + q = torch.randn(B, T, HQK, D, dtype=dtype, device=device) + k = torch.randn(B, T, HQK, D, dtype=dtype, device=device) + v = torch.randn(B, T, HV, D, dtype=dtype, device=device) + g = torch.randn(B, T, HV, D, dtype=torch.float, device=device) + beta = torch.randn(B, T, HV, dtype=torch.float, device=device).sigmoid() + + with pytest.raises((AssertionError, RuntimeError), match=r"multiple|h_v"): + cula_chunk_kda_fwd_intra( + q=q, k=k, v=v, gk=g, beta=beta, scale=D ** -0.5, + cu_seqlens=cu_seqlens, chunk_size=chunk_size, + safe_gate=True, disable_recompute=False, + ) From 7960942ce1d18bd214805a32f36b975d5e03900b Mon Sep 17 00:00:00 2001 From: Emre-Dinc <80748100+Emre-Dinc@users.noreply.github.com> Date: Sun, 24 May 2026 02:39:55 +0300 Subject: [PATCH 16/34] fix: wire h0_indices into Lightning Attention decode for state-pool indexing (#75) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 无言独上机房 <88866917+sjmshsh@users.noreply.github.com> --- benchmarks/bench_la_decode_vs_fla.py | 2 +- cula/ops/la_decode.py | 29 +-- tests/test_la_decode_pool.py | 263 +++++++++++++++++++++++++++ 3 files changed, 282 insertions(+), 12 deletions(-) create mode 100644 tests/test_la_decode_pool.py diff --git a/benchmarks/bench_la_decode_vs_fla.py b/benchmarks/bench_la_decode_vs_fla.py index 471bc226..18a0b07d 100644 --- a/benchmarks/bench_la_decode_vs_fla.py +++ b/benchmarks/bench_la_decode_vs_fla.py @@ -201,7 +201,7 @@ def kernel_fla(): # cute kernel: pre-create compiled + stream handle cute_state_k = state_init.clone().permute(0, 1, 3, 2).reshape(B * H, V, K).contiguous() out_cute_k = torch.empty(B, H, V, device=device, dtype=dtype) - cache = _get_compiled_kernel(B, 1, H, K, V, scale, USE_FAST_MATH) + cache = _get_compiled_kernel(B, 1, H, K, V, cute_state_k.shape[0], scale, USE_FAST_MATH) compiled_cute = cache["compiled"] stream_handle = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/cula/ops/la_decode.py b/cula/ops/la_decode.py index bd6f1b8c..08831dfb 100644 --- a/cula/ops/la_decode.py +++ b/cula/ops/la_decode.py @@ -127,8 +127,9 @@ def la_decode_kernel_small_batch_pretranspose( cute.arch.barrier() # Get current batch - gSrc_batch = h0_source[(batch_idx, None, None)] # (V, K) - gDst = cute.local_tile(h0_source, (1, TILE_V, TILE_K), (batch_idx, None, 0)) + pool_idx = h0_indices[i_n] * HV + i_hv + gSrc_batch = h0_source[(pool_idx, None, None)] # (V, K) + gDst = cute.local_tile(h0_source, (1, TILE_V, TILE_K), (pool_idx, None, 0)) # split tiles in V-dimension gSrc = cute.local_tile(gSrc_batch, (TILE_V, TILE_K), (None, 0)) # (TILE_V, TILE_K, num_v_tiles) @@ -289,8 +290,9 @@ def la_decode_kernel_big_batch_pretranspose( cute.arch.barrier() # Get current batch - gSrc_batch = h0_source[(batch_idx, None, None)] # (V, K) - gDst = cute.local_tile(h0_source, (1, TILE_V, TILE_K), (batch_idx, None, 0)) + pool_idx = h0_indices[i_n] * HV + i_hv + gSrc_batch = h0_source[(pool_idx, None, None)] # (V, K) + gDst = cute.local_tile(h0_source, (1, TILE_V, TILE_K), (pool_idx, None, 0)) # split tiles in V-dimension gSrc = cute.local_tile(gSrc_batch, (TILE_V, TILE_K), (None, 0)) # (TILE_V, TILE_K, num_v_tiles) @@ -418,7 +420,7 @@ def run_la_decode_kernel_big_batch_pretranspose( stream: cuda.CUstream, ): # h0_source: (B*HV, V, K) - batch_size, v_dim, _k_dim = ( + _pool_dim0, v_dim, _k_dim = ( h0_source.layout.shape[0], h0_source.layout.shape[1], h0_source.layout.shape[2], @@ -477,7 +479,7 @@ def run_la_decode_kernel_big_batch_pretranspose( TILE_V_BIG, NUM_STAGES_BIG, ).launch( - grid=(batch_size, 1, 1), + grid=(B * H, 1, 1), block=[NUM_THREADS_BIG, 1, 1], smem=smem_bytes, stream=stream, @@ -502,7 +504,7 @@ def run_la_decode_kernel_small_batch_pretranspose( stream: cuda.CUstream, ): # h0_source: (B*H, V, K) - batch_size, v_dim, _k_dim = ( + _pool_dim0, v_dim, _k_dim = ( h0_source.layout.shape[0], h0_source.layout.shape[1], h0_source.layout.shape[2], @@ -561,7 +563,7 @@ def run_la_decode_kernel_small_batch_pretranspose( TILE_V_SMALL, NUM_STAGES_SMALL, ).launch( - grid=(batch_size * NUM_BLOCKS_PER_STATE, 1, 1), + grid=(B * H * NUM_BLOCKS_PER_STATE, 1, 1), block=[NUM_THREADS_SMALL, 1, 1], smem=smem_bytes, stream=stream, @@ -569,7 +571,9 @@ def run_la_decode_kernel_small_batch_pretranspose( @functools.cache -def _get_compiled_kernel(B: int, T: int, H: int, K: int, V: int, softmax_scale: float, use_fast_math: bool = True): +def _get_compiled_kernel( + B: int, T: int, H: int, K: int, V: int, pool_dim0: int, softmax_scale: float, use_fast_math: bool = True +): """Get or create compiled kernel cache.""" return {} @@ -625,10 +629,14 @@ def linear_attention_decode( raise NotImplementedError(f"CuTe kernel doesn't support K splitting (k_dim_block={k_dim_block})") # Get compiled kernel (cached) - cache_key = (B, 1, H, HEAD_DIM, HEAD_DIM, softmax_scale, USE_FAST_MATH) + pool_dim0 = s.shape[0] + cache_key = (B, 1, H, HEAD_DIM, HEAD_DIM, pool_dim0, softmax_scale, USE_FAST_MATH) cache = _get_compiled_kernel(*cache_key) h0_source = s + + # Validate state pool dimensions + assert s.shape[0] % H == 0, f"s.shape[0] must be divisible by H={H}, got {s.shape[0]}" # First-time compilation if "compiled" not in cache: stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -644,7 +652,6 @@ def linear_attention_decode( v_view = v o_view = out - # Use s_offsets directly (pass to kernel but not actually used in current implementation) h0_indices = s_offsets # Convert to CuTe format for compilation diff --git a/tests/test_la_decode_pool.py b/tests/test_la_decode_pool.py new file mode 100644 index 00000000..c9a75796 --- /dev/null +++ b/tests/test_la_decode_pool.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Test for Lightning Attention decode state-pool indirect indexing. + +Exposes the bug where la_decode ignores s_offsets and indexes state +by flattened batch_idx directly. With identity offsets the bug is invisible. +With non-identity offsets, the kernel reads/writes wrong state slots. +""" + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.ops.la_decode import linear_attention_decode + + +def torch_la_decode_ref(q, k, v, state, decay_scales, scale): + """Pure PyTorch reference — state is [B, H, K, V] (BHKV).""" + B, H, D = q.shape + q_f = q.float() * scale + k_f = k.float() + v_f = v.float() + decay = torch.exp(-decay_scales).view(1, H, 1, 1) + state_new = state * decay + k_f.unsqueeze(-1) * v_f.unsqueeze(-2) + o = torch.einsum("bhk,bhkv->bhv", q_f, state_new) + return o.to(torch.bfloat16), state_new + + +def run_la_decode_with_pool(q, k, v, state_pool_4d, s_offsets, decay_scales, scale): + """ + Run la_decode with a state pool and arbitrary offsets. + + state_pool_4d: [pool_size, H, K, V] — the full pool (BHKV layout) + s_offsets: [B] — which pool slot each batch element uses + """ + B, H, D = q.shape + pool_size = state_pool_4d.shape[0] + + # la_decode expects BHVK layout: [pool_size*H, V, K] + state_cute = state_pool_4d.clone().transpose(-1, -2).contiguous().reshape(pool_size * H, D, D) + out = torch.zeros(B, H, D, device=q.device, dtype=torch.bfloat16) + + linear_attention_decode( + q, + k, + v, + state_cute, + out, + softmax_scale=scale, + stride_q=0, + stride_k=0, + stride_v=0, + stride_s=0, + stride_o=0, + s_offsets=s_offsets, + decay_scales=decay_scales, + HEAD_DIM=D, + K_SPLIT_DIM=D, + V_SPLIT_DIM=D, + ) + + state_out = state_cute.reshape(pool_size, H, D, D).transpose(-1, -2).contiguous() + return out, state_out + + +# --------------------------------------------------------------------------- +# Test 1: Identity offsets (baseline — should always pass) +# --------------------------------------------------------------------------- +def test_identity_offsets(): + """Identity offsets: s_offsets=[0,1,2,3]. Bug is invisible.""" + B, H, D = 4, 8, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + torch.manual_seed(42) + q = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + state_4d = torch.randn(B, H, D, D, device="cuda", dtype=torch.float32) * 0.1 + + s_offsets = torch.arange(B, device="cuda", dtype=torch.int32) + out, _ = run_la_decode_with_pool(q, k, v, state_4d, s_offsets, decay_scales, scale) + + o_ref, _ = torch_la_decode_ref(q, k, v, state_4d, decay_scales, scale) + rmse = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + rel_err = rmse / (max_ref + 1e-8) + + assert rel_err < 0.01, f"Identity offsets: rel_err={rel_err:.6f}" + + +# --------------------------------------------------------------------------- +# Test 2: Non-identity offsets (exposes the bug) +# --------------------------------------------------------------------------- +def test_non_identity_offsets(): + """ + pool_size=6, batch=4, offsets=[2, 0, 5, 1]. + Each batch reads a different, non-sequential pool slot. + Bug: kernel reads slots [0,1,2,3] instead of [2,0,5,1]. + """ + B = 4 + POOL_SIZE = 6 + H, D = 8, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + torch.manual_seed(42) + q = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + + # Large state magnitude so wrong-slot reads produce clearly different outputs + state_pool = torch.randn(POOL_SIZE, H, D, D, device="cuda", dtype=torch.float32) * 0.1 + + offsets = [2, 0, 5, 1] + s_offsets = torch.tensor(offsets, device="cuda", dtype=torch.int32) + + out, _ = run_la_decode_with_pool(q, k, v, state_pool, s_offsets, decay_scales, scale) + + # Reference: manually select the correct state for each batch element + state_selected = state_pool[s_offsets.long()] # [B, H, D, D] + o_ref, _ = torch_la_decode_ref(q, k, v, state_selected, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + rel_err = rmse / (max_ref + 1e-8) + + assert rel_err < 0.01, f"Non-identity offsets {offsets}: rel_err={rel_err:.6f}" + + +# --------------------------------------------------------------------------- +# Test 3: Reversed offsets (another non-identity pattern) +# --------------------------------------------------------------------------- +def test_reversed_offsets(): + """ + pool_size=B, offsets=[3,2,1,0] (reversed). + Batch 0 reads slot 3, batch 3 reads slot 0. + """ + B, H, D = 4, 8, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + torch.manual_seed(42) + q = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + state_pool = torch.randn(B, H, D, D, device="cuda", dtype=torch.float32) * 0.1 + + offsets = list(reversed(range(B))) + s_offsets = torch.tensor(offsets, device="cuda", dtype=torch.int32) + + out, _ = run_la_decode_with_pool(q, k, v, state_pool, s_offsets, decay_scales, scale) + + state_selected = state_pool[s_offsets.long()] + o_ref, _ = torch_la_decode_ref(q, k, v, state_selected, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + rel_err = rmse / (max_ref + 1e-8) + + assert rel_err < 0.01, f"Reversed offsets {offsets}: rel_err={rel_err:.6f}" + + +# --------------------------------------------------------------------------- +# Test 4: State writeback with non-identity offsets +# --------------------------------------------------------------------------- +def test_state_writeback_non_identity(): + """ + Verify that state updates go to the correct pool slots. + After decode, pool slot offsets[i] should have the updated state for batch i. + Other pool slots should be unchanged. + """ + B = 4 + POOL_SIZE = 6 + H, D = 8, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + torch.manual_seed(42) + q = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + state_pool = torch.randn(POOL_SIZE, H, D, D, device="cuda", dtype=torch.float32) * 0.1 + state_pool_orig = state_pool.clone() + + offsets = [2, 0, 5, 1] + s_offsets = torch.tensor(offsets, device="cuda", dtype=torch.int32) + + _, state_out = run_la_decode_with_pool(q, k, v, state_pool, s_offsets, decay_scales, scale) + + # Reference: compute expected new state for each active batch element + state_selected = state_pool_orig[s_offsets.long()] + _, state_ref = torch_la_decode_ref(q, k, v, state_selected, decay_scales, scale) + + # Check that active slots were updated correctly + for b_idx, pool_slot in enumerate(offsets): + slot_rmse = torch.sqrt(torch.mean((state_out[pool_slot].float() - state_ref[b_idx].float()) ** 2)).item() + slot_max = torch.abs(state_ref[b_idx].float()).max().item() + slot_rel = slot_rmse / (slot_max + 1e-8) + assert slot_rel < 0.001, f"State writeback: pool slot {pool_slot} (batch {b_idx}) rel_err={slot_rel:.6f}" + + # Check that inactive slots (3, 4) were NOT touched + inactive = set(range(POOL_SIZE)) - set(offsets) + for slot in inactive: + diff = torch.abs(state_out[slot] - state_pool_orig[slot]).max().item() + assert diff < 1e-8, f"Inactive pool slot {slot} was modified! max_diff={diff}" + + +# --------------------------------------------------------------------------- +# Test 5: Big batch (B > 32) with non-identity offsets +# --------------------------------------------------------------------------- +def test_big_batch_non_identity_offsets(): + """ + B=33 triggers the big-batch kernel path (B > 32). + pool_size=40, shifted offsets so batch i reads slot (i + 7) % 40. + """ + B = 33 + POOL_SIZE = 40 + H, D = 8, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + torch.manual_seed(42) + q = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + + state_pool = torch.randn(POOL_SIZE, H, D, D, device="cuda", dtype=torch.float32) * 0.1 + + offsets = [(i + 7) % POOL_SIZE for i in range(B)] + s_offsets = torch.tensor(offsets, device="cuda", dtype=torch.int32) + + out, _ = run_la_decode_with_pool(q, k, v, state_pool, s_offsets, decay_scales, scale) + + state_selected = state_pool[s_offsets.long()] + o_ref, _ = torch_la_decode_ref(q, k, v, state_selected, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + rel_err = rmse / (max_ref + 1e-8) + + assert rel_err < 0.01, f"Big batch non-identity offsets: rel_err={rel_err:.6f}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--tb=short"]) From e871b966a2d0e1fe0aa9c88a724df145d400a1fe Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Sun, 24 May 2026 09:26:42 +0800 Subject: [PATCH 17/34] Unify the naming convention (#79) * chore: rename sm100-specific cutedsl kernels * chore: revert cula/kda naming changes * chore: drop debugging codes * chore: fix lint errors --- REPO_LAYOUT.md | 10 +- benchmarks/bench_chunk_delta_h.py | 2 +- benchmarks/bench_fwd_o.py | 2 +- benchmarks/bench_kda_fused_fwd.py | 2 +- benchmarks/bench_lightning_attn.py | 2 +- benchmarks/bench_linear_attn.py | 2 +- cula/__init__.py | 2 +- cula/kda/__init__.py | 2 + cula/kda/blackwell_fused_fwd.py | 74 +---------- cula/kda/chunk_bwd.py | 2 +- cula/kda/chunk_fwd.py | 4 +- cula/lightning/__init__.py | 2 +- ...hunk_delta_h.py => chunk_delta_h_sm100.py} | 0 cula/ops/{fwd_o.py => fwd_o_sm100.py} | 0 ...ed_wip.py => kda_fully_fused_sm100_wip.py} | 117 ------------------ ...htning_attn.py => lightning_attn_sm100.py} | 0 .../{linear_attn.py => linear_attn_sm100.py} | 2 +- cula/utils.py | 10 +- pyproject.toml | 2 +- tests/test_chunk_delta_h.py | 2 +- tests/test_compare_with_fla.py | 2 +- tests/test_fwd_o.py | 2 +- tests/test_la_decode.py | 2 +- tests/test_lightning_attn.py | 2 +- 24 files changed, 29 insertions(+), 218 deletions(-) rename cula/ops/{chunk_delta_h.py => chunk_delta_h_sm100.py} (100%) rename cula/ops/{fwd_o.py => fwd_o_sm100.py} (100%) rename cula/ops/{kda_fully_fused_wip.py => kda_fully_fused_sm100_wip.py} (97%) rename cula/ops/{lightning_attn.py => lightning_attn_sm100.py} (100%) rename cula/ops/{linear_attn.py => linear_attn_sm100.py} (99%) diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 26109017..6a9d5866 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -12,11 +12,11 @@ cuLA/ │ ├── lightning/ # Lightning Attention operators │ │ └── la_decode.py # Single-token decode kernel (CuTe DSL) │ ├── ops/ # CuTe DSL kernel implementations -│ │ ├── chunk_delta_h.py # Chunk delta-H kernel -│ │ ├── fwd_o.py # Forward output kernel -│ │ ├── lightning_attn.py # Lightning Attention prefill kernel -│ │ ├── linear_attn.py # Generic linear attention kernel -│ │ ├── kda_fully_fused_wip.py # WIP fully fused KDA kernel +│ │ ├── chunk_delta_h_sm100.py # Chunk delta-H kernel (SM100) +│ │ ├── fwd_o_sm100.py # Forward output kernel (SM100) +│ │ ├── lightning_attn_sm100.py # Lightning Attention prefill kernel (SM100) +│ │ ├── linear_attn_sm100.py # Generic linear attention kernel (SM100) +│ │ ├── kda_fully_fused_sm100_wip.py # WIP fully fused KDA kernel (SM100) │ └── utils.py # Shared utilities │ ├── csrc/ # CUDA C++ / CUTLASS kernels diff --git a/benchmarks/bench_chunk_delta_h.py b/benchmarks/bench_chunk_delta_h.py index f14c866e..09fde953 100644 --- a/benchmarks/bench_chunk_delta_h.py +++ b/benchmarks/bench_chunk_delta_h.py @@ -48,7 +48,7 @@ import torch # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── -_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h") +_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h # ─── FLA baseline imports ─── diff --git a/benchmarks/bench_fwd_o.py b/benchmarks/bench_fwd_o.py index 7aa40b05..6fcbca00 100644 --- a/benchmarks/bench_fwd_o.py +++ b/benchmarks/bench_fwd_o.py @@ -45,7 +45,7 @@ import torch # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── -_fwd_o_mod = importlib.import_module("cula.ops.fwd_o") +_fwd_o_mod = importlib.import_module("cula.ops.fwd_o_sm100") chunk_gla_fwd_o = _fwd_o_mod.chunk_gla_fwd_o build_chunk_indices = _fwd_o_mod.build_chunk_indices diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index 0b2dd538..3953ccca 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -18,7 +18,7 @@ Automatically selects the cuLA fully-fused implementation based on the current GPU architecture: - - sm100 (Blackwell) → cula.kda.blackwell_fused_fwd.flash_kda_prefill + - sm100 (Blackwell) → cula.kda.blackwell_fused_fwd.flash_kda_prefill - sm90 (Hopper) → cula.kda.hopper_fused_fwd.cula_kda_prefill Compares: diff --git a/benchmarks/bench_lightning_attn.py b/benchmarks/bench_lightning_attn.py index f4668a69..07733009 100644 --- a/benchmarks/bench_lightning_attn.py +++ b/benchmarks/bench_lightning_attn.py @@ -52,7 +52,7 @@ from fla.ops.simple_gla.chunk import chunk_simple_gla_fwd -from cula.ops.lightning_attn import lightning_attn_fwd, lightning_attn_fwd_varlen +from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen # ============================================================================= # Constants diff --git a/benchmarks/bench_linear_attn.py b/benchmarks/bench_linear_attn.py index 23049bf9..08bbb520 100644 --- a/benchmarks/bench_linear_attn.py +++ b/benchmarks/bench_linear_attn.py @@ -26,7 +26,7 @@ # from fla.ops.linear_attn.naive import naive_recurrent_linear_attn from fla.utils import assert_close, device -from cula.ops.linear_attn import LinearAttentionChunkwise +from cula.ops.linear_attn_sm100 import LinearAttentionChunkwise os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison diff --git a/cula/__init__.py b/cula/__init__.py index e6d63f27..7272e289 100644 --- a/cula/__init__.py +++ b/cula/__init__.py @@ -14,7 +14,7 @@ __version__ = "0.1.0" -from cula.ops.lightning_attn import LinearAttentionChunkwiseDecay +from cula.ops.lightning_attn_sm100 import LinearAttentionChunkwiseDecay __all__ = [ "LinearAttentionChunkwiseDecay", diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index 01190143..ee1a2bb9 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -12,12 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from cula.kda.blackwell_fused_fwd import flash_kda_prefill as kda_prefill_blackwell from cula.kda.chunk import chunk_kda from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode __all__ = [ "chunk_kda", + "kda_prefill_blackwell", "kda_decode", "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", diff --git a/cula/kda/blackwell_fused_fwd.py b/cula/kda/blackwell_fused_fwd.py index 291d9664..c0533d24 100644 --- a/cula/kda/blackwell_fused_fwd.py +++ b/cula/kda/blackwell_fused_fwd.py @@ -32,7 +32,7 @@ from fla.ops.utils.constant import RCP_LN2 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard -from cula.ops.kda_fully_fused_wip import KDAChunkwise +from cula.ops.kda_fully_fused_sm100_wip import KDAChunkwise from cula.utils import USE_FAST_MATH, assert_blackwell # Global kernel cache @@ -80,10 +80,6 @@ def forward( else: num_seqs = B - # No output padding needed — the kernel handles tail tiles via - # TMA descriptor modification (like flashkda), preventing writes - # into the next sequence's output region. - g_org = None if use_gate_in_kernel: try: @@ -110,7 +106,6 @@ def forward( A_log=A_log, dt_bias=dt_bias, ) - # only in safe_gate && use_gate_in_kernel, cumsum is fused into kda_gate_chunk_cumsum if not (safe_gate and use_gate_in_kernel): g = chunk_local_cumsum( g=g, chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices @@ -126,7 +121,6 @@ def forward( g_cute = from_dlpack(g.detach()) beta_cute = from_dlpack(beta.detach()) - # FIXME: support return final_states o = torch.empty_like(q) o_cute = from_dlpack(o.detach()) @@ -135,12 +129,10 @@ def forward( has_initial_state = initial_state is not None cache_key = (has_initial_state, output_final_state, safe_gate, is_varlen, scale, chunk_size, D, USE_FAST_MATH) - # Prepare cu_seqlens as int32 for kernel if is_varlen: cu_seqlens_i32 = cu_seqlens.to(torch.int32).contiguous() cu_seqlens_cute = from_dlpack(cu_seqlens_i32.detach()) else: - # Use cached dummy cu_seqlens to avoid per-call allocation overhead dev = q.device if dev not in _dummy_cache: _dummy_cu = torch.zeros(2, dtype=torch.int32, device=dev) @@ -155,9 +147,6 @@ def forward( cu_seqlens_i32 = dc["cu_seqlens"] cu_seqlens_cute = dc["cu_seqlens_cute"] - # Workspace buffer for TMA descriptor modification (varlen tail tiles) - # Same approach as flashkda: per-CTA slot for modified TMA descriptors - # 128 bytes per TMA descriptor, indexed by bidx (sequence index) dev = q.device if dev not in _dummy_cache: _dummy_cu = torch.zeros(2, dtype=torch.int32, device=dev) @@ -171,7 +160,6 @@ def forward( dc = _dummy_cache[dev] if is_varlen: ws_size = num_seqs * 128 - # Allocate/reuse workspace (grow if needed) if "workspace" not in dc or dc["workspace"].numel() < ws_size: ws_buf = torch.zeros(ws_size, dtype=torch.uint8, device=dev) dc["workspace"] = ws_buf @@ -184,13 +172,10 @@ def forward( dc["workspace_cute"] = from_dlpack(ws_buf.detach()) workspace_cute = dc["workspace_cute"] - # State shape: [num_seqs, H, D, D] - # Prepare initial_state and final_state tensors if has_initial_state: initial_state_f32 = initial_state.to(torch.float32).contiguous() initial_state_cute = from_dlpack(initial_state_f32.detach()) else: - # Use cached tiny dummy (pointer won't be dereferenced when has_initial_state=False) initial_state_f32 = None initial_state_cute = _dummy_cache[q.device]["state_cute"] @@ -198,11 +183,9 @@ def forward( final_state_f32 = torch.zeros(num_seqs, H, D, D, dtype=torch.float32, device=q.device) final_state_cute = from_dlpack(final_state_f32.detach()) else: - # Use cached tiny dummy (pointer won't be dereferenced when output_final_state=False) final_state_f32 = None final_state_cute = _dummy_cache[q.device]["state_cute"] - # problem_size: (num_seqs, total_tokens_or_seq_len, H, D) problem_size = (num_seqs, S, H, D) if cache_key in compiled_kernel_cache: @@ -270,7 +253,6 @@ def backward( raise NotImplementedError("Backward pass is not implemented yet.") -# TODO: Blackwell fused prefill is still under development @torch.compiler.disable def flash_kda_prefill( q: torch.Tensor, @@ -289,61 +271,8 @@ def flash_kda_prefill( chunk_indices: torch.IntTensor | None = None, **kwargs, ): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, H, K]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - scale (Optional[float]): - Scale factor for the KDA attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2norm to the q,k tensor internally. Default: `False`. - use_gate_in_kernel (bool): - Whether to compute the log-space KDA decay internally. - - If `True`: - The passed `g` acts as the raw input for `-exp(A_log).view(H, -1) * softplus(g + dt_bias.view(H, K))`. - Note that as part of the input arguments, - `A_log` (shape `[H]`) and the optional `dt_bias` (shape `[H * K]`) should be provided. - - If `False`, `g` is expected to be the pre-computed decay value. - Default: `False`. - safe_gate (bool): - Whether the kernel can assume the input gate values `g` are in a safe range. - When `True`, the kernel can use M=16 TensorCore acceleration. - The safe range is approximately [-5, 0). Default: `False`. - lower_bound (Optional[float]): - Lower bound for the forget gate activation function when `use_gate_in_kernel=True`. - This parameter modifies the internal forget gate activation and is recommended - to be set to `-5` when `safe_gate` is enabled. Default: `None`. - cu_seqlens (torch.IntTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - chunk_indices (torch.IntTensor): - Chunk indices used for variable-length training, - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - """ assert_blackwell() - # initial_state is now supported assert cu_seqlens is None or q.shape[0] == 1, "For varlen, batch size must be 1. Flatten sequences first." - # assert output_final_state == False, "output_final_state=True is not supported in cutedsl_kda_prefill yet." if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -355,7 +284,6 @@ def flash_kda_prefill( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - # Non-aligned sequence lengths are handled natively by the kernel if initial_state is not None: assert initial_state.dtype == torch.float32, "initial_state must be in float32." diff --git a/cula/kda/chunk_bwd.py b/cula/kda/chunk_bwd.py index 859b6be9..4dbd3317 100644 --- a/cula/kda/chunk_bwd.py +++ b/cula/kda/chunk_bwd.py @@ -39,7 +39,7 @@ from cula.kda.chunk_intra import chunk_kda_bwd_intra from cula.utils import prepare_uniform_cu_seqlens -_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h") +_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h BK_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/cula/kda/chunk_fwd.py b/cula/kda/chunk_fwd.py index 9fab2351..f49364bb 100644 --- a/cula/kda/chunk_fwd.py +++ b/cula/kda/chunk_fwd.py @@ -30,9 +30,9 @@ from cula.utils import assert_blackwell # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── -_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h") +_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h -_fwd_o_mod = importlib.import_module("cula.ops.fwd_o") +_fwd_o_mod = importlib.import_module("cula.ops.fwd_o_sm100") chunk_gla_fwd_o = _fwd_o_mod.chunk_gla_fwd_o diff --git a/cula/lightning/__init__.py b/cula/lightning/__init__.py index 973f7707..fb5e5635 100644 --- a/cula/lightning/__init__.py +++ b/cula/lightning/__init__.py @@ -13,7 +13,7 @@ # limitations under the License. from cula.ops.la_decode import linear_attention_decode -from cula.ops.lightning_attn import ( +from cula.ops.lightning_attn_sm100 import ( LinearAttentionChunkwiseDecay, lightning_attn_fwd, lightning_attn_fwd_varlen, diff --git a/cula/ops/chunk_delta_h.py b/cula/ops/chunk_delta_h_sm100.py similarity index 100% rename from cula/ops/chunk_delta_h.py rename to cula/ops/chunk_delta_h_sm100.py diff --git a/cula/ops/fwd_o.py b/cula/ops/fwd_o_sm100.py similarity index 100% rename from cula/ops/fwd_o.py rename to cula/ops/fwd_o_sm100.py diff --git a/cula/ops/kda_fully_fused_wip.py b/cula/ops/kda_fully_fused_sm100_wip.py similarity index 97% rename from cula/ops/kda_fully_fused_wip.py rename to cula/ops/kda_fully_fused_sm100_wip.py index adcf809b..ab09f0b2 100644 --- a/cula/ops/kda_fully_fused_wip.py +++ b/cula/ops/kda_fully_fused_sm100_wip.py @@ -282,9 +282,6 @@ def _plan_tmem_offsets( acc_shape_pv = tiled_mma_pv.partition_shape_C(tile_shape_mnk_pv[:2]) tCtAccPV_fake = tiled_mma_pv.make_fragment_C(cute.append(acc_shape_pv, acc_stages)) num_pv_acc_cols = tcgen05.find_tmem_tensor_col_offset(tCtAccPV_fake) - if cutlass.const_expr(PRINT_DEBUG): - print(f"tCtAccPV_fake={tCtAccPV_fake}, num_pv_acc_cols={num_pv_acc_cols}") - # No stage for linear state. acc_shape_kv = tiled_mma_kv.partition_shape_C(tile_shape_mnk_kv[:2]) tCtAccKV_fake = tiled_mma_kv.make_fragment_C(cute.append(acc_shape_kv, 1)) @@ -292,15 +289,11 @@ def _plan_tmem_offsets( # Cannot reuse KV since we need to accumulate KV in FP32. # We setup a separated tmem space for KV16 as operand A for mma. num_kv16_acc_cols = num_kv_acc_cols // 2 # BF16 has half columns - if cutlass.const_expr(PRINT_DEBUG): - print(f"tCtAccKV_fake={tCtAccKV_fake}, num_kv_acc_cols={num_kv_acc_cols}, num_kv16_acc_cols={num_kv16_acc_cols}") acc_shape_sq = tiled_mma_sq.partition_shape_C(tile_shape_mnk_sq[:2]) # No Stage for QS since state has no stages. tCtAccSQ_fake = tiled_mma_sq.make_fragment_C(cute.append(acc_shape_sq, 1)) num_qs_acc_cols = tcgen05.find_tmem_tensor_col_offset(tCtAccSQ_fake) - if cutlass.const_expr(PRINT_DEBUG): - print(f"tCtAccSQ_fake={tCtAccSQ_fake}, num_qs_acc_cols={num_qs_acc_cols}") num_qk_acc_cols_offset = 0 num_pv_acc_cols_offset = num_qk_acc_cols_offset + num_qk_acc_cols @@ -316,14 +309,6 @@ def _plan_tmem_offsets( num_tmem_cols_total *= 2 assert num_tmem_cols_total <= SM100_TMEM_CAPACITY_COLS - if cutlass.const_expr(PRINT_DEBUG): - print(f"num_qk_acc_cols_offset: {num_qk_acc_cols_offset}") - print(f"num_pv_acc_cols_offset: {num_pv_acc_cols_offset}") - print(f"num_kv_acc_cols_offset: {num_kv_acc_cols_offset}") - print(f"num_kv16_acc_cols_offset: {num_kv16_acc_cols_offset}") - print(f"num_qs_acc_cols_offset: {num_qs_acc_cols_offset}") - print(f"num_tmem_cols_total: {num_tmem_cols_total}") - return ( num_qk_acc_cols_offset, num_pv_acc_cols_offset, @@ -649,8 +634,6 @@ def __call__( self.g_dtype, self.g_stage, ) - if PRINT_DEBUG: - print(f"g_smem_layout_staged: {g_smem_layout_staged}") # V^T*P p_smem_layout_staged = sm100_utils.make_smem_layout_b( vp_tiled_mma, @@ -735,10 +718,6 @@ def __call__( qk_tiled_mma, cluster_layout_vmnk.shape, ) - if PRINT_DEBUG: - print(f"tma_atom_g: {cute.pretty_str(tma_atom_g)}") - print(f"g_smem_layout: {cute.pretty_str(g_smem_layout)}") - # NOTE: G's last row will be extracted from sG in CUDA warp after TMA load # No separate TMA needed for G last row - we extract it from the full G tile @@ -752,61 +731,13 @@ def __call__( q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) - v_copy_size = cute.size_in_bytes(self.v_dtype, v_smem_layout) g_copy_size = cute.size_in_bytes(self.g_dtype, g_smem_layout) # NEW for KDA - if PRINT_DEBUG: - print( - f"q_copy_size: {q_copy_size}, k_copy_size: {k_copy_size}, v_copy_size: {v_copy_size}, g_copy_size: {g_copy_size}" - ) self.tma_copy_q_bytes = q_copy_size self.tma_copy_k_bytes = k_copy_size # self.tma_copy_v_bytes = v_copy_size self.tma_copy_v_bytes = k_copy_size self.tma_copy_g_bytes = g_copy_size # NEW for KDA - if cutlass.const_expr(PRINT_DEBUG): - print(f"q_layout: {cute.pretty_str(q_layout)}") - print(f"q: {cute.pretty_str(q)}") - print(f"k_layout: {cute.pretty_str(k_layout)}") - print(f"k: {cute.pretty_str(k)}") - print(f"v_layout: {cute.pretty_str(v_layout)}") - print(f"v: {cute.pretty_str(v)}") - print(f"o_layout: {cute.pretty_str(o_layout)}") - print(f"o: {cute.pretty_str(o)}") - print(f"qk_tiled_mma: {cute.pretty_str(qk_tiled_mma)}") - print(f"kv_tiled_mma: {cute.pretty_str(kv_tiled_mma)}") - print(f"vp_tiled_mma: {cute.pretty_str(vp_tiled_mma)}") - print(f"sq_tiled_mma: {cute.pretty_str(sq_tiled_mma)}") - print(f"cluster_layout_vmnk: {cute.pretty_str(cluster_layout_vmnk)}") - print(f"epi_tile: {cute.pretty_str(self.epi_tile)}") - print(f"q_smem_layout: {cute.pretty_str(q_smem_layout)}") - print(f"k_smem_layout: {cute.pretty_str(k_smem_layout)}") - print(f"v_smem_layout: {cute.pretty_str(v_smem_layout)}") - print(f"q_smem_layout_staged: {cute.pretty_str(q_smem_layout_staged)}") - print(f"k_smem_layout_staged: {cute.pretty_str(k_smem_layout_staged)}") - print(f"k_smem_layout_staged.swzzle: {cute.pretty_str(k_smem_layout_staged.inner)}") - print(f"k_smem_layout_staged.outer: {cute.pretty_str(k_smem_layout_staged.outer)}") - print(f"kv_k_smem_layout_staged: {cute.pretty_str(kv_k_smem_layout_staged)}") - print(f"kv_k_smem_layout_staged.swzzle: {cute.pretty_str(kv_k_smem_layout_staged.inner)}") - print(f"kv_k_smem_layout_staged.outer: {cute.pretty_str(kv_k_smem_layout_staged.outer)}") - print(f"v_smem_layout_staged: {cute.pretty_str(v_smem_layout_staged)}") - print(f"o_smem_layout_staged: {cute.pretty_str(o_smem_layout_staged)}") - print(f"p_smem_layout_staged: {cute.pretty_str(p_smem_layout_staged)}") - print(f"tma_atom_q: {cute.pretty_str(tma_atom_q)}") - print(f"tma_atom_k: {cute.pretty_str(tma_atom_k)}") - print(f"tma_atom_v: {cute.pretty_str(tma_atom_v)}") - print(f"tma_tensor_q: {cute.pretty_str(tma_tensor_q)}") - print(f"tma_tensor_k: {cute.pretty_str(tma_tensor_k)}") - print(f"tma_tensor_v: {cute.pretty_str(tma_tensor_v)}") - - print(f"tma_atom_o: {cute.pretty_str(tma_atom_o)}") - print(f"o_smem_layout: {cute.pretty_str(o_smem_layout)}") - print(f"tma_tensor_o: {cute.pretty_str(tma_tensor_o)}") - - print(f"q_copy_size: {q_copy_size}") - print(f"k_copy_size: {k_copy_size}") - print(f"v_copy_size: {v_copy_size}") - beta_layout = cute.make_layout((Constant.C, self.beta_stage), stride=(1, Constant.C)) g_last_layout = cute.make_layout((Constant.D, self.g_stage), stride=(1, Constant.D)) @@ -900,10 +831,6 @@ class SharedStorage: ] self.shared_storage = SharedStorage - if PRINT_DEBUG: - print(f"size of storage: {SharedStorage.__sizeof__()}") - print(f"m_smem_layout_staged: {m_smem_layout_staged}") - if cutlass.const_expr(self.is_varlen): self.grid = (1, H, B) # TensorMapManager for TMA descriptor modification in varlen tail tiles @@ -914,9 +841,6 @@ class SharedStorage: o_shape=cute.shape(o), chunk_size=self.chunk_size, ) - if PRINT_DEBUG: - print(f"grid: {self.grid}") - self.kernel( qk_tiled_mma, kk_tiled_mma, @@ -1318,23 +1242,15 @@ def kernel( # swizzle_=k_smem_layout_staged.inner, # dtype=self.io_dtype), # layout=sK_neg_g_f32.layout) - if PRINT_DEBUG: - print(f"sK_neg_g: {cute.pretty_str(sK_neg_g)}") - print(f"sK_neg_g_b: {cute.pretty_str(sK_neg_g_b)}") - print(f"sK_g: {cute.pretty_str(sK_g)}") # (((64,2),16),1,4,2):(((1,4096),64),0,1024,8192)> sV = storage.sV.get_tensor(v_smem_layout_staged.outer, swizzle=v_smem_layout_staged.inner) # G (gate/g_cumsum) - NEW for KDA sG = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) # No swizzling for last row of exp(G) sG_last = self.get_smem_tensor_sG_last(storage, g_last_layout) - if PRINT_DEBUG: - print(f"sG_last: {sG_last}") # NOTE: optimize swizzle sBeta = storage.sBeta.get_tensor(beta_layout, swizzle=None) - if PRINT_DEBUG: - print(f"sBeta: {sBeta}") # (MMA, MMA_N, MMA_K, STAGE) sP = storage.sP.get_tensor(p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner) @@ -1459,14 +1375,6 @@ def kernel( layout=k_smem_layout_bf16_fixed, ) - if cutlass.const_expr(PRINT_DEBUG): - print(f"sQ: {cute.pretty_str(sQ)}") - print(f"sK: {cute.pretty_str(sK)}") - print(f"sV: {cute.pretty_str(sV)}") - print(f"sO: {cute.pretty_str(sO)}") - print(f"sP: {cute.pretty_str(sP)}") - print(f"sQK: {cute.pretty_str(sQK)}") - (_, hidx, bidx) = cute.arch.block_idx() B, S, H, D = problem_size C = self.chunk_size @@ -1580,9 +1488,6 @@ def kernel( # (MMA, MMA_M, MMA_K, INPUT_STAGE) # (MMA, MMA_N, MMA_K, INPUT_STAGE) # (MMA, MMA_M, MMA_N, ACC_STAGE) - if PRINT_DEBUG: - print(f"sP: {cute.pretty_str(sP)}") - print(f"sM: {cute.pretty_str(sM)}") tCrV_corr, tCrM, tCtAccMV = self.mma_partition_ss( mv_tiled_mma, self.mv_mma_tiler, @@ -1608,18 +1513,6 @@ def kernel( ) # ROW MAJOR sG_flat = storage.sG.get_tensor(g_smem_layout_coalesce.outer, swizzle=g_smem_layout_coalesce.inner) - # HALF SPACE - CRITICAL FIX: Double the stage stride for BF16 - sG_flat_layout = sG_flat.layout - sG_flat_bf16_layout = cute.make_layout( - sG_flat_layout.shape, stride=(*sG_flat_layout.stride[:-1], sG_flat_layout.stride[-1] * 2) - ) - sG_flat_as_bf16 = cute.make_tensor(cute.recast_ptr(sG_flat.iterator, dtype=self.io_dtype), layout=sG_flat_bf16_layout) - if PRINT_DEBUG: - print(f"sG_flat: {cute.pretty_str(sG_flat)}") - print(f"sG_flat_as_bf16: {cute.pretty_str(sG_flat_as_bf16)}") - print(f"g_smem_layout_epi: {g_smem_layout_epi}") - print(f"g_smem_layout_coalesce: {g_smem_layout_coalesce}") - # /////////////////////////////////////////////////////////////////////////////// # LOAD WARP # /////////////////////////////////////////////////////////////////////////////// @@ -4645,16 +4538,6 @@ def smem_store_acc_as_ab_and_partition_x( # ((V, R), M, N) tXrX_r2s = thr_r2s_x.retile(tXrX_t2r) - if show_debug_info: - print(f"-------------------- SMEM STORE: {debug_name}") - print(f"copy_atom_r2s_x: {copy_atom_r2s_x}") - print(f"tiled_t2r_x: {tiled_t2r_x}") - print(f"thr_t2r_x: {thr_r2s_x}") - print(f"before partition_D: {smem_x}") - print(f"after partition_D, tXsX_r2s: {tXsX_r2s}") - print(f"before retile tXrX_t2r: {tXrX_t2r}") - print(f"after retile tXrX_r2s: {tXrX_r2s}") - return tiled_r2s_x, thr_r2s_x, tXrX_r2s, tXsX_r2s def epilog_tmem_load_and_partition_acc(self, local_tidx, tIntra, smem_y): diff --git a/cula/ops/lightning_attn.py b/cula/ops/lightning_attn_sm100.py similarity index 100% rename from cula/ops/lightning_attn.py rename to cula/ops/lightning_attn_sm100.py diff --git a/cula/ops/linear_attn.py b/cula/ops/linear_attn_sm100.py similarity index 99% rename from cula/ops/linear_attn.py rename to cula/ops/linear_attn_sm100.py index 5cfa0248..64a11755 100644 --- a/cula/ops/linear_attn.py +++ b/cula/ops/linear_attn_sm100.py @@ -42,7 +42,7 @@ .. code-block:: bash - python examples/blackwell/linear_attn.py \\ + python examples/blackwell/linear_attn_sm100.py \\ --batch_size 4 --seq_len 1024 --num_heads 8 --head_dim 64 \\ --chunk_size 64 --decay 0.95 diff --git a/cula/utils.py b/cula/utils.py index bd70730b..eaa094da 100644 --- a/cula/utils.py +++ b/cula/utils.py @@ -83,7 +83,7 @@ def assert_hopper(device: torch.device | str | int | None = None) -> None: def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callable: """Return the appropriate ``kda_prefill`` implementation for *device*. - - sm100/sm103 (Blackwell) → cula.kda.kda_prefill_blackwell (not yet available) + - sm100/sm103 (Blackwell) → cula.kda.blackwell_fused_fwd.flash_kda_prefill - sm90 (Hopper) → cula.kda.kda_prefill_hopper Args: @@ -94,11 +94,9 @@ def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callabl """ major, minor = get_device_sm_version(device) if major == 10 and minor in (0, 3): - # TODO - raise NotImplementedError( - "The Blackwell implementation of fused prefill is not yet available. " - "Please use a sm90a (Hopper) device or wait for future updates." - ) + from cula.kda import kda_prefill_blackwell + + return kda_prefill_blackwell elif major == 9 and minor == 0: from cula.kda import kda_prefill_hopper diff --git a/pyproject.toml b/pyproject.toml index 377dffe2..ff2bf950 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,7 +80,7 @@ force-sort-within-sections = false "__init__.py" = ["F401"] "cula/_version.py" = ["UP007"] # TODO: fix undefined names (exp_g, chunk_kda_bwd_dqkwg) — WIP code -"cula/ops/kda_fully_fused_wip.py" = ["F821"] +"cula/ops/kda_fully_fused_sm100_wip.py" = ["F821"] "cula/kda/blackwell_fused_fwd.py" = ["F821"] [tool.setuptools_scm] diff --git a/tests/test_chunk_delta_h.py b/tests/test_chunk_delta_h.py index 868a4f8a..01cdb157 100644 --- a/tests/test_chunk_delta_h.py +++ b/tests/test_chunk_delta_h.py @@ -22,7 +22,7 @@ import importlib.util _spec = importlib.util.spec_from_file_location( - "chunk_delta_h", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "chunk_delta_h.py") + "chunk_delta_h_sm100", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "chunk_delta_h_sm100.py") ) _mod = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_mod) diff --git a/tests/test_compare_with_fla.py b/tests/test_compare_with_fla.py index 93663c31..3468c9c4 100644 --- a/tests/test_compare_with_fla.py +++ b/tests/test_compare_with_fla.py @@ -33,7 +33,7 @@ from cutlass.cute.runtime import from_dlpack # noqa: E402 # Our implementation -from cula.ops.chunk_delta_h import ChunkDeltaRuleFwdH # noqa: E402 +from cula.ops.chunk_delta_h_sm100 import ChunkDeltaRuleFwdH # noqa: E402 def fla_reference_chunk_fwd_h( diff --git a/tests/test_fwd_o.py b/tests/test_fwd_o.py index a4a2c827..51b30003 100644 --- a/tests/test_fwd_o.py +++ b/tests/test_fwd_o.py @@ -37,7 +37,7 @@ import importlib.util _spec = importlib.util.spec_from_file_location( - "fwd_o", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "fwd_o.py") + "fwd_o_sm100", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "fwd_o_sm100.py") ) _mod = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_mod) diff --git a/tests/test_la_decode.py b/tests/test_la_decode.py index 688900f7..5b57ac58 100644 --- a/tests/test_la_decode.py +++ b/tests/test_la_decode.py @@ -233,7 +233,7 @@ def test_vs_fla(B): # --------------------------------------------------------------------------- def test_prefill_decode_e2e(): """Verify prefill output state passes directly into decode without transpose.""" - from cula.ops.lightning_attn import lightning_attn_fwd + from cula.ops.lightning_attn_sm100 import lightning_attn_fwd B, S, H, D = 2, 64, 8, 128 scale = D**-0.5 diff --git a/tests/test_lightning_attn.py b/tests/test_lightning_attn.py index 26fcc16e..5e52f86c 100644 --- a/tests/test_lightning_attn.py +++ b/tests/test_lightning_attn.py @@ -35,7 +35,7 @@ warnings.filterwarnings("ignore", category=DeprecationWarning) sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.ops.lightning_attn import lightning_attn_fwd, lightning_attn_fwd_varlen # noqa: E402 +from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen # noqa: E402 try: from fla.ops.simple_gla import chunk_simple_gla From 7564280f541dfdd0fcebea6a3d30b0ecc5719db5 Mon Sep 17 00:00:00 2001 From: Cavan Yu <95609028+icavanyu@users.noreply.github.com> Date: Mon, 25 May 2026 11:13:00 +0800 Subject: [PATCH 18/34] Consolidate duplicated benchmark helper logic into `benchmarks/utils.py`, including shared timing, reporting, and accuracy utilities. (#80) * chore: rename sm100-specific cutedsl kernels * chore: revert cula/kda naming changes * chore: drop debugging codes * chore: fix lint errors * Refine benchmark accuracy reporting * Fix native GVA chunk intra path * Align GVA intra path with main * Revert GVA test files to main * Refine benchmark error helpers --------- Co-authored-by: icavan --- benchmarks/bench_chunk_delta_h.py | 56 +++------ benchmarks/bench_fwd_o.py | 62 ++++------ benchmarks/bench_kda.py | 85 +++++++------ benchmarks/bench_kda_chunk_intra.py | 43 +++---- benchmarks/bench_kda_decode.py | 159 ++++++++++++------------ benchmarks/bench_kda_fused_fwd.py | 81 ++++++------- benchmarks/bench_kda_fwd_bwd_e2e.py | 109 +++++++++-------- benchmarks/bench_la_decode_vs_fla.py | 47 ++------ benchmarks/bench_lightning_attn.py | 173 ++++++++++++--------------- benchmarks/bench_linear_attn.py | 25 +++- benchmarks/bench_recompute_wu.py | 59 +++++---- benchmarks/utils.py | 164 +++++++++++++++++++++++++ pyproject.toml | 2 +- 13 files changed, 575 insertions(+), 490 deletions(-) diff --git a/benchmarks/bench_chunk_delta_h.py b/benchmarks/bench_chunk_delta_h.py index 09fde953..d4761dc4 100644 --- a/benchmarks/bench_chunk_delta_h.py +++ b/benchmarks/bench_chunk_delta_h.py @@ -18,7 +18,7 @@ for chunk_delta_rule_fwd_h (inter-chunk recurrent state) Compares: - - Accuracy: max_diff, mean_diff between CuTe DSL and FLA Triton outputs + - Accuracy: relative_rms_error, max_diff, mean_diff between CuTe DSL and FLA Triton outputs - Performance: kernel execution time (ms) with CUDA events Both non-varlen and varlen modes are supported. @@ -46,6 +46,7 @@ import numpy as np import torch +from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_mean_abs # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── _delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") @@ -70,29 +71,6 @@ # ============================================================ # Helpers # ============================================================ -def time_kernel(fn, warmup=None, n_iters=None): - if warmup is None: - warmup = 1 if NCU_MODE else WARMUP - if n_iters is None: - n_iters = 1 if NCU_MODE else N_ITERS - """Time a kernel using CUDA events. Returns ms/call.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start_evt = torch.cuda.Event(enable_timing=True) - end_evt = torch.cuda.Event(enable_timing=True) - start_evt.record() - for _ in range(n_iters): - fn() - end_evt.record() - torch.cuda.synchronize() - return start_evt.elapsed_time(end_evt) / n_iters - - -def accuracy_stats(ref, out): - """Compute max and mean absolute difference.""" - diff = (ref.float() - out.float()).abs() - return diff.max().item(), diff.mean().item() # ============================================================ @@ -148,7 +126,7 @@ def bench_non_varlen(configs): h_out = cute_result[0] torch.cuda.synchronize() - max_diff, mean_diff = accuracy_stats(h_fla, h_out) + relative_rms_error, max_diff, mean_diff = relative_rms_error_max_mean_abs(h_fla, h_out) # ---- Performance timing ---- def run_fla(k=k, w=w, u=u, gk=gk, h0=h0): @@ -177,8 +155,8 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0): save_new_value=save_vnew, ) - ms_fla = time_kernel(run_fla) - ms_cute = time_kernel(run_cute) + ms_fla = benchmark_cuda_mode_fn(run_fla, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) + ms_cute = benchmark_cuda_mode_fn(run_cute, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) speedup = ms_fla / ms_cute if ms_cute > 0 else float("inf") flags = [] @@ -197,6 +175,7 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0): "T": T, "H": H, "flags": flag_str, + "relative_rms_error": relative_rms_error, "max_diff": max_diff, "mean_diff": mean_diff, "ms_fla": ms_fla, @@ -206,7 +185,7 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0): results.append(r) print( f" B={B:2d} T={T:5d} H={H:3d}{flag_str:<16s} | " - f"max_diff={max_diff:.6f} mean_diff={mean_diff:.8f} | " + f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" ) @@ -309,7 +288,7 @@ def bench_varlen(configs): h_out = cute_result[0] torch.cuda.synchronize() - max_diff, mean_diff = accuracy_stats(h_fla, h_out) + relative_rms_error, max_diff, mean_diff = relative_rms_error_max_mean_abs(h_fla, h_out) # ---- Performance timing ---- def run_fla(k=k, w=w, u=u, gk=gk, h0=h0, cu=cu_seqlens_long): @@ -340,8 +319,8 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0, cu=cu_seqlens): cu_seqlens=cu, ) - ms_fla = time_kernel(run_fla) - ms_cute = time_kernel(run_cute) + ms_fla = benchmark_cuda_mode_fn(run_fla, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) + ms_cute = benchmark_cuda_mode_fn(run_cute, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) speedup = ms_fla / ms_cute if ms_cute > 0 else float("inf") min_l, max_l = min(seq_lens), max(seq_lens) @@ -365,6 +344,7 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0, cu=cu_seqlens): "H": H, "n_seqs": num_seqs, "flags": flag_str, + "relative_rms_error": relative_rms_error, "max_diff": max_diff, "mean_diff": mean_diff, "ms_fla": ms_fla, @@ -374,7 +354,7 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0, cu=cu_seqlens): results.append(r) print( f" {tag:40s} H={H:3d}{flag_str:<16s} | " - f"max_diff={max_diff:.6f} mean_diff={mean_diff:.8f} | " + f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" ) @@ -401,7 +381,7 @@ def print_report(nv_results, vl_results): print("\n [Non-Varlen]") print(f" {'─' * 100}") print( - f" {'Config':<35s} │ {'max_diff':>10s} {'mean_diff':>12s}" + f" {'Config':<35s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 100}") @@ -409,19 +389,19 @@ def print_report(nv_results, vl_results): label = f"B={r['B']:2d} T={r['T']:5d} H={r['H']:3d}{r['flags']}" print( f" {label:<35s} │ " - f"{r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 100}") speedups = [r["speedup"] for r in nv_results] geo = math.exp(sum(math.log(s) for s in speedups) / len(speedups)) - print(f" {'Geometric mean':<35s} │ {'':>10s} {'':>12s} │ {'':>9s} {'':>9s} {geo:7.2f}x") + print(f" {'Geometric mean':<35s} │ {'':>18s} {'':>10s} {'':>12s} │ {'':>9s} {'':>9s} {geo:7.2f}x") if vl_results: print("\n [Varlen]") print(f" {'─' * 115}") print( - f" {'Config':>55s} │ {'max_diff':>10s} {'mean_diff':>12s}" + f" {'Config':>55s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 115}") @@ -429,13 +409,13 @@ def print_report(nv_results, vl_results): label = f"{r['tag']} H={r['H']:3d}{r['flags']}" print( f" {label:>55s} │ " - f"{r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 115}") speedups = [r["speedup"] for r in vl_results] geo = math.exp(sum(math.log(s) for s in speedups) / len(speedups)) - print(f" {'Geometric mean':>55s} │ {'':>10s} {'':>12s} │ {'':>9s} {'':>9s} {geo:7.2f}x") + print(f" {'Geometric mean':>55s} │ {'':>18s} {'':>10s} {'':>12s} │ {'':>9s} {'':>9s} {geo:7.2f}x") print(f"\n{sep}\n") diff --git a/benchmarks/bench_fwd_o.py b/benchmarks/bench_fwd_o.py index 6fcbca00..94ecb4af 100644 --- a/benchmarks/bench_fwd_o.py +++ b/benchmarks/bench_fwd_o.py @@ -18,7 +18,7 @@ for chunk_gla_fwd_o (KDA forward output) Compares: - - Accuracy: max_diff, mean_diff between CuTe DSL and FLA Triton outputs + - Accuracy: relative_rms_error, max_diff, mean_diff between CuTe DSL and FLA Triton outputs - Performance: kernel execution time (ms) with CUDA events Both non-varlen and varlen modes are supported. @@ -44,14 +44,16 @@ import torch +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── _fwd_o_mod = importlib.import_module("cula.ops.fwd_o_sm100") chunk_gla_fwd_o = _fwd_o_mod.chunk_gla_fwd_o build_chunk_indices = _fwd_o_mod.build_chunk_indices # ─── FLA baseline imports ─── -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison +from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_rel_mean_abs from fla.ops.gla.chunk import chunk_gla_fwd_o_gk # noqa: E402 # ============================================================ @@ -69,34 +71,6 @@ # ============================================================ # Helpers # ============================================================ -def time_kernel(fn, warmup=None, n_iters=None): - if warmup is None: - warmup = 1 if NCU_MODE else WARMUP - if n_iters is None: - n_iters = 1 if NCU_MODE else N_ITERS - """Time a kernel using CUDA events. Returns ms/call.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start_evt = torch.cuda.Event(enable_timing=True) - end_evt = torch.cuda.Event(enable_timing=True) - start_evt.record() - for _ in range(n_iters): - fn() - end_evt.record() - torch.cuda.synchronize() - return start_evt.elapsed_time(end_evt) / n_iters - - -def accuracy_stats(ref, out): - """Compute max, relative max, and mean absolute difference.""" - ref_f = ref.float() - diff = (ref_f - out.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - denom = ref_f.abs().max().item() - rel_max_diff = max_diff / denom if denom > 0 else 0.0 - return max_diff, rel_max_diff, mean_diff # ============================================================ @@ -150,7 +124,7 @@ def bench_non_varlen(configs): ) torch.cuda.synchronize() - max_diff, rel_max_diff, mean_diff = accuracy_stats(o_fla, o_cute_t) + relative_rms_error, max_diff, rel_max_diff, mean_diff = relative_rms_error_max_rel_mean_abs(o_fla, o_cute_t) # ---- Performance timing ---- def run_fla(q=q, v=v, g=g, A=A, h=h, scale=scale): @@ -179,14 +153,15 @@ def run_cute(q=q, v=v, g=g, h=h, o=o_cute_t, A=A, scale=scale): persistent=True, ) - ms_fla = time_kernel(run_fla) - ms_cute = time_kernel(run_cute) + ms_fla = benchmark_cuda_mode_fn(run_fla, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) + ms_cute = benchmark_cuda_mode_fn(run_cute, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) speedup = ms_fla / ms_cute if ms_cute > 0 else float("inf") r = { "B": B, "T": T, "H": H, + "relative_rms_error": relative_rms_error, "max_diff": max_diff, "rel_max_diff": rel_max_diff, "mean_diff": mean_diff, @@ -197,7 +172,7 @@ def run_cute(q=q, v=v, g=g, h=h, o=o_cute_t, A=A, scale=scale): results.append(r) print( f" B={B:2d} T={T:5d} H={H:2d} | " - f"max_diff={max_diff:.6f} rel_max={rel_max_diff:.6f} mean_diff={mean_diff:.8f} | " + f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} rel_max={rel_max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" ) @@ -289,7 +264,9 @@ def bench_varlen(configs): torch.cuda.synchronize() # Both outputs are [1, T_total, H, V]; squeeze to [T_total, H, V] for comparison - max_diff, rel_max_diff, mean_diff = accuracy_stats(o_fla.squeeze(0), o_cute_flat.squeeze(0)) + relative_rms_error, max_diff, rel_max_diff, mean_diff = relative_rms_error_max_rel_mean_abs( + o_fla.squeeze(0), o_cute_flat.squeeze(0) + ) # ---- Performance timing ---- def run_fla(q_flat=q_flat, v_flat=v_flat, g_flat=g_flat, A_flat=A_flat, h_flat=h_flat, cu_fla=cu_fla, scale=scale): @@ -331,8 +308,8 @@ def run_cute( persistent=True, ) - ms_fla = time_kernel(run_fla) - ms_cute = time_kernel(run_cute) + ms_fla = benchmark_cuda_mode_fn(run_fla, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) + ms_cute = benchmark_cuda_mode_fn(run_cute, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) speedup = ms_fla / ms_cute if ms_cute > 0 else float("inf") n_seqs = len(seq_lens) @@ -344,6 +321,7 @@ def run_cute( "T_total": T_total, "H": H, "n_seqs": n_seqs, + "relative_rms_error": relative_rms_error, "max_diff": max_diff, "rel_max_diff": rel_max_diff, "mean_diff": mean_diff, @@ -354,7 +332,7 @@ def run_cute( results.append(r) print( f" {tag:45s} H={H:2d} | " - f"max_diff={max_diff:.6f} rel_max={rel_max_diff:.6f} mean_diff={mean_diff:.8f} | " + f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} rel_max={rel_max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" ) @@ -380,7 +358,7 @@ def print_report(nv_results, vl_results): if nv_results: print("\n [Non-Varlen]") hdr = ( - f" {'B':>3s} {'T':>5s} {'H':>3s} │ {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" + f" {'B':>3s} {'T':>5s} {'H':>3s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 90}") @@ -389,7 +367,7 @@ def print_report(nv_results, vl_results): for r in nv_results: print( f" {r['B']:3d} {r['T']:5d} {r['H']:3d} │ " - f"{r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 90}") @@ -397,7 +375,7 @@ def print_report(nv_results, vl_results): if vl_results: print("\n [Varlen]") hdr = ( - f" {'Config':>45s} {'H':>3s} │ {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" + f" {'Config':>45s} {'H':>3s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 117}") @@ -406,7 +384,7 @@ def print_report(nv_results, vl_results): for r in vl_results: print( f" {r['tag']:>45s} {r['H']:3d} │ " - f"{r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 117}") diff --git a/benchmarks/bench_kda.py b/benchmarks/bench_kda.py index dc31d11d..78bbe241 100644 --- a/benchmarks/bench_kda.py +++ b/benchmarks/bench_kda.py @@ -18,7 +18,7 @@ for chunk_kda (KDA forward) Compares: - - Accuracy: RMSE, relative max diff between cuLA and FLA outputs + - Accuracy: relative_rms_error, relative max diff between cuLA and FLA outputs - Performance: kernel execution time (ms) with CUDA events Modes: @@ -45,9 +45,11 @@ from benchmarks.utils import ( SEED, + benchmark_cuda_mode_fn, build_varlen_configs, exclusive_cumsum, prepare_safe_gate_inputs, + relative_rms_error_rel_max_mean_abs, set_seed, ) from cula.kda import chunk_kda as cula_chunk_kda @@ -72,37 +74,6 @@ def generate_balanced_seqlens(total_tokens, num_seqs): return [base] * (num_seqs - 1) + [base + remainder] -def time_kernel(fn, warmup=None, n_iters=None): - if warmup is None: - warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP - if n_iters is None: - n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start_evt = torch.cuda.Event(enable_timing=True) - end_evt = torch.cuda.Event(enable_timing=True) - start_evt.record() - for _ in range(n_iters): - fn() - end_evt.record() - torch.cuda.synchronize() - return start_evt.elapsed_time(end_evt) / n_iters - - -def accuracy_stats(ref, out): - """Compute RMSE, relative max diff, and mean absolute difference.""" - ref_f = ref.float() - out_f = out.float() - diff = (ref_f - out_f).abs() - rmse = diff.pow(2).mean().sqrt().item() - max_diff = diff.max().item() - denom = ref_f.abs().max().item() - rel_max = max_diff / denom if denom > 0 else 0.0 - mean_diff = diff.mean().item() - return rmse, rel_max, mean_diff - - def run_kda(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound, fn): return fn( q=q, @@ -203,7 +174,7 @@ def bench_fixed(configs): o_cula, _ = run_kda(**common, fn=cula_chunk_kda) torch.cuda.synchronize() - rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) # Performance def fn_fla(**common_kw): @@ -212,14 +183,26 @@ def fn_fla(**common_kw): def fn_cula(**common_kw): return lambda: run_kda(**common_kw, fn=cula_chunk_kda) - ms_fla = time_kernel(fn_fla(**common)) - ms_cula = time_kernel(fn_cula(**common)) + ms_fla = benchmark_cuda_mode_fn( + fn_fla(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + fn_cula(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") r = { "B": B, "T": T, - "rmse": rmse, + "relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff, "ms_fla": ms_fla, @@ -276,7 +259,7 @@ def bench_varlen(configs): o_cula, _ = run_kda(**common, fn=cula_chunk_kda) torch.cuda.synchronize() - rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) # Performance def fn_fla(**common_kw): @@ -285,8 +268,20 @@ def fn_fla(**common_kw): def fn_cula(**common_kw): return lambda: run_kda(**common_kw, fn=cula_chunk_kda) - ms_fla = time_kernel(fn_fla(**common)) - ms_cula = time_kernel(fn_cula(**common)) + ms_fla = benchmark_cuda_mode_fn( + fn_fla(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + fn_cula(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") n_seqs = len(seq_lens) @@ -299,7 +294,7 @@ def fn_cula(**common_kw): "dist": dist, "T_total": T, "n_seqs": n_seqs, - "rmse": rmse, + "relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff, "ms_fla": ms_fla, @@ -334,14 +329,14 @@ def print_report(fixed_results, varlen_results): print("\n [Fixed-Length]") print(f" {'─' * 85}") print( - f" {'B':>3s} {'T':>5s} │ {'RMSE':>10s} {'rel_max':>10s}" + f" {'B':>3s} {'T':>5s} │ {'rel_rmse':>18s} {'rel_max':>10s}" f" │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s}" ) print(f" {'─' * 85}") for r in fixed_results: print( f" {r['B']:3d} {r['T']:5d} │ " - f"{r['rmse']:10.6f} {r['rel_max']:10.6f} │ " + f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 85}") @@ -349,12 +344,14 @@ def print_report(fixed_results, varlen_results): if varlen_results: print("\n [Varlen]") print(f" {'─' * 100}") - print(f" {'Config':>45s} │ {'RMSE':>10s} {'rel_max':>10s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s}") + print( + f" {'Config':>45s} │ {'rel_rmse':>18s} {'rel_max':>10s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s}" + ) print(f" {'─' * 100}") for r in varlen_results: print( f" {r['tag']:>45s} │ " - f"{r['rmse']:10.6f} {r['rel_max']:10.6f} │ " + f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 100}") diff --git a/benchmarks/bench_kda_chunk_intra.py b/benchmarks/bench_kda_chunk_intra.py index 4c2718d5..c44fe63d 100644 --- a/benchmarks/bench_kda_chunk_intra.py +++ b/benchmarks/bench_kda_chunk_intra.py @@ -28,14 +28,20 @@ import sys import torch -import triton sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra -from benchmarks.utils import SEED, exclusive_cumsum, generate_random_seq_lens, prepare_intra_inputs +from benchmarks.utils import ( + SEED, + exclusive_cumsum, + generate_random_seq_lens, + prepare_intra_inputs, + relative_rms_error_rel_max_mean_abs_rhs, + triton_bench_fn, +) from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra # Constant params @@ -51,19 +57,6 @@ DISABLE_RECOMPUTE = False # Whether to disable recompute (compute QG in forward) - -def accuracy_stats(a, b): - """Compute RMSE, relative max diff, and mean absolute difference.""" - a, b = a.float(), b.float() - diff = a - b - rmse = diff.pow(2).mean().sqrt().item() - max_diff = diff.abs().max().item() - denom = b.abs().max().item() - rel_max = max_diff / denom if denom > 0 else 0.0 - mean_diff = diff.abs().mean().item() - return rmse, rel_max, mean_diff - - # ============================================================================== # Unified uniform seqlen benchmark (handles both standard and GVA) # ============================================================================== @@ -83,7 +76,7 @@ def benchmark_chunk_intra_uniform(): ) print("=" * 100) print( - f"{'B':>4} {'T':>7} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'B':>4} {'T':>7} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 100) @@ -106,15 +99,15 @@ def benchmark_chunk_intra_uniform(): out_cula = cula_chunk_kda_fwd_intra(**common) o_fla = out_fla[0] if isinstance(out_fla, (tuple, list)) else out_fla o_cula = out_cula[0] if isinstance(out_cula, (tuple, list)) else out_cula - rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs_rhs(o_fla, o_cula) # Performance - ms_fla = triton.testing.do_bench(lambda: fla_chunk_kda_fwd_intra(**common)) - ms_cula = triton.testing.do_bench(lambda: cula_chunk_kda_fwd_intra(**common)) + ms_fla = triton_bench_fn(lambda: fla_chunk_kda_fwd_intra(**common)) + ms_cula = triton_bench_fn(lambda: cula_chunk_kda_fwd_intra(**common)) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{B:>4} {T:>7} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{B:>4} {T:>7} │ {relative_rms_error:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 100) @@ -140,7 +133,7 @@ def benchmark_chunk_intra_varlen(): ) print("=" * 110) print( - f"{'total_len':>10} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'total_len':>10} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 110) @@ -164,15 +157,15 @@ def benchmark_chunk_intra_varlen(): out_cula = cula_chunk_kda_fwd_intra(**common) o_fla = out_fla[0] if isinstance(out_fla, (tuple, list)) else out_fla o_cula = out_cula[0] if isinstance(out_cula, (tuple, list)) else out_cula - rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs_rhs(o_fla, o_cula) # Performance - ms_fla = triton.testing.do_bench(lambda: fla_chunk_kda_fwd_intra(**common)) - ms_cula = triton.testing.do_bench(lambda: cula_chunk_kda_fwd_intra(**common)) + ms_fla = triton_bench_fn(lambda: fla_chunk_kda_fwd_intra(**common)) + ms_cula = triton_bench_fn(lambda: cula_chunk_kda_fwd_intra(**common)) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{total_len:>10} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{total_len:>10} │ {relative_rms_error:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 110) diff --git a/benchmarks/bench_kda_decode.py b/benchmarks/bench_kda_decode.py index ed5209c4..5da9afb2 100644 --- a/benchmarks/bench_kda_decode.py +++ b/benchmarks/bench_kda_decode.py @@ -31,6 +31,8 @@ python benchmarks/bench_kda_decode.py --batch-sizes 1 4 16 64 128 256 python benchmarks/bench_kda_decode.py --Hs 8 32 64 python benchmarks/bench_kda_decode.py --ncu + python benchmarks/bench_kda_decode.py --output + python benchmarks/bench_kda_decode.py --output tmp/kda_decode_bench.md Note: - This benchmark is currently restricted to K=128 and V=128. @@ -51,44 +53,11 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +from benchmarks.utils import benchmark_cuda_fn, relative_rms_error_rel_max from cula.kda import fused_sigmoid_gating_delta_rule_update as cula_fused from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as fla_fused -# ────────────────────────────────────────────────────────────────────── -# Timing utility -# ────────────────────────────────────────────────────────────────────── -def benchmark_fn(fn, *, setup_fn=None, warmup=30, rep=200): - """Benchmark using CUDA events. - - If provided, setup_fn runs before each iteration and is excluded from the - timing window. This is used to reset the mutable recurrent state fairly. - """ - for _ in range(warmup): - if setup_fn is not None: - setup_fn() - fn() - torch.cuda.synchronize() - - starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - - for i in range(rep): - if setup_fn is not None: - setup_fn() - starts[i].record() - fn() - ends[i].record() - - torch.cuda.synchronize() - times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends)) - n = len(times) - if n <= 2: - return sum(times) / max(len(times), 1) - iqr = times[n // 4 : 3 * n // 4] - return sum(iqr) / len(iqr) - - # ────────────────────────────────────────────────────────────────────── # Input generation # ────────────────────────────────────────────────────────────────────── @@ -106,19 +75,6 @@ def make_inputs(N, H, HV, K, V, device="cuda", seed=42): return q, k, v, a, b, A_log, dt_bias, state -# ────────────────────────────────────────────────────────────────────── -# Accuracy check -# ────────────────────────────────────────────────────────────────────── -def accuracy_stats(ref, out): - ref_f, out_f = ref.float(), out.float() - diff = (ref_f - out_f).abs() - rmse = diff.pow(2).mean().sqrt().item() - max_diff = diff.max().item() - denom = ref_f.abs().max().item() - rel_max = max_diff / denom if denom > 0 else 0.0 - return rmse, rel_max - - def to_v_last_state(state: torch.Tensor, layout: str) -> torch.Tensor: if layout == "kv": return state @@ -152,7 +108,14 @@ def normalize_gpu_type(gpu_name: str) -> str: return "_".join(tokens) if tokens else "UNKNOWN_GPU" -def write_markdown_report(args, gpu_name: str, sections: list[tuple[int, int, list[dict]]], output_path: pathlib.Path): +def write_markdown_report( + args, + gpu_name: str, + sections: list[tuple[int, int, list[dict]]], + output_path: pathlib.Path, + *, + generator_name: str = "benchmarks/bench_kda_decode.py", +): """Write benchmark results into a markdown report.""" now = datetime.now().strftime("%Y-%m-%d %H:%M:%S") py_ver = platform.python_version() @@ -170,7 +133,7 @@ def summary(vals): lines = [] lines.append("# Benchmark Results - KDA Decode") lines.append("") - lines.append(f"> Auto-generated by `benchmarks/bench_kda_decode.py` on {now}.") + lines.append(f"> Auto-generated by `{generator_name}` on {now}.") lines.append("") lines.append(f"> **GPU:** {gpu_name} | **CUDA:** {cuda_ver} | **PyTorch:** {torch_ver} | **Python:** {py_ver}") lines.append("") @@ -202,23 +165,25 @@ def summary(vals): lines.append("### Accuracy (Output)") lines.append("") - lines.append("| N | cuLA v out RMSE | cuLA v out rel | cuLA k out RMSE | cuLA k out rel |") - lines.append("|---|----------------:|---------------:|----------------:|---------------:|") + lines.append("| N | cuLA v out rel_rmse | cuLA v out rel_max | cuLA k out rel_rmse | cuLA k out rel_max |") + lines.append("|---|------------------------------:|-------------------:|------------------------------:|-------------------:|") for r in results: lines.append( - f"| {r['N']} | {r['out_v_last_rmse']:.3e} | {r['out_v_last_rel']:.3e} | " - f"{r['out_k_last_rmse']:.3e} | {r['out_k_last_rel']:.3e} |" + f"| {r['N']} | {r['out_v_last_relative_rms_error']:.3e} | {r['out_v_last_rel_max']:.3e} | " + f"{r['out_k_last_relative_rms_error']:.3e} | {r['out_k_last_rel_max']:.3e} |" ) lines.append("") lines.append("### Accuracy (State)") lines.append("") - lines.append("| N | cuLA v state RMSE | cuLA v state rel | cuLA k state RMSE | cuLA k state rel |") - lines.append("|---|------------------:|-----------------:|------------------:|-----------------:|") + lines.append( + "| N | cuLA v state rel_rmse | cuLA v state rel_max | cuLA k state rel_rmse | cuLA k state rel_max |" + ) + lines.append("|---|--------------------------------:|---------------------:|--------------------------------:|---------------------:|") for r in results: lines.append( - f"| {r['N']} | {r['state_v_last_rmse']:.3e} | {r['state_v_last_rel']:.3e} | " - f"{r['state_k_last_rmse']:.3e} | {r['state_k_last_rel']:.3e} |" + f"| {r['N']} | {r['state_v_last_relative_rms_error']:.3e} | {r['state_v_last_rel_max']:.3e} | " + f"{r['state_k_last_relative_rms_error']:.3e} | {r['state_k_last_rel_max']:.3e} |" ) lines.append("") @@ -232,6 +197,10 @@ def summary(vals): output_path.write_text("\n".join(lines), encoding="utf-8") +def default_report_path(gpu_name: str) -> pathlib.Path: + return pathlib.Path(f"BENCHMARK_KDA_DECODE_{normalize_gpu_type(gpu_name)}.md") + + # ────────────────────────────────────────────────────────────────────── # Run one config # ────────────────────────────────────────────────────────────────────── @@ -316,10 +285,12 @@ def call_fla_v_last(state_buf): o_cula_k_last = call_cula_k_last(state_cula_k_last) o_fla_v_last = call_fla_v_last(state_fla_v_last) - out_v_last_rmse, out_v_last_rel = accuracy_stats(o_fla_v_last, o_cula_v_last) - out_k_last_rmse, out_k_last_rel = accuracy_stats(o_fla_v_last, o_cula_k_last) - state_v_last_rmse, state_v_last_rel = accuracy_stats(state_fla_v_last, state_cula_v_last) - state_k_last_rmse, state_k_last_rel = accuracy_stats(state_fla_v_last, to_v_last_state(state_cula_k_last, "vk")) + out_v_last_relative_rms_error, out_v_last_rel_max = relative_rms_error_rel_max(o_fla_v_last, o_cula_v_last) + out_k_last_relative_rms_error, out_k_last_rel_max = relative_rms_error_rel_max(o_fla_v_last, o_cula_k_last) + state_v_last_relative_rms_error, state_v_last_rel_max = relative_rms_error_rel_max(state_fla_v_last, state_cula_v_last) + state_k_last_relative_rms_error, state_k_last_rel_max = relative_rms_error_rel_max( + state_fla_v_last, to_v_last_state(state_cula_k_last, "vk") + ) if ncu_mode: w, r = 1, 1 @@ -340,13 +311,13 @@ def setup_fla_v_last(): state_bench_fla_v_last.copy_(state_init_v_last) with torch.no_grad(): - t_cula_v_last = benchmark_fn( + t_cula_v_last = benchmark_cuda_fn( lambda: call_cula_v_last(state_bench_cula_v_last), setup_fn=setup_cula_v_last, warmup=w, rep=r ) - t_cula_k_last = benchmark_fn( + t_cula_k_last = benchmark_cuda_fn( lambda: call_cula_k_last(state_bench_cula_k_last), setup_fn=setup_cula_k_last, warmup=w, rep=r ) - t_fla_v_last = benchmark_fn( + t_fla_v_last = benchmark_cuda_fn( lambda: call_fla_v_last(state_bench_fla_v_last), setup_fn=setup_fla_v_last, warmup=w, rep=r ) @@ -361,21 +332,21 @@ def setup_fla_v_last(): "t_fla_v_last_ms": t_fla_v_last, "speedup_v_last": t_fla_v_last / t_cula_v_last if t_cula_v_last > 0 else float("inf"), "speedup_k_last": t_fla_v_last / t_cula_k_last if t_cula_k_last > 0 else float("inf"), - "out_v_last_rmse": out_v_last_rmse, - "out_v_last_rel": out_v_last_rel, - "out_k_last_rmse": out_k_last_rmse, - "out_k_last_rel": out_k_last_rel, - "state_v_last_rmse": state_v_last_rmse, - "state_v_last_rel": state_v_last_rel, - "state_k_last_rmse": state_k_last_rmse, - "state_k_last_rel": state_k_last_rel, + "out_v_last_relative_rms_error": out_v_last_relative_rms_error, + "out_v_last_rel_max": out_v_last_rel_max, + "out_k_last_relative_rms_error": out_k_last_relative_rms_error, + "out_k_last_rel_max": out_k_last_rel_max, + "state_v_last_relative_rms_error": state_v_last_relative_rms_error, + "state_v_last_rel_max": state_v_last_rel_max, + "state_k_last_relative_rms_error": state_k_last_relative_rms_error, + "state_k_last_rel_max": state_k_last_rel_max, } # ────────────────────────────────────────────────────────────────────── # Main # ────────────────────────────────────────────────────────────────────── -def main(): +def build_parser(): parser = argparse.ArgumentParser(description="Benchmark KDA decode: cuLA vs FLA") parser.add_argument("--batch-sizes", nargs="+", type=int, default=[1, 4, 8, 16, 32, 64, 128, 256]) parser.add_argument("--Hs", nargs="+", type=int, default=[8, 32, 64], help="Q/K head counts to benchmark") @@ -385,7 +356,19 @@ def main(): parser.add_argument("--warmup", type=int, default=30) parser.add_argument("--rep", type=int, default=200) parser.add_argument("--ncu", action="store_true", help="NCU mode: warmup=1, rep=1") - args = parser.parse_args() + parser.add_argument( + "--output", + nargs="?", + const="__AUTO__", + default=None, + help="Write markdown report. Omit the value to use the default BENCHMARK_KDA_DECODE_.md filename.", + ) + return parser + + +def main(argv=None): + parser = build_parser() + args = parser.parse_args(argv) if args.K != 128 or args.V != 128: raise ValueError(f"bench_kda_decode.py currently only supports K=128 and V=128, got K={args.K}, V={args.V}") @@ -427,23 +410,29 @@ def print_section(h_dim: int, v_dim: int): ) print() - hdr_out = f"{'N':>5} | {'cuLA v out RMSE':>16} | {'rel':>10} | {'cuLA k out RMSE':>16} | {'rel':>10}" + hdr_out = ( + f"{'N':>5} | {'cuLA v out rel_rmse':>30} | {'rel_max':>10} | " + f"{'cuLA k out rel_rmse':>30} | {'rel_max':>10}" + ) print(hdr_out) print("-" * len(hdr_out)) for res in results: print( - f"{res['N']:5d} | {res['out_v_last_rmse']:16.3e} | {res['out_v_last_rel']:10.3e} | " - f"{res['out_k_last_rmse']:16.3e} | {res['out_k_last_rel']:10.3e}" + f"{res['N']:5d} | {res['out_v_last_relative_rms_error']:30.3e} | {res['out_v_last_rel_max']:10.3e} | " + f"{res['out_k_last_relative_rms_error']:30.3e} | {res['out_k_last_rel_max']:10.3e}" ) print() - hdr_state = f"{'N':>5} | {'cuLA v state RMSE':>18} | {'rel':>10} | {'cuLA k state RMSE':>18} | {'rel':>10}" + hdr_state = ( + f"{'N':>5} | {'cuLA v state rel_rmse':>32} | {'rel_max':>10} | " + f"{'cuLA k state rel_rmse':>32} | {'rel_max':>10}" + ) print(hdr_state) print("-" * len(hdr_state)) for res in results: print( - f"{res['N']:5d} | {res['state_v_last_rmse']:18.3e} | {res['state_v_last_rel']:10.3e} | " - f"{res['state_k_last_rmse']:18.3e} | {res['state_k_last_rel']:10.3e}" + f"{res['N']:5d} | {res['state_v_last_relative_rms_error']:32.3e} | {res['state_v_last_rel_max']:10.3e} | " + f"{res['state_k_last_relative_rms_error']:32.3e} | {res['state_k_last_rel_max']:10.3e}" ) all_sections.append((h_dim, v_dim, results)) @@ -452,10 +441,12 @@ def print_section(h_dim: int, v_dim: int): print_section(h_dim, args.V) print() - gpu_type = normalize_gpu_type(gpu_name) - report_path = pathlib.Path(__file__).resolve().parent.parent / f"BENCHMARK_KDA_DECODE_{gpu_type}.md" - write_markdown_report(args, gpu_name, all_sections, report_path) - print(f"Markdown report written to: {report_path}") + if args.output is not None: + output_path = default_report_path(gpu_name) if args.output == "__AUTO__" else pathlib.Path(args.output) + write_markdown_report(args, gpu_name, all_sections, output_path) + print(f"Markdown report written to: {output_path.resolve()}") + + return args, gpu_name, all_sections if __name__ == "__main__": diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index 3953ccca..6a602767 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -22,7 +22,7 @@ - sm90 (Hopper) → cula.kda.hopper_fused_fwd.cula_kda_prefill Compares: - - Accuracy: RMSE, relative max diff between cuLA fully-fused and FLA Triton + - Accuracy: relative_rms_error, relative max diff between cuLA fully-fused and FLA Triton - Performance: kernel execution time (ms) with CUDA events Modes: @@ -53,9 +53,11 @@ from benchmarks.utils import ( SEED, + benchmark_cuda_mode_fn, build_varlen_configs, exclusive_cumsum, prepare_safe_gate_inputs, + relative_rms_error_rel_max_mean_abs, set_seed, ) from cula.utils import get_device_sm_version, get_kda_fused_fwd @@ -86,35 +88,6 @@ # ============================================================ # Helpers # ============================================================ -def time_kernel(fn, warmup=None, n_iters=None): - if warmup is None: - warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP - if n_iters is None: - n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start_evt = torch.cuda.Event(enable_timing=True) - end_evt = torch.cuda.Event(enable_timing=True) - start_evt.record() - for _ in range(n_iters): - fn() - end_evt.record() - torch.cuda.synchronize() - return start_evt.elapsed_time(end_evt) / n_iters - - -def accuracy_stats(ref, out): - """Compute RMSE, relative max diff, and mean absolute difference.""" - ref_f = ref.float() - out_f = out.float() - diff = (ref_f - out_f).abs() - rmse = diff.pow(2).mean().sqrt().item() - max_diff = diff.max().item() - denom = ref_f.abs().max().item() - rel_max = max_diff / denom if denom > 0 else 0.0 - mean_diff = diff.mean().item() - return rmse, rel_max, mean_diff def run_fla(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound): @@ -209,11 +182,23 @@ def bench_fixed(configs): o_cula, _ = run_cula(**common) torch.cuda.synchronize() - rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) # Performance - ms_fla = time_kernel(lambda: run_fla(**common)) - ms_cula = time_kernel(lambda: run_cula(**common)) + ms_fla = benchmark_cuda_mode_fn( + lambda: run_fla(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + lambda: run_cula(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") results.append( @@ -222,7 +207,7 @@ def bench_fixed(configs): "T": T, "H": H, "HV": HV, - "rmse": rmse, + "relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff, "ms_fla": ms_fla, @@ -288,11 +273,23 @@ def bench_varlen(configs): o_cula, _ = run_cula(**common) torch.cuda.synchronize() - rmse, rel_max, mean_diff = accuracy_stats(o_fla, o_cula) + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) # Performance - ms_fla = time_kernel(lambda: run_fla(**common)) - ms_cula = time_kernel(lambda: run_cula(**common)) + ms_fla = benchmark_cuda_mode_fn( + lambda: run_fla(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + lambda: run_cula(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") n_seqs = len(seq_lens) @@ -308,7 +305,7 @@ def bench_varlen(configs): "n_seqs": n_seqs, "H": H, "HV": HV, - "rmse": rmse, + "relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff, "ms_fla": ms_fla, @@ -345,7 +342,7 @@ def print_report(fixed_results, varlen_results): print(f" {'─' * 110}") print( f" {'B':>3s} {'T':>6s} {'H':>3s} {'HV':>3s} {'GVA':>4s} │ " - f"{'RMSE':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'rel_rmse':>18s} {'rel_max':>10s} {'mean_diff':>10s} │ " f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" ) print(f" {'─' * 110}") @@ -353,7 +350,7 @@ def print_report(fixed_results, varlen_results): gva_tag = f"{r['HV'] // r['H']}x" if r["HV"] > r["H"] else "no" print( f" {r['B']:3d} {r['T']:6d} {r['H']:3d} {r['HV']:3d} {gva_tag:>4s} │ " - f"{r['rmse']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " + f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 110}") @@ -363,7 +360,7 @@ def print_report(fixed_results, varlen_results): print(f" {'─' * 120}") print( f" {'Config':>45s} {'H':>3s} {'HV':>3s} {'GVA':>4s} │ " - f"{'RMSE':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'rel_rmse':>18s} {'rel_max':>10s} {'mean_diff':>10s} │ " f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" ) print(f" {'─' * 120}") @@ -371,7 +368,7 @@ def print_report(fixed_results, varlen_results): gva_tag = f"{r['HV'] // r['H']}x" if r["HV"] > r["H"] else "no" print( f" {r['tag']:>45s} {r['H']:3d} {r['HV']:3d} {gva_tag:>4s} │ " - f"{r['rmse']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " + f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 120}") diff --git a/benchmarks/bench_kda_fwd_bwd_e2e.py b/benchmarks/bench_kda_fwd_bwd_e2e.py index 24c67213..2a1f1d65 100644 --- a/benchmarks/bench_kda_fwd_bwd_e2e.py +++ b/benchmarks/bench_kda_fwd_bwd_e2e.py @@ -18,7 +18,7 @@ for chunk_kda forward + backward (end-to-end) Compares: - - Accuracy: err_ratio, relative max diff between cuLA and FLA outputs & gradients + - Accuracy: relative_rms_error, relative max diff between cuLA and FLA outputs & gradients - Performance: kernel execution time (ms) with CUDA events Modes: @@ -49,10 +49,12 @@ from benchmarks.utils import ( SEED, + benchmark_cuda_mode_fn, build_varlen_configs, exclusive_cumsum, generate_random_seq_lens, prepare_safe_gate_inputs, + relative_rms_error_rel_max_mean_abs, set_seed, ) from cula.kda import chunk_kda as cula_chunk_kda @@ -72,37 +74,6 @@ # ============================================================ # Helpers # ============================================================ -def time_kernel(fn, warmup=None, n_iters=None): - if warmup is None: - warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP - if n_iters is None: - n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start_evt = torch.cuda.Event(enable_timing=True) - end_evt = torch.cuda.Event(enable_timing=True) - start_evt.record() - for _ in range(n_iters): - fn() - end_evt.record() - torch.cuda.synchronize() - return start_evt.elapsed_time(end_evt) / n_iters - - -def accuracy_stats(ref, out): - """Compute err_ratio, relative max diff, and mean absolute difference.""" - ref_f = ref.float() - out_f = out.float() - diff = (ref_f - out_f).abs() - err = diff.flatten().pow(2).mean().sqrt().item() - base = ref_f.flatten().pow(2).mean().sqrt().item() - err_ratio = err / (base + 1e-8) - max_diff = diff.max().item() - denom = ref_f.abs().max().item() - rel_max = max_diff / denom if denom > 0 else 0.0 - mean_diff = diff.mean().item() - return err_ratio, rel_max, mean_diff def run_kda_e2e(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound, do, dht, fn): @@ -279,16 +250,18 @@ def bench_fixed(configs): torch.cuda.synchronize() for name in ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"): - err_ratio, rel_max, mean_diff = accuracy_stats(fla_results[name], cula_results[name]) - acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs( + fla_results[name], cula_results[name] + ) + acc[name] = {"relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff} else: # forward-only accuracy o_fla, ht_fla = run_kda_e2e(**common, fn=fla_chunk_kda) o_cula, ht_cula = run_kda_e2e(**common, fn=cula_chunk_kda) torch.cuda.synchronize() for name, ref, out in [("o", o_fla, o_cula), ("ht", ht_fla, ht_cula)]: - err_ratio, rel_max, mean_diff = accuracy_stats(ref, out) - acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(ref, out) + acc[name] = {"relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff} # For timing, use leaf tensors with requires_grad q_t = q.detach().clone().requires_grad_(True) @@ -320,8 +293,20 @@ def fn_fla(**kw): def fn_cula(**kw): return lambda: run_kda_e2e(**kw, fn=cula_chunk_kda) - ms_fla = time_kernel(fn_fla(**timing_common)) - ms_cula = time_kernel(fn_cula(**timing_common)) + ms_fla = benchmark_cuda_mode_fn( + fn_fla(**timing_common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + fn_cula(**timing_common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") r = { @@ -391,15 +376,17 @@ def bench_varlen(configs): torch.cuda.synchronize() for name in ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0"): - err_ratio, rel_max, mean_diff = accuracy_stats(fla_results[name], cula_results[name]) - acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs( + fla_results[name], cula_results[name] + ) + acc[name] = {"relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff} else: o_fla, ht_fla = run_kda_e2e(**common, fn=fla_chunk_kda) o_cula, ht_cula = run_kda_e2e(**common, fn=cula_chunk_kda) torch.cuda.synchronize() for name, ref, out in [("o", o_fla, o_cula), ("ht", ht_fla, ht_cula)]: - err_ratio, rel_max, mean_diff = accuracy_stats(ref, out) - acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(ref, out) + acc[name] = {"relative_rms_error": relative_rms_error, "rel_max": rel_max, "mean_diff": mean_diff} # For timing, use leaf tensors with requires_grad q_t = q.detach().clone().requires_grad_(True) @@ -431,8 +418,20 @@ def fn_fla(**kw): def fn_cula(**kw): return lambda: run_kda_e2e(**kw, fn=cula_chunk_kda) - ms_fla = time_kernel(fn_fla(**timing_common)) - ms_cula = time_kernel(fn_cula(**timing_common)) + ms_fla = benchmark_cuda_mode_fn( + fn_fla(**timing_common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + fn_cula(**timing_common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") n_seqs = len(seq_lens) @@ -493,15 +492,20 @@ def print_report(fixed_results, varlen_results): for r in fixed_results: rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in acc_keys) - err_ratio_vals = " ".join(f"{r['accuracy'].get(k, {}).get('err_ratio', 0.0):10.6f}" for k in acc_keys) + relative_rms_error_vals = " ".join( + f"{r['accuracy'].get(k, {}).get('relative_rms_error', 0.0):10.6f}" for k in acc_keys + ) # Line 1: timing + rel_max print( f" {r['B']:3d} {r['T']:5d} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x │ " f"{'rel_max:':>10s}{rel_max_vals}" ) - # Line 2: err_ratio (no timing columns) - print(f" {'':3s} {'':5s} │ {'':9s} {'':11s} {'':8s} │ {'err_ratio:':>10s}{err_ratio_vals}") + # Line 2: rel_rmse (no timing columns) + print( + f" {'':3s} {'':5s} │ {'':9s} {'':11s} {'':8s} │ " + f"{'rel_rmse:':>10s}{relative_rms_error_vals}" + ) print(f" {'─' * 125}") if varlen_results: @@ -513,15 +517,20 @@ def print_report(fixed_results, varlen_results): for r in varlen_results: rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in acc_keys) - err_ratio_vals = " ".join(f"{r['accuracy'].get(k, {}).get('err_ratio', 0.0):10.6f}" for k in acc_keys) + relative_rms_error_vals = " ".join( + f"{r['accuracy'].get(k, {}).get('relative_rms_error', 0.0):10.6f}" for k in acc_keys + ) # Line 1: timing + rel_max print( f" {r['tag']:>45s} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x │ " f"{'rel_max:':>10s}{rel_max_vals}" ) - # Line 2: err_ratio (no config/timing columns) - print(f" {'':>45s} │ {'':9s} {'':11s} {'':8s} │ {'err_ratio:':>10s}{err_ratio_vals}") + # Line 2: rel_rmse (no config/timing columns) + print( + f" {'':>45s} │ {'':9s} {'':11s} {'':8s} │ " + f"{'rel_rmse:':>10s}{relative_rms_error_vals}" + ) print(f" {'─' * 140}") print(f"\n{sep}\n") diff --git a/benchmarks/bench_la_decode_vs_fla.py b/benchmarks/bench_la_decode_vs_fla.py index 18a0b07d..850bca58 100644 --- a/benchmarks/bench_la_decode_vs_fla.py +++ b/benchmarks/bench_la_decode_vs_fla.py @@ -53,36 +53,13 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +from benchmarks.utils import benchmark_cuda_fn, relative_rms_error from fla.ops.common.fused_recurrent import fused_recurrent_fwd, fused_recurrent_fwd_kernel from cula.ops.la_decode import _get_compiled_kernel, linear_attention_decode from cula.utils import USE_FAST_MATH -# ───────────────────────────────────────────────────────────────────────────── -# Timing utility -# ───────────────────────────────────────────────────────────────────────────── -def benchmark_fn(fn, warmup=30, rep=200): - """Benchmark using CUDA events. Returns IQR-mean time in ms.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - - for i in range(rep): - starts[i].record() - fn() - ends[i].record() - - torch.cuda.synchronize() - times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends)) - n = len(times) - iqr = times[n // 4 : 3 * n // 4] - return sum(iqr) / len(iqr) - - # ───────────────────────────────────────────────────────────────────────────── # Core benchmark for one configuration # ───────────────────────────────────────────────────────────────────────────── @@ -147,12 +124,12 @@ def run_config(B, H, K, V, layer_idx, num_layers): # ── Correctness ──────────────────────────────────────────────────────── o_fla_cmp = o_fla.squeeze(1).float() o_cute_cmp = out_cute.float() - rmse = torch.sqrt(torch.mean((o_cute_cmp - o_fla_cmp) ** 2)).item() + output_relative_rms_error = relative_rms_error(o_fla_cmp, o_cute_cmp) max_ref = torch.abs(o_fla_cmp).max().item() rel_maxdiff = torch.abs(o_cute_cmp - o_fla_cmp).max().item() / (max_ref + 1e-8) state_cute_back = state_cute.reshape(B, H, V, K).permute(0, 1, 3, 2).contiguous() - state_rmse = torch.sqrt(torch.mean((state_cute_back - ht_fla.float()) ** 2)).item() + state_relative_rms_error = relative_rms_error(ht_fla, state_cute_back) # ================================================================== # Mode 1: KERNEL-ONLY (pre-allocated everything, minimal host overhead) @@ -209,8 +186,8 @@ def kernel_cute(): compiled_cute(cute_state_k, decay_scales, q_3d, k_3d, v_3d, out_cute_k, s_offsets, stream_handle) with torch.no_grad(): - kernel_fla_ms = benchmark_fn(kernel_fla) - kernel_cute_ms = benchmark_fn(kernel_cute) + kernel_fla_ms = benchmark_cuda_fn(kernel_fla) + kernel_cute_ms = benchmark_cuda_fn(kernel_cute) # ================================================================== # Mode 2: WRAPPER (full call path as used in production) @@ -251,8 +228,8 @@ def wrapper_cute(): ) with torch.no_grad(): - wrap_fla_ms = benchmark_fn(wrapper_fla) - wrap_cute_ms = benchmark_fn(wrapper_cute) + wrap_fla_ms = benchmark_cuda_fn(wrapper_fla) + wrap_cute_ms = benchmark_cuda_fn(wrapper_cute) return { "B": B, @@ -262,9 +239,9 @@ def wrapper_cute(): "wrap_fla_ms": wrap_fla_ms, "wrap_cute_ms": wrap_cute_ms, "wrap_speedup": wrap_fla_ms / wrap_cute_ms, - "rmse": rmse, + "output_relative_rms_error": output_relative_rms_error, "rel_maxdiff": rel_maxdiff, - "state_rmse": state_rmse, + "state_relative_rms_error": state_relative_rms_error, } @@ -299,7 +276,7 @@ def main(): print(f"{'=' * 100}") print( f"{'B':>5} | {'fla (ms)':>10} | {'cute (ms)':>10} | " - f"{'speedup':>8} | {'RMSE':>10} | {'Rel MaxDiff':>12} | {'State RMSE':>12}" + f"{'speedup':>8} | {'rel_rmse':>18} | {'Rel MaxDiff':>12} | {'State rel_rmse':>24}" ) print("─" * 90) @@ -309,8 +286,8 @@ def main(): results.append(r) print( f"{r['B']:>5} | {r['kernel_fla_ms']:>10.4f} | {r['kernel_cute_ms']:>10.4f} | " - f"{r['kernel_speedup']:>7.2f}x | {r['rmse']:>10.6f} | " - f"{r['rel_maxdiff']:>12.6f} | {r['state_rmse']:>12.8f}" + f"{r['kernel_speedup']:>7.2f}x | {r['output_relative_rms_error']:>18.6f} | " + f"{r['rel_maxdiff']:>12.6f} | {r['state_relative_rms_error']:>24.8f}" ) # ── Wrapper comparison ────────────────────────────────────────────── diff --git a/benchmarks/bench_lightning_attn.py b/benchmarks/bench_lightning_attn.py index 07733009..0534fe11 100644 --- a/benchmarks/bench_lightning_attn.py +++ b/benchmarks/bench_lightning_attn.py @@ -50,6 +50,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +from benchmarks.utils import gen_random, gen_skewed, gen_uniform, relative_rms_error, time_cuda_fn from fla.ops.simple_gla.chunk import chunk_simple_gla_fwd from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen @@ -105,53 +106,6 @@ def torch_naive_lightning_attn(Q, K, V, decay, scale=1.0, initial_state=None, ou return O, ht -# ============================================================================= -# Sequence length generators (for varlen) -# ============================================================================= -def gen_uniform(N, T): - """All sequences have equal length.""" - per = T // N - return [per] * N - - -def gen_skewed(N, T): - """One long sequence + many short ones.""" - if N == 1: - return [T] - short = max(1, T // (2 * (N - 1))) - long_len = T - short * (N - 1) - return [long_len] + [short] * (N - 1) - - -def gen_random(N, T, seed=42): - """Random sequence lengths summing to ~T.""" - rng = np.random.RandomState(seed) - raw = rng.dirichlet(np.ones(N)) - lens = np.maximum(1, np.round(raw * T).astype(int)) - diff = T - lens.sum() - lens[0] += diff - lens = np.maximum(1, lens) - return lens.tolist() - - -# ============================================================================= -# Timing helper -# ============================================================================= -def time_fn(fn, warmup, iters): - """Time a CUDA function using events. Returns ms/call.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - fn() - end.record() - torch.cuda.synchronize() - return start.elapsed_time(end) / iters - - # ============================================================================= # Runners # ============================================================================= @@ -173,7 +127,7 @@ def fn(): ) fn() # compile - ms = time_fn(fn, warmup, iters) + ms = time_cuda_fn(fn, warmup, iters) o, ht = fn() return o, ht, ms @@ -198,7 +152,7 @@ def fn(): fn() compile_ms = (time.time() - t0) * 1000 - ms = time_fn(fn, warmup, iters) + ms = time_cuda_fn(fn, warmup, iters) O, ht = fn() return O, ht, ms, compile_ms @@ -222,7 +176,7 @@ def fn(): fn() compile_ms = (time.time() - t0) * 1000 - ms = time_fn(fn, warmup, iters) + ms = time_cuda_fn(fn, warmup, iters) O, sp = fn() return ms, O, sp, compile_ms @@ -246,7 +200,7 @@ def fn(): ) fn() # compile - ms = time_fn(fn, warmup, iters) + ms = time_cuda_fn(fn, warmup, iters) return ms @@ -312,23 +266,18 @@ def benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, i for label, o_test, ht_test in [("fla", o_fla, ht_fla), ("cute", o_cute, ht_cute)]: if o_test is not None: diff = o_naive - o_test.float() - rms = o_naive.pow(2).mean().sqrt().item() - rmse = diff.pow(2).mean().sqrt().item() - result[f"{label}_o_rmse_ratio"] = rmse / (rms + 1e-8) + result[f"{label}_o_relative_rms_error"] = relative_rms_error(o_naive, o_test) result[f"{label}_o_maxdiff"] = diff.abs().max().item() if output_ht and ht_naive is not None and ht_test is not None: # CuTe kernel outputs BHVK state; transpose to BHKV for comparison ht_cmp = ht_test.transpose(-1, -2).float() if label == "cute" else ht_test.float() - diff_ht = ht_naive - ht_cmp - ht_rms = ht_naive.pow(2).mean().sqrt().item() - ht_rmse = diff_ht.pow(2).mean().sqrt().item() - result[f"{label}_ht_rmse_ratio"] = ht_rmse / (ht_rms + 1e-8) + result[f"{label}_ht_relative_rms_error"] = relative_rms_error(ht_naive, ht_cmp) else: - result[f"{label}_ht_rmse_ratio"] = float("nan") + result[f"{label}_ht_relative_rms_error"] = float("nan") else: - result[f"{label}_o_rmse_ratio"] = float("nan") + result[f"{label}_o_relative_rms_error"] = float("nan") result[f"{label}_o_maxdiff"] = float("nan") - result[f"{label}_ht_rmse_ratio"] = float("nan") + result[f"{label}_ht_relative_rms_error"] = float("nan") # --- Speedup --- fla_ok = _valid(result["fla_ms"]) @@ -397,22 +346,20 @@ def benchmark_varlen_config(N, seq_lens, H, D, warmup, iters, dist=""): diff_o = O_p.float() - O_np.float() result["p_vs_np_O_diff"] = diff_o.abs().max().item() o_rmse = diff_o.pow(2).mean().sqrt().item() - o_rms = O_p.float().pow(2).mean().sqrt().item() result["p_vs_np_O_rmse"] = o_rmse - result["p_vs_np_O_rmse_ratio"] = o_rmse / (o_rms + 1e-8) + result["p_vs_np_O_relative_rms_error"] = relative_rms_error(O_p, O_np) diff_ht = sp_p.float() - sp_np.float() result["p_vs_np_ht_diff"] = diff_ht.abs().max().item() ht_rmse = diff_ht.pow(2).mean().sqrt().item() - ht_rms = sp_p.float().pow(2).mean().sqrt().item() result["p_vs_np_ht_rmse"] = ht_rmse - result["p_vs_np_ht_rmse_ratio"] = ht_rmse / (ht_rms + 1e-8) + result["p_vs_np_ht_relative_rms_error"] = relative_rms_error(sp_p, sp_np) else: result["p_vs_np_O_diff"] = float("nan") result["p_vs_np_ht_diff"] = float("nan") result["p_vs_np_O_rmse"] = float("nan") - result["p_vs_np_O_rmse_ratio"] = float("nan") + result["p_vs_np_O_relative_rms_error"] = float("nan") result["p_vs_np_ht_rmse"] = float("nan") - result["p_vs_np_ht_rmse_ratio"] = float("nan") + result["p_vs_np_ht_relative_rms_error"] = float("nan") # --- Speedups --- p_ms = result["persistent_ms"] @@ -439,8 +386,8 @@ def print_standard_header(): hdr = ( f"{'Config':<28} {'Mode':<10} " f"{'FLA(ms)':>9} {'CuteDSL(ms)':>12} {'Speedup':>8} " - f"{'FLA_O_RMSE%':>12} {'Cute_O_RMSE%':>13} " - f"{'FLA_Ht_RMSE%':>13} {'Cute_Ht_RMSE%':>14}" + f"{'FLA_O_rel_rmse%':>26} {'Cute_O_rel_rmse%':>27} " + f"{'FLA_Ht_rel_rmse%':>27} {'Cute_Ht_rel_rmse%':>28}" ) print(hdr) print("-" * len(hdr)) @@ -452,10 +399,26 @@ def print_standard_result(r): fla = f"{r['fla_ms']:.3f}" if _valid(r.get("fla_ms", float("nan"))) else "ERR" dsl = f"{r['cutedsl_ms']:.3f}" if _valid(r.get("cutedsl_ms", float("nan"))) else "ERR" sp = f"{r['speedup']:.2f}x" if _valid(r.get("speedup", float("nan"))) else "-" - fla_o = f"{r['fla_o_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("fla_o_rmse_ratio", float("nan"))) else "-" - cute_o = f"{r['cute_o_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("cute_o_rmse_ratio", float("nan"))) else "-" - fla_ht = f"{r['fla_ht_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("fla_ht_rmse_ratio", float("nan"))) else "-" - cute_ht = f"{r['cute_ht_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("cute_ht_rmse_ratio", float("nan"))) else "-" + fla_o = ( + f"{r['fla_o_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("fla_o_relative_rms_error", float("nan"))) + else "-" + ) + cute_o = ( + f"{r['cute_o_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("cute_o_relative_rms_error", float("nan"))) + else "-" + ) + fla_ht = ( + f"{r['fla_ht_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("fla_ht_relative_rms_error", float("nan"))) + else "-" + ) + cute_ht = ( + f"{r['cute_ht_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("cute_ht_relative_rms_error", float("nan"))) + else "-" + ) print(f"{cfg:<28} {r['mode']:<10} {fla:>9} {dsl:>12} {sp:>8} {fla_o:>12} {cute_o:>13} {fla_ht:>13} {cute_ht:>14}") if r.get("fla_err"): @@ -469,7 +432,7 @@ def print_varlen_header(): f"{'Config':<24} {'Dist':<8} " f"{'Persist(ms)':>12} {'NonPer(ms)':>11} {'FLA_vl(ms)':>11} " f"{'P/NP':>6} {'P/FLAvl':>8} " - f"{'O diff':>10} {'O_RMSE%':>9} {'ht diff':>10} {'ht_RMSE%':>10}" + f"{'O diff':>10} {'O_rel_rmse%':>22} {'ht diff':>10} {'ht_rel_rmse%':>23}" ) print(hdr) print("-" * len(hdr)) @@ -487,9 +450,17 @@ def print_varlen_result(r): pvfla_vl = f"{r['p_vs_fla_vl_speedup']:.2f}x" if _valid(r.get("p_vs_fla_vl_speedup", float("nan"))) else "-" od = f"{r['p_vs_np_O_diff']:.1e}" if not np.isnan(r.get("p_vs_np_O_diff", float("nan"))) else "-" - ormse = f"{r['p_vs_np_O_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("p_vs_np_O_rmse_ratio", float("nan"))) else "-" + ormse = ( + f"{r['p_vs_np_O_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("p_vs_np_O_relative_rms_error", float("nan"))) + else "-" + ) hd = f"{r['p_vs_np_ht_diff']:.1e}" if not np.isnan(r.get("p_vs_np_ht_diff", float("nan"))) else "-" - htrmse = f"{r['p_vs_np_ht_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("p_vs_np_ht_rmse_ratio", float("nan"))) else "-" + htrmse = ( + f"{r['p_vs_np_ht_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("p_vs_np_ht_relative_rms_error", float("nan"))) + else "-" + ) print( f"{cfg:<24} {dist:<8} {p_ms:>12} {np_ms:>11} {fla_vl:>11} {pvnp:>6} {pvfla_vl:>8} {od:>10} {ormse:>9} {hd:>10} {htrmse:>10}" @@ -619,10 +590,12 @@ def run_benchmark_suite(args): if od: print(f" P vs NP O diff: max={max(od):.2e} (bit-exact={all(x == 0 for x in od)})") ormse = [ - r["p_vs_np_O_rmse_ratio"] * 100 for r in mode_r if not np.isnan(r.get("p_vs_np_O_rmse_ratio", float("nan"))) + r["p_vs_np_O_relative_rms_error"] * 100 + for r in mode_r + if not np.isnan(r.get("p_vs_np_O_relative_rms_error", float("nan"))) ] if ormse: - print(f" P vs NP O RMSE ratio: avg={np.mean(ormse):.4f}% max={np.max(ormse):.4f}%") + print(f" P vs NP O relative_rms_error: avg={np.mean(ormse):.4f}% max={np.max(ormse):.4f}%") else: speedups = [r["speedup"] for r in mode_r] print(f"\n [{mode}] ({len(mode_r)} configs)") @@ -631,19 +604,19 @@ def run_benchmark_suite(args): ) for label, name in [("fla", "FLA"), ("cute", "CuteDSL")]: o_rmses = [ - r[f"{label}_o_rmse_ratio"] * 100 + r[f"{label}_o_relative_rms_error"] * 100 for r in mode_r - if not np.isnan(r.get(f"{label}_o_rmse_ratio", float("nan"))) + if not np.isnan(r.get(f"{label}_o_relative_rms_error", float("nan"))) ] if o_rmses: - print(f" {name} O RMSE% (vs naive): avg={np.mean(o_rmses):.4f} max={np.max(o_rmses):.4f}") + print(f" {name} O relative_rms_error% (vs naive): avg={np.mean(o_rmses):.4f} max={np.max(o_rmses):.4f}") ht_rmses = [ - r[f"{label}_ht_rmse_ratio"] * 100 + r[f"{label}_ht_relative_rms_error"] * 100 for r in mode_r - if not np.isnan(r.get(f"{label}_ht_rmse_ratio", float("nan"))) + if not np.isnan(r.get(f"{label}_ht_relative_rms_error", float("nan"))) ] if ht_rmses: - print(f" {name} Ht RMSE% (vs naive): avg={np.mean(ht_rmses):.4f} max={np.max(ht_rmses):.4f}") + print(f" {name} Ht relative_rms_error% (vs naive): avg={np.mean(ht_rmses):.4f} max={np.max(ht_rmses):.4f}") # --- Plot --- if args.plot: @@ -767,36 +740,38 @@ def generate_report(all_results, modes, args): has_ht = mode == "h0_ht" if has_ht: f.write( - "| Config | FLA(ms) | CuteDSL(ms) | Speedup | FLA_O_RMSE% | Cute_O_RMSE% | FLA_Ht_RMSE% | Cute_Ht_RMSE% |\n" + "| Config | FLA(ms) | CuteDSL(ms) | Speedup | FLA_O_rel_rmse% | Cute_O_rel_rmse% | FLA_Ht_rel_rmse% | Cute_Ht_rel_rmse% |\n" ) f.write( - "|--------|---------|-------------|---------|-------------|--------------|--------------|---------------|\n" + "|--------|---------|-------------|---------|---------------------------|----------------------------|----------------------------|-----------------------------|\n" ) else: - f.write("| Config | FLA(ms) | CuteDSL(ms) | Speedup | FLA_O_RMSE% | Cute_O_RMSE% |\n") - f.write("|--------|---------|-------------|---------|-------------|---------------|\n") + f.write("| Config | FLA(ms) | CuteDSL(ms) | Speedup | FLA_O_rel_rmse% | Cute_O_rel_rmse% |\n") + f.write("|--------|---------|-------------|---------|---------------------------|----------------------------|\n") for r in mr: cfg = f"B={r['B']},T={r['T']},H={r['H']}" sp = f"{r['speedup']:.2f}x" if _valid(r.get("speedup", float("nan"))) else "-" fla = f"{r['fla_ms']:.3f}" if _valid(r.get("fla_ms", float("nan"))) else "-" dsl = f"{r['cutedsl_ms']:.3f}" if _valid(r.get("cutedsl_ms", float("nan"))) else "-" fla_o = ( - f"{r['fla_o_rmse_ratio'] * 100:.3f}%" if not np.isnan(r.get("fla_o_rmse_ratio", float("nan"))) else "-" + f"{r['fla_o_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("fla_o_relative_rms_error", float("nan"))) + else "-" ) cute_o = ( - f"{r['cute_o_rmse_ratio'] * 100:.3f}%" - if not np.isnan(r.get("cute_o_rmse_ratio", float("nan"))) + f"{r['cute_o_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("cute_o_relative_rms_error", float("nan"))) else "-" ) if has_ht: fla_ht = ( - f"{r['fla_ht_rmse_ratio'] * 100:.3f}%" - if not np.isnan(r.get("fla_ht_rmse_ratio", float("nan"))) + f"{r['fla_ht_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("fla_ht_relative_rms_error", float("nan"))) else "-" ) cute_ht = ( - f"{r['cute_ht_rmse_ratio'] * 100:.3f}%" - if not np.isnan(r.get("cute_ht_rmse_ratio", float("nan"))) + f"{r['cute_ht_relative_rms_error'] * 100:.3f}%" + if not np.isnan(r.get("cute_ht_relative_rms_error", float("nan"))) else "-" ) f.write(f"| {cfg} | {fla} | {dsl} | {sp} | {fla_o} | {cute_o} | {fla_ht} | {cute_ht} |\n") @@ -825,12 +800,14 @@ def generate_report(all_results, modes, args): ) for label, name in [("fla", "FLA"), ("cute", "CuteDSL")]: o_rmses = [ - r[f"{label}_o_rmse_ratio"] * 100 + r[f"{label}_o_relative_rms_error"] * 100 for r in mr - if not np.isnan(r.get(f"{label}_o_rmse_ratio", float("nan"))) + if not np.isnan(r.get(f"{label}_o_relative_rms_error", float("nan"))) ] if o_rmses: - f.write(f" - {name} O RMSE% (vs naive): avg {np.mean(o_rmses):.4f}, max {np.max(o_rmses):.4f}\n") + f.write( + f" - {name} O relative_rms_error% (vs naive): avg {np.mean(o_rmses):.4f}, max {np.max(o_rmses):.4f}\n" + ) f.write("\n---\n*Generated by bench_lightning_attn.py*\n") print(f"\nReport saved to {path}") diff --git a/benchmarks/bench_linear_attn.py b/benchmarks/bench_linear_attn.py index 08bbb520..5079b7ce 100644 --- a/benchmarks/bench_linear_attn.py +++ b/benchmarks/bench_linear_attn.py @@ -21,7 +21,6 @@ from cutlass.cute.runtime import from_dlpack from einops import rearrange from fla.ops.linear_attn import fused_chunk_linear_attn -from fla.ops.linear_attn.utils import normalize_output # from fla.ops.linear_attn.naive import naive_recurrent_linear_attn from fla.utils import assert_close, device @@ -33,6 +32,30 @@ PRINT_DEBUG = False +def normalize_output(q: torch.Tensor, k: torch.Tensor, o: torch.Tensor) -> torch.Tensor: + """Backward-compatible normalization for old linear-attn benchmark paths. + + Supports both the historical `[B, T, H, D]` layout and the chunked + `[B, H, N, C, D]` layout used by `naive_chunk_linear_attn`. + """ + if q.ndim == 4: + k_cum = k.cumsum(1) + z = (q * k_cum).sum(-1, keepdim=True) + return o / (z + 1e-10) + + if q.ndim == 5: + batch, heads, num_chunks, chunk_size, depth = q.shape + q_flat = q.reshape(batch, heads, num_chunks * chunk_size, depth) + k_flat = k.reshape(batch, heads, num_chunks * chunk_size, depth) + o_flat = o.reshape(batch, heads, num_chunks * chunk_size, o.shape[-1]) + k_cum = k_flat.cumsum(2) + z = (q_flat * k_cum).sum(-1, keepdim=True) + o_flat = o_flat / (z + 1e-10) + return o_flat.reshape_as(o) + + raise ValueError(f"Unsupported normalize_output layout with ndim={q.ndim}") + + def print_chunkwise(t, name): if not PRINT_DEBUG: return diff --git a/benchmarks/bench_recompute_wu.py b/benchmarks/bench_recompute_wu.py index c40dac31..b67d1c07 100644 --- a/benchmarks/bench_recompute_wu.py +++ b/benchmarks/bench_recompute_wu.py @@ -28,17 +28,23 @@ import sys import torch -import triton sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra from fla.ops.kda.wy_fast import recompute_w_u_fwd as fla_recompute_w_u_fwd -from fla.utils import get_abs_err, get_err_ratio import cula.cudac as cula_cuda -from benchmarks.utils import SEED, exclusive_cumsum, generate_random_seq_lens, prepare_intra_inputs +from benchmarks.utils import ( + SEED, + exclusive_cumsum, + generate_random_seq_lens, + prepare_intra_inputs, + relative_rms_error, + relative_rms_error_rel_max_mean_abs_rhs, + triton_bench_fn, +) from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra # Constant params @@ -55,17 +61,8 @@ DISABLE_RECOMPUTE = False # Whether to disable recompute (compute QG in forward) -def accuracy_stats(a, b): - """Compute RMSE, relative max diff, and mean absolute difference.""" - a, b = a.float(), b.float() - diff = a - b - rmse = diff.pow(2).mean().sqrt().item() - max_diff = diff.abs().max().item() - denom = b.abs().max().item() - rel_max = max_diff / denom if denom > 0 else 0.0 - mean_diff = diff.abs().mean().item() - return rmse, rel_max, mean_diff - +def get_abs_err(ref: torch.Tensor, out: torch.Tensor) -> float: + return (ref.float() - out.float()).abs().max().item() def prepare_recompute_wu_inputs(B, T, device, cu_seqlens=None, chunk_size=BT): """Prepare inputs for recompute_w_u benchmarking (handles both MHA and GVA). @@ -131,7 +128,7 @@ def benchmark_recompute_wu_uniform(): ) print("=" * 100) print( - f"{'B':>4} {'T':>7} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'B':>4} {'T':>7} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 100) @@ -156,22 +153,23 @@ def benchmark_recompute_wu_uniform(): ("w", w_fla, w_cula), ("u", u_fla, u_cula), ("qg", qg_fla, qg_cula), ("kg", kg_fla, kg_cula), ]: if t_fla is not None and t_cula is not None: - stats[name] = accuracy_stats(t_fla, t_cula) - rmse = max(s[0] for s in stats.values()) + stats[name] = relative_rms_error_rel_max_mean_abs_rhs(t_fla, t_cula) + # Use max across all outputs for display. + relative_rms_error_value = max(s[0] for s in stats.values()) rel_max = max(s[1] for s in stats.values()) mean_diff = max(s[2] for s in stats.values()) # Performance - ms_fla = triton.testing.do_bench( + ms_fla = triton_bench_fn( lambda: run_fla_recompute_wu(k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, DISABLE_RECOMPUTE), ) - ms_cula = triton.testing.do_bench( + ms_cula = triton_bench_fn( lambda: run_cula_recompute_wu(k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, chunk_size, DISABLE_RECOMPUTE), ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{B:>4} {T:>7} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{B:>4} {T:>7} │ {relative_rms_error_value:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 100) @@ -195,7 +193,7 @@ def benchmark_recompute_wu_varlen(): ) print("=" * 110) print( - f"{'total_len':>10} │ {'RMSE':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'total_len':>10} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 110) @@ -221,22 +219,23 @@ def benchmark_recompute_wu_varlen(): ("w", w_fla, w_cula), ("u", u_fla, u_cula), ("qg", qg_fla, qg_cula), ("kg", kg_fla, kg_cula), ]: if t_fla is not None and t_cula is not None: - stats[name] = accuracy_stats(t_fla, t_cula) - rmse = max(s[0] for s in stats.values()) + stats[name] = relative_rms_error_rel_max_mean_abs_rhs(t_fla, t_cula) + # Use max across all outputs for display. + relative_rms_error_value = max(s[0] for s in stats.values()) rel_max = max(s[1] for s in stats.values()) mean_diff = max(s[2] for s in stats.values()) # Performance - ms_fla = triton.testing.do_bench( + ms_fla = triton_bench_fn( lambda: run_fla_recompute_wu(k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, DISABLE_RECOMPUTE), ) - ms_cula = triton.testing.do_bench( + ms_cula = triton_bench_fn( lambda: run_cula_recompute_wu(k, v, beta, Akk, q, g, cu_seqlens, chunk_indices, chunk_size, DISABLE_RECOMPUTE), ) speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{total_len:>10} │ {rmse:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{total_len:>10} │ {relative_rms_error_value:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 110) @@ -265,19 +264,19 @@ def check_determinism(num_seqs=NUM_SEQS, T=2001, H=H, iters=1000): if not torch.equal(w, ref_w): print(f"Iteration {i}: w mismatch") - print(f"{get_abs_err(ref_w, w):.6f} absolute error, {get_err_ratio(ref_w, w):.6f} relative error") + print(f"{get_abs_err(ref_w, w):.6f} absolute error, {relative_rms_error(ref_w, w):.6f} relative_rms_error") raise AssertionError("Non-deterministic output detected in w") if not torch.equal(u, ref_u): print(f"Iteration {i}: u mismatch") - print(f"{get_abs_err(ref_u, u):.6f} absolute error, {get_err_ratio(ref_u, u):.6f} relative error") + print(f"{get_abs_err(ref_u, u):.6f} absolute error, {relative_rms_error(ref_u, u):.6f} relative_rms_error") raise AssertionError("Non-deterministic output detected in u") if kg is not None and not torch.equal(kg, ref_kg): print(f"Iteration {i}: kg mismatch") - print(f"{get_abs_err(ref_kg, kg):.6f} absolute error, {get_err_ratio(ref_kg, kg):.6f} relative error") + print(f"{get_abs_err(ref_kg, kg):.6f} absolute error, {relative_rms_error(ref_kg, kg):.6f} relative_rms_error") raise AssertionError("Non-deterministic output detected in kg") if qg is not None and not torch.equal(qg, ref_qg): print(f"Iteration {i}: qg mismatch") - print(f"{get_abs_err(ref_qg, qg):.6f} absolute error, {get_err_ratio(ref_qg, qg):.6f} relative error") + print(f"{get_abs_err(ref_qg, qg):.6f} absolute error, {relative_rms_error(ref_qg, qg):.6f} relative_rms_error") raise AssertionError("Non-deterministic output detected in qg") diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 64591662..485dc0d7 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -60,6 +60,170 @@ def set_seed(seed: int): torch.cuda.manual_seed(seed) +def benchmark_cuda_fn(fn, *, setup_fn=None, warmup=30, rep=200, aggregate="iqr_mean"): + """Benchmark a CUDA callable with events and return milliseconds per call.""" + for _ in range(warmup): + if setup_fn is not None: + setup_fn() + fn() + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + + for i in range(rep): + if setup_fn is not None: + setup_fn() + starts[i].record() + fn() + ends[i].record() + + torch.cuda.synchronize() + times = [s.elapsed_time(e) for s, e in zip(starts, ends)] + if not times: + return 0.0 + if aggregate == "mean": + return sum(times) / len(times) + if aggregate == "iqr_mean": + times = sorted(times) + if len(times) <= 2: + return sum(times) / len(times) + iqr = times[len(times) // 4 : 3 * len(times) // 4] + return sum(iqr) / len(iqr) + raise ValueError(f"Unsupported aggregate={aggregate}") + + +def resolve_benchmark_repeats(default_warmup, default_rep, *, ncu_mode=False, sanitizer_mode=False): + """Resolve benchmark warmup and repeat counts for normal vs profiling runs.""" + if ncu_mode or sanitizer_mode: + return 1, 1 + return default_warmup, default_rep + + +def benchmark_cuda_mode_fn( + fn, + *, + default_warmup, + default_rep, + ncu_mode=False, + sanitizer_mode=False, + setup_fn=None, +): + """Benchmark a CUDA callable using standard repo warmup/repeat mode rules.""" + warmup, rep = resolve_benchmark_repeats( + default_warmup, + default_rep, + ncu_mode=ncu_mode, + sanitizer_mode=sanitizer_mode, + ) + return benchmark_cuda_fn(fn, setup_fn=setup_fn, warmup=warmup, rep=rep, aggregate="mean") + + +def triton_bench_fn(fn, **kwargs): + """Benchmark a callable with Triton's do_bench helper.""" + import triton + + return triton.testing.do_bench(fn, **kwargs) + + +def time_cuda_fn(fn, warmup, iters): + """Time a CUDA callable and return milliseconds per call.""" + return benchmark_cuda_fn(fn, warmup=warmup, rep=iters, aggregate="mean") + + +def _error_stats(ref: torch.Tensor, out: torch.Tensor): + """Return shared float-cast tensors and basic absolute/RMS error stats.""" + ref_f = ref.float() + out_f = out.float() + diff = (ref_f - out_f).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + rmse = diff.pow(2).mean().sqrt().item() + ref_rms = ref_f.pow(2).mean().sqrt().item() + return ref_f, out_f, max_diff, mean_diff, rmse, ref_rms + + +def _relative_max(max_diff: float, denom: float): + return max_diff / denom if denom > 0 else 0.0 + + +def rmse_rel_max(ref: torch.Tensor, out: torch.Tensor): + """Return RMSE and relative max error between two tensors.""" + ref_f, _out_f, max_diff, _mean_diff, rmse, _ref_rms = _error_stats(ref, out) + rel_max = _relative_max(max_diff, ref_f.abs().max().item()) + return rmse, rel_max + + +def relative_rms_error(ref: torch.Tensor, out: torch.Tensor): + """Return relative RMS error between two tensors.""" + _ref_f, _out_f, _max_diff, _mean_diff, rmse, ref_rms = _error_stats(ref, out) + return rmse / (ref_rms + 1e-8) + + +def relative_rms_error_rel_max(ref: torch.Tensor, out: torch.Tensor): + """Return relative RMS error and relative max error.""" + ref_f, _out_f, max_diff, _mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + relative_rms = relative_rms_error(ref, out) + rel_max = _relative_max(max_diff, ref_f.abs().max().item()) + return relative_rms, rel_max + + +def rmse_rel_max_mean_abs(ref: torch.Tensor, out: torch.Tensor): + """Return RMSE, relative max error, and mean absolute difference.""" + ref_f, _out_f, max_diff, mean_diff, rmse, _ref_rms = _error_stats(ref, out) + rel_max = _relative_max(max_diff, ref_f.abs().max().item()) + return rmse, rel_max, mean_diff + + +def rmse_rel_max_mean_abs_rhs(ref: torch.Tensor, out: torch.Tensor): + """Return RMSE, relative max error vs rhs magnitude, and mean absolute difference.""" + _ref_f, out_f, max_diff, mean_diff, rmse, _ref_rms = _error_stats(ref, out) + rel_max = _relative_max(max_diff, out_f.abs().max().item()) + return rmse, rel_max, mean_diff + + +def relative_rms_error_rel_max_mean_abs(ref: torch.Tensor, out: torch.Tensor): + """Return relative RMS error, relative max error, and mean absolute difference.""" + ref_f, _out_f, max_diff, mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + relative_rms = relative_rms_error(ref, out) + rel_max = _relative_max(max_diff, ref_f.abs().max().item()) + return relative_rms, rel_max, mean_diff + + +def relative_rms_error_rel_max_mean_abs_rhs(ref: torch.Tensor, out: torch.Tensor): + """Return relative RMS error, rhs-relative max error, and mean absolute difference.""" + _ref_f, out_f, max_diff, mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + relative_rms = relative_rms_error(ref, out) + rel_max = _relative_max(max_diff, out_f.abs().max().item()) + return relative_rms, rel_max, mean_diff + + +def relative_rms_error_max_mean_abs(ref: torch.Tensor, out: torch.Tensor): + """Return relative RMS error, max error, and mean absolute difference.""" + _ref_f, _out_f, max_diff, mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + return relative_rms_error(ref, out), max_diff, mean_diff + + +def relative_rms_error_max_rel_mean_abs(ref: torch.Tensor, out: torch.Tensor): + """Return relative RMS error, max error, relative max error, and mean absolute difference.""" + ref_f, _out_f, max_diff, mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + rel_max_diff = _relative_max(max_diff, ref_f.abs().max().item()) + return relative_rms_error(ref, out), max_diff, rel_max_diff, mean_diff + + +def max_mean_abs_diff(ref: torch.Tensor, out: torch.Tensor): + """Return max and mean absolute difference.""" + _ref_f, _out_f, max_diff, mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + return max_diff, mean_diff + + +def max_rel_mean_abs_diff(ref: torch.Tensor, out: torch.Tensor): + """Return max error, relative max error, and mean absolute difference.""" + ref_f, _out_f, max_diff, mean_diff, _rmse, _ref_rms = _error_stats(ref, out) + rel_max_diff = _relative_max(max_diff, ref_f.abs().max().item()) + return max_diff, rel_max_diff, mean_diff + + def exclusive_cumsum(a: list[int]): r = [0] for v in a: diff --git a/pyproject.toml b/pyproject.toml index ff2bf950..b70e04fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ dependencies = [ "nvidia-cutlass-dsl==4.4.2", "apache-tvm-ffi==0.1.9", ] -license = "Apache-2.0" +license = { text = "Apache-2.0" } [project.optional-dependencies] dev = [ From 4571c7d2e4724556562d8394a907c8229aa84d91 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Mon, 25 May 2026 12:49:36 +0800 Subject: [PATCH 19/34] [KDA] support GVA for SM100 end-to-end (#73) * upgrade fla and update b200 bench, update readme and fix lightning test param * update h200 bench result with fla bug fixed * update b200 bench * update b200 bench * fix readme * remove useless repeat_interleave for fla * fix readme * gva for delta_h * bench for delta_h * fix ref delta_h * add gva for fwd_o * code lint * add gva test * integrate with fla v0.5.0 and pass e2e gva tests * add check * delete * refactor * add sanity check * delete useless tests * add gva for e2e bench * code lint * adjust print layout --------- Co-authored-by: boyu.zbw --- benchmarks/bench_chunk_delta_h.py | 84 ++- benchmarks/bench_fwd_o.py | 90 ++- benchmarks/bench_kda.py | 36 +- benchmarks/bench_kda_chunk_intra.py | 41 +- benchmarks/bench_kda_decode.py | 16 +- benchmarks/bench_kda_fused_fwd.py | 8 +- benchmarks/bench_kda_fwd_bwd_e2e.py | 38 +- benchmarks/bench_la_decode_vs_fla.py | 8 +- benchmarks/bench_lightning_attn.py | 18 +- benchmarks/bench_recompute_wu.py | 38 +- benchmarks/generate_benchmark_hopper_md.py | 27 +- benchmarks/generate_benchmark_md.py | 23 +- benchmarks/utils.py | 4 +- csrc/api/kda_sm100.cu | 4 +- csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp | 22 +- .../sm100/kda_fwd_intra_mainloop_sm100.hpp | 2 +- .../kda_fwd_recomp_w_u_mainloop_sm100.hpp | 5 +- csrc/kda/sm100/tile_scheduler.hpp | 6 +- cula/kda/blackwell_fused_fwd.py | 13 +- cula/kda/chunk.py | 13 +- cula/kda/chunk_bwd.py | 126 ++-- cula/kda/chunk_intra.py | 572 ++---------------- cula/ops/chunk_delta_h_sm100.py | 247 ++++---- cula/ops/fwd_o_sm100.py | 194 +++--- tests/test_kda.py | 76 ++- tests/test_kda_compare_fla.py | 76 ++- tests/test_kda_gva_intra_sm100.py | 386 ------------ 27 files changed, 803 insertions(+), 1370 deletions(-) delete mode 100644 tests/test_kda_gva_intra_sm100.py diff --git a/benchmarks/bench_chunk_delta_h.py b/benchmarks/bench_chunk_delta_h.py index d4761dc4..ad26a9b3 100644 --- a/benchmarks/bench_chunk_delta_h.py +++ b/benchmarks/bench_chunk_delta_h.py @@ -46,6 +46,7 @@ import numpy as np import torch + from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_mean_abs # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── @@ -82,20 +83,20 @@ def bench_non_varlen(configs): print("=" * 80) results = [] - for B, T, H, use_gk, use_h0, store_ht, save_vnew in configs: + for B, T, H, HV, use_gk, use_h0, store_ht, save_vnew in configs: torch.manual_seed(42) torch.cuda.empty_cache() k = torch.randn(B, T, H, K, device=device, dtype=dtype) * 0.1 - w = torch.randn(B, T, H, K, device=device, dtype=dtype) * 0.1 - u = torch.randn(B, T, H, V, device=device, dtype=dtype) * 0.1 + w = torch.randn(B, T, HV, K, device=device, dtype=dtype) * 0.1 + u = torch.randn(B, T, HV, V, device=device, dtype=dtype) * 0.1 gk = None h0 = None if use_gk: - gk = -torch.abs(torch.randn(B, T, H, K, device=device, dtype=torch.float32) * 0.1).cumsum(dim=1) + gk = -torch.abs(torch.randn(B, T, HV, K, device=device, dtype=torch.float32) * 0.1).cumsum(dim=1) if use_h0: - h0 = torch.randn(B, H, K, V, device=device, dtype=torch.float32) * 0.01 + h0 = torch.randn(B, HV, K, V, device=device, dtype=torch.float32) * 0.01 # ---- FLA baseline ---- fla_result = fla_fwd_h( @@ -170,10 +171,13 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0): flags.append("vn") flag_str = f" [{','.join(flags)}]" if flags else "" + hv_str = f"/{HV}" if HV != H else "" r = { "B": B, "T": T, "H": H, + "HV": HV, + "hv_str": hv_str, "flags": flag_str, "relative_rms_error": relative_rms_error, "max_diff": max_diff, @@ -184,7 +188,7 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0): } results.append(r) print( - f" B={B:2d} T={T:5d} H={H:3d}{flag_str:<16s} | " + f" B={B:2d} T={T:5d} H={H:3d}{hv_str:<4s}{flag_str:<16s} | " f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" @@ -222,7 +226,7 @@ def bench_varlen(configs): print("=" * 80) results = [] - for num_seqs, total_T, H, ratio, use_gk, use_h0, store_ht, save_vnew in configs: + for num_seqs, total_T, H, HV, ratio, use_gk, use_h0, store_ht, save_vnew in configs: seq_lens = generate_seq_lens(num_seqs, total_T, ratio) cu_seqlens_list = [0] for sl in seq_lens: @@ -240,22 +244,22 @@ def bench_varlen(configs): torch.manual_seed(42) torch.cuda.empty_cache() - # Both FLA and CuTe DSL use [1, total_T, H, ...] (4D with B=1) + # Both FLA and CuTe DSL use [1, total_T, H/HV, ...] (4D with B=1) k = torch.randn(1, total_T, H, K, device=device, dtype=dtype) * 0.1 - w = torch.randn(1, total_T, H, K, device=device, dtype=dtype) * 0.1 - u = torch.randn(1, total_T, H, V, device=device, dtype=dtype) * 0.1 + w = torch.randn(1, total_T, HV, K, device=device, dtype=dtype) * 0.1 + u = torch.randn(1, total_T, HV, V, device=device, dtype=dtype) * 0.1 gk = None h0 = None if use_gk: - gk_raw = torch.randn(1, total_T, H, K, device=device, dtype=torch.float32) * 0.1 + gk_raw = torch.randn(1, total_T, HV, K, device=device, dtype=torch.float32) * 0.1 gk = torch.zeros_like(gk_raw) for i in range(num_seqs): bos = cu_seqlens[i].item() eos = cu_seqlens[i + 1].item() gk[:, bos:eos] = -torch.abs(gk_raw[:, bos:eos]).cumsum(dim=1) if use_h0: - h0 = torch.randn(num_seqs, H, K, V, device=device, dtype=torch.float32) * 0.01 + h0 = torch.randn(num_seqs, HV, K, V, device=device, dtype=torch.float32) * 0.01 # ---- FLA baseline ---- fla_result = fla_fwd_h( @@ -338,10 +342,13 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0, cu=cu_seqlens): flags.append("vn") flag_str = f" [{','.join(flags)}]" if flags else "" + hv_str = f"/{HV}" if HV != H else "" r = { "tag": tag, "T_total": total_T, "H": H, + "HV": HV, + "hv_str": hv_str, "n_seqs": num_seqs, "flags": flag_str, "relative_rms_error": relative_rms_error, @@ -353,7 +360,7 @@ def run_cute(k=k, w=w, u=u, gk=gk, h0=h0, cu=cu_seqlens): } results.append(r) print( - f" {tag:40s} H={H:3d}{flag_str:<16s} | " + f" {tag:40s} H={H:3d}{hv_str:<4s}{flag_str:<16s} | " f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" @@ -381,15 +388,15 @@ def print_report(nv_results, vl_results): print("\n [Non-Varlen]") print(f" {'─' * 100}") print( - f" {'Config':<35s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'mean_diff':>12s}" + f" {'Config':<35s} │ {'rel_rmse':>10s} {'max_diff':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 100}") for r in nv_results: - label = f"B={r['B']:2d} T={r['T']:5d} H={r['H']:3d}{r['flags']}" + label = f"B={r['B']:2d} T={r['T']:5d} H={r['H']:3d}{r['hv_str']}{r['flags']}" print( f" {label:<35s} │ " - f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f"{r['relative_rms_error']:10.6f} {r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 100}") @@ -401,15 +408,15 @@ def print_report(nv_results, vl_results): print("\n [Varlen]") print(f" {'─' * 115}") print( - f" {'Config':>55s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'mean_diff':>12s}" + f" {'Config':>55s} │ {'rel_rmse':>10s} {'max_diff':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 115}") for r in vl_results: - label = f"{r['tag']} H={r['H']:3d}{r['flags']}" + label = f"{r['tag']} H={r['H']:3d}{r['hv_str']}{r['flags']}" print( f" {label:>55s} │ " - f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f"{r['relative_rms_error']:10.6f} {r['max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 115}") @@ -432,6 +439,18 @@ def main(): choices=["non-varlen", "varlen", "both"], help="Which benchmark mode to run (default: both)", ) + parser.add_argument( + "--heads", + type=int, + default=64, + help="Number of QK heads H (default: 64)", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of value heads HV (default: same as --heads, i.e. no GVA)", + ) parser.add_argument( "--ncu", action="store_true", @@ -444,22 +463,25 @@ def main(): NCU_MODE = True print("[NCU mode] warmup=1, iters=1") - # (B, T, H, use_gk, use_h0, store_ht, save_vnew) + H = args.heads + HV = args.hv if args.hv is not None else H + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + + # (B, T, H, HV, use_gk, use_h0, store_ht, save_vnew) non_varlen_configs = [ - # Sweep B × H with all features (gk, h0, ht, vnew) - (1, 8192, 64, True, True, True, True), - (2, 8192, 64, True, True, True, True), - (4, 8192, 64, True, True, True, True), - (8, 8192, 64, True, True, True, True), + (1, 8192, H, HV, True, True, True, True), + (2, 8192, H, HV, True, True, True, True), + (4, 8192, H, HV, True, True, True, True), + (8, 8192, H, HV, True, True, True, True), ] - # (num_seqs, total_T, H, ratio, use_gk, use_h0, store_ht, save_vnew) + # (num_seqs, total_T, H, HV, ratio, use_gk, use_h0, store_ht, save_vnew) varlen_configs = [ - (20, 8192, 64, 2.0, True, True, True, True), - (25, 8192, 64, 3.0, True, True, True, True), - (20, 8192, 64, 4.0, True, True, True, True), - (20, 32768, 64, 2.0, True, True, True, True), - (25, 32768, 64, 3.0, True, True, True, True), + (20, 8192, H, HV, 2.0, True, True, True, True), + (25, 8192, H, HV, 3.0, True, True, True, True), + (20, 8192, H, HV, 4.0, True, True, True, True), + (20, 32768, H, HV, 2.0, True, True, True, True), + (25, 32768, H, HV, 3.0, True, True, True, True), ] nv_res, vl_res = [], [] diff --git a/benchmarks/bench_fwd_o.py b/benchmarks/bench_fwd_o.py index 94ecb4af..d9ab9bbe 100644 --- a/benchmarks/bench_fwd_o.py +++ b/benchmarks/bench_fwd_o.py @@ -53,9 +53,10 @@ # ─── FLA baseline imports ─── os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison -from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_rel_mean_abs from fla.ops.gla.chunk import chunk_gla_fwd_o_gk # noqa: E402 +from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_rel_mean_abs # noqa: E402 + # ============================================================ # Constants # ============================================================ @@ -82,17 +83,17 @@ def bench_non_varlen(configs): print("=" * 80) results = [] - for B, T, H in configs: + for B, T, H, HV in configs: scale = K**-0.5 NT = (T + BT - 1) // BT torch.manual_seed(42) torch.cuda.empty_cache() q = torch.randn(B, T, H, K, dtype=dtype, device=device) - v = torch.randn(B, T, H, V, dtype=dtype, device=device) - g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - h = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 - A = torch.randn(B, T, H, BT, dtype=dtype, device=device) * 0.1 + v = torch.randn(B, T, HV, V, dtype=dtype, device=device) + g = torch.randn(B, T, HV, K, dtype=torch.float32, device=device) * 0.1 + h = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + A = torch.randn(B, T, HV, BT, dtype=dtype, device=device) * 0.1 # ---- FLA baseline (accuracy) ---- o_fla = chunk_gla_fwd_o_gk( @@ -107,7 +108,7 @@ def bench_non_varlen(configs): ) # ---- CuTe DSL (accuracy) ---- - o_cute_t = torch.zeros(B, T, H, V, dtype=dtype, device=device) + o_cute_t = torch.zeros(B, T, HV, V, dtype=dtype, device=device) # Warmup / first call triggers compilation via cache chunk_gla_fwd_o( @@ -157,10 +158,13 @@ def run_cute(q=q, v=v, g=g, h=h, o=o_cute_t, A=A, scale=scale): ms_cute = benchmark_cuda_mode_fn(run_cute, default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE) speedup = ms_fla / ms_cute if ms_cute > 0 else float("inf") + hv_str = f"/{HV}" if HV != H else "" r = { "B": B, "T": T, "H": H, + "HV": HV, + "hv_str": hv_str, "relative_rms_error": relative_rms_error, "max_diff": max_diff, "rel_max_diff": rel_max_diff, @@ -171,7 +175,7 @@ def run_cute(q=q, v=v, g=g, h=h, o=o_cute_t, A=A, scale=scale): } results.append(r) print( - f" B={B:2d} T={T:5d} H={H:2d} | " + f" B={B:2d} T={T:5d} H={H:2d}{hv_str:<4s} | " f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} rel_max={rel_max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" @@ -207,7 +211,7 @@ def bench_varlen(configs): print("=" * 80) results = [] - for seq_lens, H in configs: + for seq_lens, H, HV in configs: scale = K**-0.5 T_total = sum(seq_lens) cu_seqlens_list = [0] @@ -219,12 +223,12 @@ def bench_varlen(configs): torch.cuda.empty_cache() # Flat token-indexed tensors (shared data for both kernels) - # 4D with B=1: [1, T_total, H, *] + # q uses H (QK heads), g/v/h/A/o use HV (value heads) q_flat = torch.randn(1, T_total, H, K, dtype=dtype, device=device) - v_flat = torch.randn(1, T_total, H, V, dtype=dtype, device=device) - g_flat = torch.randn(1, T_total, H, K, dtype=torch.float32, device=device) * 0.1 - h_flat = torch.randn(1, total_nt_val, H, K, V, dtype=dtype, device=device) * 0.01 - A_flat = torch.randn(1, T_total, H, BT, dtype=dtype, device=device) * 0.1 + v_flat = torch.randn(1, T_total, HV, V, dtype=dtype, device=device) + g_flat = torch.randn(1, T_total, HV, K, dtype=torch.float32, device=device) * 0.1 + h_flat = torch.randn(1, total_nt_val, HV, K, V, dtype=dtype, device=device) * 0.01 + A_flat = torch.randn(1, T_total, HV, BT, dtype=dtype, device=device) * 0.1 # ---- FLA baseline (needs [1, T_total, H, *] + cu_seqlens int64) ---- cu_fla = torch.tensor(cu_seqlens_list, dtype=torch.long, device=device) @@ -242,7 +246,7 @@ def bench_varlen(configs): ) # ---- CuTe DSL varlen ---- - o_cute_flat = torch.zeros(1, T_total, H, V, dtype=dtype, device=device) + o_cute_flat = torch.zeros(1, T_total, HV, V, dtype=dtype, device=device) cu_cute = torch.tensor(cu_seqlens_list, dtype=torch.int32, device=device) ci_cute = build_chunk_indices(seq_lens, BT=BT, device=device) @@ -316,10 +320,13 @@ def run_cute( min_l, max_l = min(seq_lens), max(seq_lens) avg_l = T_total // n_seqs tag = f"{n_seqs}seqs T={T_total} [{min_l}..{max_l}] avg={avg_l}" + hv_str = f"/{HV}" if HV != H else "" r = { "tag": tag, "T_total": T_total, "H": H, + "HV": HV, + "hv_str": hv_str, "n_seqs": n_seqs, "relative_rms_error": relative_rms_error, "max_diff": max_diff, @@ -331,7 +338,7 @@ def run_cute( } results.append(r) print( - f" {tag:45s} H={H:2d} | " + f" {tag:45s} H={H:2d}{hv_str:<4s} | " f"relative_rms_error={relative_rms_error:.6f} max_diff={max_diff:.6f} rel_max={rel_max_diff:.6f} mean_diff={mean_diff:.8f} | " f"FLA={ms_fla:.4f}ms CuTe={ms_cute:.4f}ms | " f"speedup={speedup:.2f}x" @@ -358,16 +365,17 @@ def print_report(nv_results, vl_results): if nv_results: print("\n [Non-Varlen]") hdr = ( - f" {'B':>3s} {'T':>5s} {'H':>3s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" + f" {'B':>3s} {'T':>5s} {'H':>7s} │ {'rel_rmse':>10s} {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 90}") print(hdr) print(f" {'─' * 90}") for r in nv_results: + h_label = f"{r['H']}{r['hv_str']}" print( - f" {r['B']:3d} {r['T']:5d} {r['H']:3d} │ " - f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f" {r['B']:3d} {r['T']:5d} {h_label:>7s} │ " + f"{r['relative_rms_error']:10.6f} {r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 90}") @@ -375,16 +383,17 @@ def print_report(nv_results, vl_results): if vl_results: print("\n [Varlen]") hdr = ( - f" {'Config':>45s} {'H':>3s} │ {'rel_rmse':>18s} {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" + f" {'Config':>45s} {'H':>7s} │ {'rel_rmse':>10s} {'max_diff':>10s} {'rel_max':>10s} {'mean_diff':>12s}" f" │ {'FLA(ms)':>9s} {'CuTe(ms)':>9s} {'Speedup':>8s}" ) print(f" {'─' * 117}") print(hdr) print(f" {'─' * 117}") for r in vl_results: + h_label = f"{r['H']}{r['hv_str']}" print( - f" {r['tag']:>45s} {r['H']:3d} │ " - f"{r['relative_rms_error']:18.6f} {r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " + f" {r['tag']:>45s} {h_label:>7s} │ " + f"{r['relative_rms_error']:10.6f} {r['max_diff']:10.6f} {r['rel_max_diff']:10.6f} {r['mean_diff']:12.8f} │ " f"{r['ms_fla']:9.4f} {r['ms_cute']:9.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 117}") @@ -409,6 +418,18 @@ def main(): action="store_true", help="NCU profiling mode: warmup=1, iters=1", ) + parser.add_argument( + "--heads", + type=int, + default=64, + help="Number of QK heads H (default: 64)", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of value heads HV (default: same as --heads, i.e. no GVA)", + ) args = parser.parse_args() global NCU_MODE @@ -416,21 +437,24 @@ def main(): NCU_MODE = True print("[NCU mode] warmup=1, iters=1") + H = args.heads + HV = args.hv if args.hv is not None else H + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + non_varlen_configs = [ - # (B, T, H) - (2, 8192, 64), - (2, 32768, 64), - (4, 8192, 64), - (4, 32768, 64), + # (B, T, H, HV) + (2, 8192, H, HV), + (2, 32768, H, HV), + (4, 8192, H, HV), + (4, 32768, H, HV), ] varlen_configs = [ - # (seq_lens, H) — realistic serving scenarios - # ~20-25 seqs, total 8k/32k, lengths vary 2-3x, H=64 - (gen_varlen_seqs(8192, 20, seed=1), 64), - (gen_varlen_seqs(8192, 25, seed=2), 64), - (gen_varlen_seqs(32768, 20, seed=3), 64), - (gen_varlen_seqs(32768, 25, seed=4), 64), + # (seq_lens, H, HV) + (gen_varlen_seqs(8192, 20, seed=1), H, HV), + (gen_varlen_seqs(8192, 25, seed=2), H, HV), + (gen_varlen_seqs(32768, 20, seed=3), H, HV), + (gen_varlen_seqs(32768, 25, seed=4), H, HV), ] nv_res, vl_res = [], [] diff --git a/benchmarks/bench_kda.py b/benchmarks/bench_kda.py index 78bbe241..ea79102f 100644 --- a/benchmarks/bench_kda.py +++ b/benchmarks/bench_kda.py @@ -58,6 +58,7 @@ # Constants # ============================================================ H, D = 64, 128 +HV = H WARMUP = 10 N_ITERS = 30 NCU_MODE = False @@ -102,7 +103,7 @@ def check_determinism(H=4, total_T=8192, num_seqs=10, iters=10000): seq_lens = generate_balanced_seqlens(total_T, num_seqs) cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens) + inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens, num_v_heads=HV) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -150,7 +151,7 @@ def bench_fixed(configs): seq_lens = [T] * B cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens) + inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens, num_v_heads=HV) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -235,7 +236,7 @@ def bench_varlen(configs): T = total_len cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens) + inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, num_v_heads=HV) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -329,14 +330,14 @@ def print_report(fixed_results, varlen_results): print("\n [Fixed-Length]") print(f" {'─' * 85}") print( - f" {'B':>3s} {'T':>5s} │ {'rel_rmse':>18s} {'rel_max':>10s}" + f" {'B':>3s} {'T':>5s} │ {'rel_rmse':>10s} {'rel_max':>10s}" f" │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s}" ) print(f" {'─' * 85}") for r in fixed_results: print( f" {r['B']:3d} {r['T']:5d} │ " - f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} │ " + f"{r['relative_rms_error']:10.6f} {r['rel_max']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 85}") @@ -345,13 +346,13 @@ def print_report(fixed_results, varlen_results): print("\n [Varlen]") print(f" {'─' * 100}") print( - f" {'Config':>45s} │ {'rel_rmse':>18s} {'rel_max':>10s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s}" + f" {'Config':>45s} │ {'rel_rmse':>10s} {'rel_max':>10s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>11s} {'Speedup':>8s}" ) print(f" {'─' * 100}") for r in varlen_results: print( f" {r['tag']:>45s} │ " - f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} │ " + f"{r['relative_rms_error']:10.6f} {r['rel_max']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:11.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 100}") @@ -386,9 +387,23 @@ def main(): action="store_true", help="Disable recompute in both FLA and cuLA (pre-compute QG)", ) + global H + parser.add_argument( + "--heads", + type=int, + default=H, + help=f"Number of Q/K heads (H). Default: {H}", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of V heads (HV). Default: same as H (no GVA). Set HV > H to run in GVA mode.", + ) args = parser.parse_args() global NCU_MODE, SANITIZER_MODE, DISABLE_RECOMPUTE + H = args.heads if args.ncu: NCU_MODE = True print("[NCU mode] warmup=1, iters=1") @@ -398,6 +413,13 @@ def main(): if args.disable_recompute: DISABLE_RECOMPUTE = True print("[Disable recompute] pre-compute QG in forward") + HV = H + if args.hv is not None: + if args.hv < H or args.hv % H != 0: + raise ValueError(f"--hv must be a positive multiple of H ({H}), got {args.hv}") + HV = args.hv + if HV > H: + print(f"[GVA] HV={HV} (H={H}, ratio={HV // H}x)") fixed_configs = [ # (B, T) diff --git a/benchmarks/bench_kda_chunk_intra.py b/benchmarks/bench_kda_chunk_intra.py index c44fe63d..33fec5ce 100644 --- a/benchmarks/bench_kda_chunk_intra.py +++ b/benchmarks/bench_kda_chunk_intra.py @@ -46,7 +46,7 @@ # Constant params B, H, D = 2, 64, 128 -HV = H # overridable via --hv; HV > H enables GVA mode +HV = H # overridable via --hv; HV > H enables GVA mode BT = 64 # chunk size # Varlen benchmark params @@ -57,6 +57,7 @@ DISABLE_RECOMPUTE = False # Whether to disable recompute (compute QG in forward) + # ============================================================================== # Unified uniform seqlen benchmark (handles both standard and GVA) # ============================================================================== @@ -76,7 +77,7 @@ def benchmark_chunk_intra_uniform(): ) print("=" * 100) print( - f"{'B':>4} {'T':>7} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'B':>4} {'T':>7} │ {'rel_rmse':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 100) @@ -89,9 +90,17 @@ def benchmark_chunk_intra_uniform(): ) common = dict( - q=q, k=k, v=v, gk=g, beta=beta, scale=scale, - cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - safe_gate=True, disable_recompute=DISABLE_RECOMPUTE, + q=q, + k=k, + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=True, + disable_recompute=DISABLE_RECOMPUTE, ) # Accuracy: run once and compare @@ -107,7 +116,7 @@ def benchmark_chunk_intra_uniform(): speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{B:>4} {T:>7} │ {relative_rms_error:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{B:>4} {T:>7} │ {relative_rms_error:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 100) @@ -133,7 +142,7 @@ def benchmark_chunk_intra_varlen(): ) print("=" * 110) print( - f"{'total_len':>10} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'total_len':>10} │ {'rel_rmse':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 110) @@ -147,9 +156,17 @@ def benchmark_chunk_intra_varlen(): ) common = dict( - q=q, k=k, v=v, gk=g, beta=beta, scale=scale, - cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - safe_gate=True, disable_recompute=DISABLE_RECOMPUTE, + q=q, + k=k, + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=True, + disable_recompute=DISABLE_RECOMPUTE, ) # Accuracy @@ -165,7 +182,7 @@ def benchmark_chunk_intra_varlen(): speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{total_len:>10} │ {relative_rms_error:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{total_len:>10} │ {relative_rms_error:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 110) @@ -188,7 +205,7 @@ def benchmark_chunk_intra_varlen(): "--hv", type=int, default=None, - help=f"Override number of V heads (HV). Default: H (no GVA). Set HV > H to enable GVA mode.", + help="Override number of V heads (HV). Default: H (no GVA). Set HV > H to enable GVA mode.", ) args = parser.parse_args() diff --git a/benchmarks/bench_kda_decode.py b/benchmarks/bench_kda_decode.py index 5da9afb2..cf68b3a6 100644 --- a/benchmarks/bench_kda_decode.py +++ b/benchmarks/bench_kda_decode.py @@ -166,7 +166,9 @@ def summary(vals): lines.append("### Accuracy (Output)") lines.append("") lines.append("| N | cuLA v out rel_rmse | cuLA v out rel_max | cuLA k out rel_rmse | cuLA k out rel_max |") - lines.append("|---|------------------------------:|-------------------:|------------------------------:|-------------------:|") + lines.append( + "|---|------------------------------:|-------------------:|------------------------------:|-------------------:|" + ) for r in results: lines.append( f"| {r['N']} | {r['out_v_last_relative_rms_error']:.3e} | {r['out_v_last_rel_max']:.3e} | " @@ -176,10 +178,10 @@ def summary(vals): lines.append("### Accuracy (State)") lines.append("") + lines.append("| N | cuLA v state rel_rmse | cuLA v state rel_max | cuLA k state rel_rmse | cuLA k state rel_max |") lines.append( - "| N | cuLA v state rel_rmse | cuLA v state rel_max | cuLA k state rel_rmse | cuLA k state rel_max |" + "|---|--------------------------------:|---------------------:|--------------------------------:|---------------------:|" ) - lines.append("|---|--------------------------------:|---------------------:|--------------------------------:|---------------------:|") for r in results: lines.append( f"| {r['N']} | {r['state_v_last_relative_rms_error']:.3e} | {r['state_v_last_rel_max']:.3e} | " @@ -410,10 +412,7 @@ def print_section(h_dim: int, v_dim: int): ) print() - hdr_out = ( - f"{'N':>5} | {'cuLA v out rel_rmse':>30} | {'rel_max':>10} | " - f"{'cuLA k out rel_rmse':>30} | {'rel_max':>10}" - ) + hdr_out = f"{'N':>5} | {'cuLA v out rel_rmse':>30} | {'rel_max':>10} | {'cuLA k out rel_rmse':>30} | {'rel_max':>10}" print(hdr_out) print("-" * len(hdr_out)) for res in results: @@ -424,8 +423,7 @@ def print_section(h_dim: int, v_dim: int): print() hdr_state = ( - f"{'N':>5} | {'cuLA v state rel_rmse':>32} | {'rel_max':>10} | " - f"{'cuLA k state rel_rmse':>32} | {'rel_max':>10}" + f"{'N':>5} | {'cuLA v state rel_rmse':>32} | {'rel_max':>10} | {'cuLA k state rel_rmse':>32} | {'rel_max':>10}" ) print(hdr_state) print("-" * len(hdr_state)) diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index 6a602767..b42a0f7e 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -342,7 +342,7 @@ def print_report(fixed_results, varlen_results): print(f" {'─' * 110}") print( f" {'B':>3s} {'T':>6s} {'H':>3s} {'HV':>3s} {'GVA':>4s} │ " - f"{'rel_rmse':>18s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'rel_rmse':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" ) print(f" {'─' * 110}") @@ -350,7 +350,7 @@ def print_report(fixed_results, varlen_results): gva_tag = f"{r['HV'] // r['H']}x" if r["HV"] > r["H"] else "no" print( f" {r['B']:3d} {r['T']:6d} {r['H']:3d} {r['HV']:3d} {gva_tag:>4s} │ " - f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " + f"{r['relative_rms_error']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 110}") @@ -360,7 +360,7 @@ def print_report(fixed_results, varlen_results): print(f" {'─' * 120}") print( f" {'Config':>45s} {'H':>3s} {'HV':>3s} {'GVA':>4s} │ " - f"{'rel_rmse':>18s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'rel_rmse':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" ) print(f" {'─' * 120}") @@ -368,7 +368,7 @@ def print_report(fixed_results, varlen_results): gva_tag = f"{r['HV'] // r['H']}x" if r["HV"] > r["H"] else "no" print( f" {r['tag']:>45s} {r['H']:3d} {r['HV']:3d} {gva_tag:>4s} │ " - f"{r['relative_rms_error']:18.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " + f"{r['relative_rms_error']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" ) print(f" {'─' * 120}") diff --git a/benchmarks/bench_kda_fwd_bwd_e2e.py b/benchmarks/bench_kda_fwd_bwd_e2e.py index 2a1f1d65..8a63c5eb 100644 --- a/benchmarks/bench_kda_fwd_bwd_e2e.py +++ b/benchmarks/bench_kda_fwd_bwd_e2e.py @@ -63,6 +63,7 @@ # Constants # ============================================================ H, D = 64, 128 +HV = H # Number of V heads (GVA: HV > H, HV % H == 0) WARMUP = 25 N_ITERS = 100 NCU_MODE = False @@ -164,7 +165,7 @@ def check_determinism(num_seqs=5, T=512, iters=20): seq_lens = generate_random_seq_lens(num_seqs, T, 63, seed=SEED) cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True) + inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True, num_v_heads=HV) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -216,7 +217,7 @@ def bench_fixed(configs): seq_lens = [T] * B cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True) + inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True, num_v_heads=HV) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -342,7 +343,7 @@ def bench_varlen(configs): T = total_len cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True) + inputs = prepare_safe_gate_inputs(1, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=True, num_v_heads=HV) q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] scale, init_state, lower_bound = inputs["scale"], inputs["init_state"], inputs["lower_bound"] @@ -466,7 +467,7 @@ def print_report(fixed_results, varlen_results): print(" BENCHMARK REPORT: chunk_kda forward+backward (E2E)") print(" cuLA CuTe DSL vs FLA Triton") print( - f" H={H} D={D} dtype=bf16 safe_gate=True phase={PHASE} disable_recompute={DISABLE_RECOMPUTE}" + f" H={H} HV={HV} D={D} dtype=bf16 safe_gate=True phase={PHASE} disable_recompute={DISABLE_RECOMPUTE}" ) wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS @@ -502,10 +503,7 @@ def print_report(fixed_results, varlen_results): f"{'rel_max:':>10s}{rel_max_vals}" ) # Line 2: rel_rmse (no timing columns) - print( - f" {'':3s} {'':5s} │ {'':9s} {'':11s} {'':8s} │ " - f"{'rel_rmse:':>10s}{relative_rms_error_vals}" - ) + print(f" {'':3s} {'':5s} │ {'':9s} {'':11s} {'':8s} │ {'rel_rmse:':>10s}{relative_rms_error_vals}") print(f" {'─' * 125}") if varlen_results: @@ -527,10 +525,7 @@ def print_report(fixed_results, varlen_results): f"{'rel_max:':>10s}{rel_max_vals}" ) # Line 2: rel_rmse (no config/timing columns) - print( - f" {'':>45s} │ {'':9s} {'':11s} {'':8s} │ " - f"{'rel_rmse:':>10s}{relative_rms_error_vals}" - ) + print(f" {'':>45s} │ {'':9s} {'':11s} {'':8s} │ {'rel_rmse:':>10s}{relative_rms_error_vals}") print(f" {'─' * 140}") print(f"\n{sep}\n") @@ -570,9 +565,26 @@ def main(): action="store_true", help="Disable recompute in both FLA and cuLA (pre-compute QG)", ) + global H + parser.add_argument( + "--heads", + type=int, + default=H, + help=f"Number of Q/K heads (H). Default: {H}", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of V heads (HV). Default: same as --heads. For GVA, set HV > H with HV %% H == 0", + ) args = parser.parse_args() - global NCU_MODE, SANITIZER_MODE, DISABLE_RECOMPUTE, PHASE + global NCU_MODE, SANITIZER_MODE, DISABLE_RECOMPUTE, PHASE, HV + H = args.heads + HV = args.hv if args.hv is not None else H + if HV < H or HV % H != 0: + raise ValueError(f"--hv ({HV}) must be >= --heads ({H}) and divisible by --heads") if args.ncu: NCU_MODE = True print("[NCU mode] warmup=1, iters=1") diff --git a/benchmarks/bench_la_decode_vs_fla.py b/benchmarks/bench_la_decode_vs_fla.py index 850bca58..27eafbea 100644 --- a/benchmarks/bench_la_decode_vs_fla.py +++ b/benchmarks/bench_la_decode_vs_fla.py @@ -53,9 +53,9 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -from benchmarks.utils import benchmark_cuda_fn, relative_rms_error from fla.ops.common.fused_recurrent import fused_recurrent_fwd, fused_recurrent_fwd_kernel +from benchmarks.utils import benchmark_cuda_fn, relative_rms_error from cula.ops.la_decode import _get_compiled_kernel, linear_attention_decode from cula.utils import USE_FAST_MATH @@ -276,7 +276,7 @@ def main(): print(f"{'=' * 100}") print( f"{'B':>5} | {'fla (ms)':>10} | {'cute (ms)':>10} | " - f"{'speedup':>8} | {'rel_rmse':>18} | {'Rel MaxDiff':>12} | {'State rel_rmse':>24}" + f"{'speedup':>8} | {'rel_rmse':>12} | {'Rel MaxDiff':>12} | {'State rel_rmse':>12}" ) print("─" * 90) @@ -286,8 +286,8 @@ def main(): results.append(r) print( f"{r['B']:>5} | {r['kernel_fla_ms']:>10.4f} | {r['kernel_cute_ms']:>10.4f} | " - f"{r['kernel_speedup']:>7.2f}x | {r['output_relative_rms_error']:>18.6f} | " - f"{r['rel_maxdiff']:>12.6f} | {r['state_relative_rms_error']:>24.8f}" + f"{r['kernel_speedup']:>7.2f}x | {r['output_relative_rms_error']:>12.6f} | " + f"{r['rel_maxdiff']:>12.6f} | {r['state_relative_rms_error']:>12.8f}" ) # ── Wrapper comparison ────────────────────────────────────────────── diff --git a/benchmarks/bench_lightning_attn.py b/benchmarks/bench_lightning_attn.py index 0534fe11..17001446 100644 --- a/benchmarks/bench_lightning_attn.py +++ b/benchmarks/bench_lightning_attn.py @@ -50,9 +50,9 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -from benchmarks.utils import gen_random, gen_skewed, gen_uniform, relative_rms_error, time_cuda_fn from fla.ops.simple_gla.chunk import chunk_simple_gla_fwd +from benchmarks.utils import gen_random, gen_skewed, gen_uniform, relative_rms_error, time_cuda_fn from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen # ============================================================================= @@ -400,9 +400,7 @@ def print_standard_result(r): dsl = f"{r['cutedsl_ms']:.3f}" if _valid(r.get("cutedsl_ms", float("nan"))) else "ERR" sp = f"{r['speedup']:.2f}x" if _valid(r.get("speedup", float("nan"))) else "-" fla_o = ( - f"{r['fla_o_relative_rms_error'] * 100:.3f}%" - if not np.isnan(r.get("fla_o_relative_rms_error", float("nan"))) - else "-" + f"{r['fla_o_relative_rms_error'] * 100:.3f}%" if not np.isnan(r.get("fla_o_relative_rms_error", float("nan"))) else "-" ) cute_o = ( f"{r['cute_o_relative_rms_error'] * 100:.3f}%" @@ -609,14 +607,18 @@ def run_benchmark_suite(args): if not np.isnan(r.get(f"{label}_o_relative_rms_error", float("nan"))) ] if o_rmses: - print(f" {name} O relative_rms_error% (vs naive): avg={np.mean(o_rmses):.4f} max={np.max(o_rmses):.4f}") + print( + f" {name} O relative_rms_error% (vs naive): avg={np.mean(o_rmses):.4f} max={np.max(o_rmses):.4f}" + ) ht_rmses = [ r[f"{label}_ht_relative_rms_error"] * 100 for r in mode_r if not np.isnan(r.get(f"{label}_ht_relative_rms_error", float("nan"))) ] if ht_rmses: - print(f" {name} Ht relative_rms_error% (vs naive): avg={np.mean(ht_rmses):.4f} max={np.max(ht_rmses):.4f}") + print( + f" {name} Ht relative_rms_error% (vs naive): avg={np.mean(ht_rmses):.4f} max={np.max(ht_rmses):.4f}" + ) # --- Plot --- if args.plot: @@ -747,7 +749,9 @@ def generate_report(all_results, modes, args): ) else: f.write("| Config | FLA(ms) | CuteDSL(ms) | Speedup | FLA_O_rel_rmse% | Cute_O_rel_rmse% |\n") - f.write("|--------|---------|-------------|---------|---------------------------|----------------------------|\n") + f.write( + "|--------|---------|-------------|---------|---------------------------|----------------------------|\n" + ) for r in mr: cfg = f"B={r['B']},T={r['T']},H={r['H']}" sp = f"{r['speedup']:.2f}x" if _valid(r.get("speedup", float("nan"))) else "-" diff --git a/benchmarks/bench_recompute_wu.py b/benchmarks/bench_recompute_wu.py index b67d1c07..f68e0592 100644 --- a/benchmarks/bench_recompute_wu.py +++ b/benchmarks/bench_recompute_wu.py @@ -32,7 +32,6 @@ sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison -from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra from fla.ops.kda.wy_fast import recompute_w_u_fwd as fla_recompute_w_u_fwd import cula.cudac as cula_cuda @@ -49,7 +48,7 @@ # Constant params B, H, D = 2, 64, 128 -HV = H # overridable via --hv; HV > H enables GVA mode +HV = H # overridable via --hv; HV > H enables GVA mode BT = 64 # chunk size # Varlen benchmark params @@ -64,6 +63,7 @@ def get_abs_err(ref: torch.Tensor, out: torch.Tensor) -> float: return (ref.float() - out.float()).abs().max().item() + def prepare_recompute_wu_inputs(B, T, device, cu_seqlens=None, chunk_size=BT): """Prepare inputs for recompute_w_u benchmarking (handles both MHA and GVA). @@ -75,9 +75,17 @@ def prepare_recompute_wu_inputs(B, T, device, cu_seqlens=None, chunk_size=BT): ) _, _, _, _, _, Akk = cula_chunk_kda_fwd_intra( - q=q, k=k, v=v, gk=g, beta=beta, scale=scale, - cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - safe_gate=True, disable_recompute=False, + q=q, + k=k, + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=True, + disable_recompute=False, ) return q, k, v, g, beta, Akk, cu_seqlens, chunk_indices @@ -128,7 +136,7 @@ def benchmark_recompute_wu_uniform(): ) print("=" * 100) print( - f"{'B':>4} {'T':>7} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'B':>4} {'T':>7} │ {'rel_rmse':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 100) @@ -150,7 +158,10 @@ def benchmark_recompute_wu_uniform(): stats = {} for name, t_fla, t_cula in [ - ("w", w_fla, w_cula), ("u", u_fla, u_cula), ("qg", qg_fla, qg_cula), ("kg", kg_fla, kg_cula), + ("w", w_fla, w_cula), + ("u", u_fla, u_cula), + ("qg", qg_fla, qg_cula), + ("kg", kg_fla, kg_cula), ]: if t_fla is not None and t_cula is not None: stats[name] = relative_rms_error_rel_max_mean_abs_rhs(t_fla, t_cula) @@ -169,7 +180,7 @@ def benchmark_recompute_wu_uniform(): speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{B:>4} {T:>7} │ {relative_rms_error_value:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{B:>4} {T:>7} │ {relative_rms_error_value:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 100) @@ -193,7 +204,7 @@ def benchmark_recompute_wu_varlen(): ) print("=" * 110) print( - f"{'total_len':>10} │ {'rel_rmse':>18} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" + f"{'total_len':>10} │ {'rel_rmse':>10} {'rel_max':>10} {'mean_diff':>12} │ {'FLA(ms)':>9} {'cuLA(ms)':>9} {'Speedup':>8}" ) print("─" * 110) @@ -216,7 +227,10 @@ def benchmark_recompute_wu_varlen(): stats = {} for name, t_fla, t_cula in [ - ("w", w_fla, w_cula), ("u", u_fla, u_cula), ("qg", qg_fla, qg_cula), ("kg", kg_fla, kg_cula), + ("w", w_fla, w_cula), + ("u", u_fla, u_cula), + ("qg", qg_fla, qg_cula), + ("kg", kg_fla, kg_cula), ]: if t_fla is not None and t_cula is not None: stats[name] = relative_rms_error_rel_max_mean_abs_rhs(t_fla, t_cula) @@ -235,7 +249,7 @@ def benchmark_recompute_wu_varlen(): speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") print( - f"{total_len:>10} │ {relative_rms_error_value:>18.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" + f"{total_len:>10} │ {relative_rms_error_value:>10.6f} {rel_max:>10.6f} {mean_diff:>12.8f} │ {ms_fla:>9.4f} {ms_cula:>9.4f} {speedup:>7.2f}x" ) print("─" * 110) @@ -297,7 +311,7 @@ def check_determinism(num_seqs=NUM_SEQS, T=2001, H=H, iters=1000): "--hv", type=int, default=None, - help=f"Override number of V heads (HV). Default: H (no GVA). Set HV > H to enable GVA mode.", + help="Override number of V heads (HV). Default: H (no GVA). Set HV > H to enable GVA mode.", ) args = parser.parse_args() diff --git a/benchmarks/generate_benchmark_hopper_md.py b/benchmarks/generate_benchmark_hopper_md.py index 13ffd1d8..d05dc0df 100644 --- a/benchmarks/generate_benchmark_hopper_md.py +++ b/benchmarks/generate_benchmark_hopper_md.py @@ -55,13 +55,18 @@ # ============================================================ -def run_kda_fused_fwd_benchmarks(has_init_state: bool = False): +def run_kda_fused_fwd_benchmarks(has_init_state: bool = False, heads=None, hv=None): """Run bench_kda_fused_fwd.main() with programmatic args and return (fixed, varlen) results.""" print("\n>>> Running KDA Fused Forward benchmarks (via bench_kda_fused_fwd.main)...") orig_argv = sys.argv - sys.argv = ["bench_kda_fused_fwd.py", "--mode", "both"] + argv = ["bench_kda_fused_fwd.py", "--mode", "both"] if has_init_state: - sys.argv.append("--init_state") + argv.append("--init_state") + if heads is not None: + argv += ["--heads", str(heads)] + if hv is not None: + argv += ["--hv", str(hv)] + sys.argv = argv try: fixed_res, varlen_res = kda_fused_fwd_main() finally: @@ -151,6 +156,18 @@ def main(): action="store_true", help="Use non-zero initial state (default: False)", ) + parser.add_argument( + "--heads", + type=int, + default=None, + help="Number of Q/K heads (H) for KDA benchmarks. Default: use bench_kda_fused_fwd default.", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of V heads (HV) for KDA benchmarks. For GVA, set HV > H with HV %% H == 0.", + ) args = parser.parse_args() env = get_env_info() @@ -162,7 +179,9 @@ def main(): kda_fused_fixed = data["kda_fused_fixed"] kda_fused_varlen = data["kda_fused_varlen"] else: - kda_fused_fixed, kda_fused_varlen = run_kda_fused_fwd_benchmarks(has_init_state=args.init_state) + kda_fused_fixed, kda_fused_varlen = run_kda_fused_fwd_benchmarks( + has_init_state=args.init_state, heads=args.heads, hv=args.hv + ) if args.save_cache: cache_path = Path(args.save_cache) diff --git a/benchmarks/generate_benchmark_md.py b/benchmarks/generate_benchmark_md.py index 96a09091..2fba9537 100644 --- a/benchmarks/generate_benchmark_md.py +++ b/benchmarks/generate_benchmark_md.py @@ -78,12 +78,17 @@ # ============================================================ -def run_kda_benchmarks(): +def run_kda_benchmarks(heads=None, hv=None): """Run bench_kda.main() with programmatic args and return (fixed, varlen) results.""" print("\n>>> Running KDA benchmarks (via bench_kda.main)...") # bench_kda.main() parses sys.argv — override it temporarily orig_argv = sys.argv - sys.argv = ["bench_kda.py", "--mode", "both"] + argv = ["bench_kda.py", "--mode", "both"] + if heads is not None: + argv += ["--heads", str(heads)] + if hv is not None: + argv += ["--hv", str(hv)] + sys.argv = argv try: fixed_res, varlen_res = kda_main() finally: @@ -289,6 +294,18 @@ def main(): help="Output markdown filename (relative to project root). Default: BENCHMARK_GB200.md", ) parser.add_argument("--save-cache", type=str, default=None, help="Save benchmark results to JSON for future --cache use.") + parser.add_argument( + "--heads", + type=int, + default=None, + help="Number of Q/K heads (H) for KDA benchmarks. Default: use bench_kda default.", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of V heads (HV) for KDA benchmarks. For GVA, set HV > H with HV %% H == 0.", + ) args = parser.parse_args() env = get_env_info() @@ -303,7 +320,7 @@ def main(): la_varlen = data["la_varlen"] la_decode = data.get("la_decode", []) else: - kda_fixed, kda_varlen = run_kda_benchmarks() + kda_fixed, kda_varlen = run_kda_benchmarks(heads=args.heads, hv=args.hv) la_standard, la_varlen = run_lightning_attn_benchmarks() la_decode = run_la_decode_benchmarks() diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 485dc0d7..406093a6 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -482,9 +482,7 @@ def prepare_safe_gate_inputs( ) -def prepare_intra_inputs( - batch_size, T, H, D, device, cu_seqlens=None, chunk_size=CHUNK_SIZE, seed=SEED, num_v_heads=None -): +def prepare_intra_inputs(batch_size, T, H, D, device, cu_seqlens=None, chunk_size=CHUNK_SIZE, seed=SEED, num_v_heads=None): """Prepare preprocessed inputs ready for chunk_kda_fwd_intra. Supports both standard (HV=H) and GVA (HV > H) layouts via ``num_v_heads``: diff --git a/csrc/api/kda_sm100.cu b/csrc/api/kda_sm100.cu index ff89887a..7edca370 100644 --- a/csrc/api/kda_sm100.cu +++ b/csrc/api/kda_sm100.cu @@ -184,8 +184,8 @@ ChunkKDAFwdRecompWU( int tile_num = chunk_indices.size(0); auto device_prop = at::cuda::getCurrentDeviceProperties(); params.num_sm = device_prop->multiProcessorCount; - params.tile_scheduler_params = StaticPersistentTileScheduler::Params{ - tile_num, params.h_v, params.heads_per_group, params.num_sm, nullptr}; + params.tile_scheduler_params = + StaticPersistentTileScheduler::Params{tile_num, params.h_v, params.heads_per_group, params.num_sm, nullptr}; kda::sm100::run_kda_fwd_recomp_w_u_sm100(params, at::cuda::getCurrentCUDAStream()); } \ No newline at end of file diff --git a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp index f928616d..3689bee3 100644 --- a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp @@ -343,19 +343,15 @@ run_kda_fwd_intra_sm100_impl_dispatch(KDA_fwd_intra_params& params, cudaStream_t typename Kernel::SmemLayoutInputFP32{}); // --- Pack TMA params --- - typename Kernel::template TmaParams< - decltype(shape_QK), - decltype(shape_VG), - decltype(tma_Q), - decltype(tma_K), - decltype(tma_G)> - tma_params = { - shape_QK, - shape_VG, - tma_Q, - tma_K, - tma_G, - }; + typename Kernel:: + template TmaParams + tma_params = { + shape_QK, + shape_VG, + tma_Q, + tma_K, + tma_G, + }; // --- Launch config --- auto kernel_fn = &kda_fwd_intra_sm100_kernel_entry; diff --git a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp index 7b624ebb..55cdc686 100644 --- a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp @@ -718,7 +718,7 @@ struct KdaChunkFwdIntraMainloopSm100 { auto blk_coord = TileScheduler::decode_tile_coord( tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); - int head_idx = get<1>(blk_coord); // v-head index + int head_idx = get<1>(blk_coord); // v-head index int tile_idx = get<2>(blk_coord); int qk_head_idx = get<3>(blk_coord); // == head_idx / heads_per_group int token_offset = cu_seqlens_ptr[batch_idx]; diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index d5ff9a9e..d702e46f 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -906,7 +906,7 @@ struct KdaChunkFwdRecompWUMainloopSm100 { auto blk_coord = TileScheduler::decode_tile_coord( tid, params.h_v, params.heads_per_group, chunk_indices_ptr, cu_seqlens_ptr); int batch_idx = get<0>(blk_coord); - int head_idx = get<1>(blk_coord); // v-head + int head_idx = get<1>(blk_coord); // v-head int tile_idx = get<2>(blk_coord); int qk_head_idx = get<3>(blk_coord); // qk-head int token_offset = cu_seqlens_ptr[batch_idx]; @@ -928,8 +928,7 @@ struct KdaChunkFwdRecompWUMainloopSm100 { [[maybe_unused]] auto mQ = [&]() { if constexpr (StoreQG) { return domain_offset( - make_coord(token_offset, _0{}, _0{}), - tma_params.tma_q.get_tma_tensor(tma_params.shape_qk)); + make_coord(token_offset, _0{}, _0{}), tma_params.tma_q.get_tma_tensor(tma_params.shape_qk)); } else { return 0; // unused placeholder } diff --git a/csrc/kda/sm100/tile_scheduler.hpp b/csrc/kda/sm100/tile_scheduler.hpp index 695bb26c..9a2691cf 100644 --- a/csrc/kda/sm100/tile_scheduler.hpp +++ b/csrc/kda/sm100/tile_scheduler.hpp @@ -37,9 +37,9 @@ // =================================================================== struct StaticPersistentTileScheduler { struct Params { - int num_blocks; // number of sequence chunks (from chunk_indices) - int num_heads; // == num_v_heads; tiles are enumerated by v-head - int heads_per_group; // == num_v_heads / num_qk_heads, precomputed on host + int num_blocks; // number of sequence chunks (from chunk_indices) + int num_heads; // == num_v_heads; tiles are enumerated by v-head + int heads_per_group; // == num_v_heads / num_qk_heads, precomputed on host int num_sm; int* tile_counter; // unused diff --git a/cula/kda/blackwell_fused_fwd.py b/cula/kda/blackwell_fused_fwd.py index c0533d24..c7dec95c 100644 --- a/cula/kda/blackwell_fused_fwd.py +++ b/cula/kda/blackwell_fused_fwd.py @@ -297,11 +297,18 @@ def flash_kda_prefill( if not (-5 <= lower_bound < 0): raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." - assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." - assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." + # Validate head dimensions for GVA + B, T, H, K, HV = *q.shape, v.shape[2] + assert q.shape == k.shape, f"q and k must have the same shape, got q={q.shape} vs k={k.shape}" assert q.dtype == k.dtype == v.dtype == torch.bfloat16, "q, k, v must be in bfloat16." + assert beta.dtype == torch.bfloat16 or beta.dtype == torch.float32, "beta must be in bfloat16 or float32." assert q.shape[-1] == k.shape[-1] == v.shape[-1] == 128, "Currently we only support head dim of 128 for KDA" + assert HV % H == 0, ( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by num_qk_heads (H={H}), but got HV % H = {HV % H}" + ) + assert g.shape == (B, T, HV, K), f"g must have shape [B, T, HV, K]={[B, T, HV, K]}, got {list(g.shape)}" + assert beta.shape == (B, T, HV), f"beta must have shape [B, T, HV]={[B, T, HV]}, got {list(beta.shape)}" + if scale is None: scale = k.shape[-1] ** -0.5 o, final_state = ChunkKDAFunction.apply( diff --git a/cula/kda/chunk.py b/cula/kda/chunk.py index b89c7806..cb2df476 100644 --- a/cula/kda/chunk.py +++ b/cula/kda/chunk.py @@ -365,6 +365,7 @@ def chunk_kda( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) + assert cu_seqlens.dtype == torch.int32, "cu_seqlens must be in int32" if initial_state is not None: assert initial_state.dtype == torch.float32, "initial_state must be in float32." @@ -379,13 +380,17 @@ def chunk_kda( if not (-5 <= lower_bound < 0): raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." - assert k.shape[-1] <= 256, "Currently we only support key headdim <=256 for KDA :-(" - assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." - assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." + # Validate head dimensions for GVA + B, T, H, K, HV = *q.shape, v.shape[2] + assert q.shape == k.shape, f"q and k must have the same shape, got q={q.shape} vs k={k.shape}" assert q.dtype == k.dtype == v.dtype == torch.bfloat16, "q, k, v must be in bfloat16." assert beta.dtype == torch.bfloat16 or beta.dtype == torch.float32, "beta must be in bfloat16 or float32." assert q.shape[-1] == k.shape[-1] == v.shape[-1] == 128, "Currently we only support head dim of 128 for KDA" + assert HV % H == 0, ( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by num_qk_heads (H={H}), but got HV % H = {HV % H}" + ) + assert g.shape == (B, T, HV, K), f"g must have shape [B, T, HV, K]={[B, T, HV, K]}, got {list(g.shape)}" + assert beta.shape == (B, T, HV), f"beta must have shape [B, T, HV]={[B, T, HV]}, got {list(beta.shape)}" if scale is None: scale = k.shape[-1] ** -0.5 diff --git a/cula/kda/chunk_bwd.py b/cula/kda/chunk_bwd.py index 4dbd3317..738b1589 100644 --- a/cula/kda/chunk_bwd.py +++ b/cula/kda/chunk_bwd.py @@ -56,7 +56,7 @@ configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in NUM_WARPS for num_stages in [2, 3, 4] ], - key=["H", "K", "V", "BT", "BK", "BV"], + key=["H", "HV", "K", "V", "BT", "BK", "BV"], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=["T"]) @@ -73,6 +73,7 @@ def chunk_kda_bwd_kernel_dAv( scale, T, H: tl.constexpr, + HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, @@ -81,7 +82,8 @@ def chunk_kda_bwd_kernel_dAv( IS_VARLEN: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) @@ -92,12 +94,12 @@ def chunk_kda_bwd_kernel_dAv( # offset calculation q += (bos * H + i_h) * K k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - do += (bos * H + i_h) * V - dv += (bos * H + i_h) * V - dA += (bos * H + i_h) * BT + v += (bos * HV + i_hv) * V + do += (bos * HV + i_hv) * V + dv += (bos * HV + i_hv) * V + dA += (bos * HV + i_hv) * BT - p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_A = tl.make_block_ptr(A + (bos * HV + i_hv) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) b_A = tl.load(p_A, boundary_check=(0, 1)) o_t = i_t * BT + tl.arange(0, BT) @@ -107,9 +109,9 @@ def chunk_kda_bwd_kernel_dAv( b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (V, T), (1, H * V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v, (V, T), (1, HV * V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) # [BV, BT] b_v = tl.load(p_v, boundary_check=(0, 1)) # [BT, BV] @@ -120,7 +122,7 @@ def chunk_kda_bwd_kernel_dAv( b_dv = tl.dot(b_A.to(b_do.dtype), b_do) tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dA = tl.make_block_ptr(dA, (T, BT), (HV * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.0) tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) @@ -137,8 +139,9 @@ def chunk_kda_bwd_kernel_dAv( for BV in BV_LIST for num_warps in NUM_WARPS for num_stages in [2, 3, 4] + if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) ], - key=["BT", "TRANSPOSE_STATE"], + key=["BT", "HV", "TRANSPOSE_STATE"], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=["T"]) @@ -165,6 +168,7 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( scale, T, H: tl.constexpr, + HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, @@ -174,7 +178,8 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( IS_VARLEN: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) if IS_VARLEN: i_tg = i_t.to(tl.int64) @@ -193,26 +198,26 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( q += (bos * H + i_h) * K k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - v_new += (bos * H + i_h) * V - g += (bos * H + i_h) * K - beta += bos * H + i_h - A += (bos * H + i_h) * BT - h += (i_tg * H + i_h) * K * V - do += (bos * H + i_h) * V - dh += (i_tg * H + i_h) * K * V - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - dv2 += (bos * H + i_h) * V - dg += (bos * H + i_h) * K - db += bos * H + i_h - dA += (bos * H + i_h) * BT - - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + v += (bos * HV + i_hv) * V + v_new += (bos * HV + i_hv) * V + g += (bos * HV + i_hv) * K + beta += bos * HV + i_hv + A += (bos * HV + i_hv) * BT + h += (i_tg * HV + i_hv) * K * V + do += (bos * HV + i_hv) * V + dh += (i_tg * HV + i_hv) * K * V + dq += (bos * HV + i_hv) * K + dk += (bos * HV + i_hv) * K + dv += (bos * HV + i_hv) * V + dv2 += (bos * HV + i_hv) * V + dg += (bos * HV + i_hv) * K + db += bos * HV + i_hv + dA += (bos * HV + i_hv) * BT + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) b_beta = tl.load(p_beta, boundary_check=(0,)) - p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_A = tl.make_block_ptr(A, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) b_A = tl.load(p_A, boundary_check=(0, 1)) b_dA = tl.zeros([BT, BT], dtype=tl.float32) @@ -223,11 +228,11 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( m_k = o_k < K p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H * K + o_k + p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * HV * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) b_dq = tl.zeros([BT, BK], dtype=tl.float32) @@ -236,15 +241,15 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dgk = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v_new = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) if TRANSPOSE_STATE: p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) else: p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) # [BT, BV] b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) b_do = tl.load(p_do, boundary_check=(0, 1)) @@ -260,8 +265,8 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) tl.debug_barrier() # DO NOT REMOVE THIS LINE! if i_k == 0: - p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) b_v = tl.load(p_v, boundary_check=(0, 1)) @@ -294,9 +299,9 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_beta[:, None] b_dk = b_dk + b_dkgb * b_gb - p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) @@ -307,8 +312,8 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) b_dA = tl.where(m_A, -b_dA, 0) - p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dA = tl.make_block_ptr(dA, (T, BT), (HV * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (HV,), (i_t * BT,), (BT,), (0,)) tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) @@ -324,7 +329,7 @@ def chunk_kda_bwd_dAv( chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, do.shape[-1] + B, T, H, K, HV, V = *k.shape, do.shape[2], do.shape[-1] BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) @@ -339,9 +344,9 @@ def chunk_kda_bwd_dAv( BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - dA = v.new_empty(B, T, H, BT, dtype=torch.float) + dA = v.new_empty(B, T, HV, BT, dtype=torch.float) dv = torch.empty_like(do) - grid = (NT, B * H) + grid = (NT, B * HV) chunk_kda_bwd_kernel_dAv[grid]( q=q, k=k, @@ -355,6 +360,7 @@ def chunk_kda_bwd_dAv( scale=scale, T=T, H=H, + HV=HV, K=K, V=V, BT=BT, @@ -382,21 +388,22 @@ def chunk_kda_bwd_wy_dqkg_fused( chunk_indices: torch.LongTensor | None = None, transpose_state_layout: bool = False, ): - B, T, H, K, V = *k.shape, v.shape[-1] + B, T, H, K, HV, V = *k.shape, v.shape[2], v.shape[-1] BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - dq = torch.empty_like(q, dtype=torch.float) - dk = torch.empty_like(k, dtype=torch.float) + # dq, dk are allocated at HV dimension; caller reduces to H if GVA + dq = g.new_empty(B, T, HV, K, dtype=torch.float) + dk = g.new_empty(B, T, HV, K, dtype=torch.float) dv2 = torch.empty_like(v) dg = torch.empty_like(g, dtype=torch.float) db = torch.empty_like(beta, dtype=torch.float) dA = torch.empty_like(A, dtype=torch.float) - grid = (NT, B * H) + grid = (NT, B * HV) chunk_kda_bwd_kernel_wy_dqkg_fused[grid]( q=q, k=k, @@ -420,6 +427,7 @@ def chunk_kda_bwd_wy_dqkg_fused( scale=scale, T=T, H=H, + HV=HV, K=K, V=V, BT=BT, @@ -456,8 +464,12 @@ def chunk_kda_bwd( **kwargs, ): assert transpose_state_layout is False, "transpose_state_layout=True is not supported for training." + + H, HV = q.shape[2], v.shape[2] + G = HV // H + if disable_recompute is False: - B, T, _, _ = k.shape + B, T, _, K = k.shape if use_gate_in_kernel: g = kda_gate_chunk_cumsum( g=g_org, @@ -475,10 +487,11 @@ def chunk_kda_bwd( cu_seqlens = prepare_uniform_cu_seqlens(B, T, q.device, torch.int32) if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - w = torch.empty_like(k) + # w, u, kg, qg all live in h_v head space. + w = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) u = torch.empty_like(v) - qg = torch.empty_like(q) if q is not None else None - kg = torch.empty_like(k) if g is not None else None + qg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) if q is not None else None + kg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) cula_cuda.recompute_w_u_cuda(k, v, beta, Akk, g, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q, qg) if cp_context is not None: # Restore the full initial_state tensor from the compressed version. @@ -590,6 +603,11 @@ def chunk_kda_bwd( safe_gate=safe_gate, ) + # For GVA, reduce dq and dk from [B, T, HV, K] back to [B, T, H, K] + if HV > H: + dq = dq.view(*dq.shape[:2], H, G, dq.shape[-1]).sum(dim=3) + dk = dk.view(*dk.shape[:2], H, G, dk.shape[-1]).sum(dim=3) + dA, dbias = None, None dg = chunk_local_cumsum( dg, diff --git a/cula/kda/chunk_intra.py b/cula/kda/chunk_intra.py index aeb063fa..afc02a50 100644 --- a/cula/kda/chunk_intra.py +++ b/cula/kda/chunk_intra.py @@ -20,336 +20,11 @@ import triton.language as tl from fla.ops.utils import prepare_chunk_indices from fla.ops.utils.op import exp2, gather -from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs +from fla.utils import IS_GATHER_SUPPORTED, autotune_cache_kwargs import cula.cudac as cula_cuda from cula.utils import prepare_uniform_cu_seqlens -if IS_TF32_SUPPORTED: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr("tf32") -else: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr("ieee") - -################################################################################ -# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass -################################################################################ - - -@triton.heuristics( - { - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - } -) -@triton.autotune( - configs=[triton.Config({"BK": BK}, num_warps=num_warps) for BK in [32, 64] for num_warps in [1, 2, 4]], - key=["H", "K", "BC"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_kda_fwd_kernel_inter_solve_fused( - q, - k, - g, - beta, - Aqk, - Akkd, - Akk, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_SAFE_GATE: tl.constexpr, -): - """ - Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. - Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. - - This kernel: - 1. Computes off-diagonal Aqk blocks -> writes to global - 2. Computes off-diagonal Akk blocks -> keeps in registers - 3. Loads diagonal Akk blocks from Akkd (fp32) - 4. Does forward substitution on diagonals - 5. Computes merged Akk_inv - 6. Writes Akk_inv to Akk - """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT >= T: - return - - i_tc0 = i_t * BT - i_tc1 = i_t * BT + BC - i_tc2 = i_t * BT + 2 * BC - i_tc3 = i_t * BT + 3 * BC - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - Aqk += (bos * H + i_h) * BT - Akk += (bos * H + i_h) * BT - Akkd += (bos * H + i_h) * BC - - o_i = tl.arange(0, BC) - m_tc1 = (i_tc1 + o_i) < T - m_tc2 = (i_tc2 + o_i) < T - m_tc3 = (i_tc3 + o_i) < T - - b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) - - b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) - - b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) - - ################################################################################ - # off-diagonal blocks - ################################################################################ - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_k0 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - p_g0 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) - - if i_tc1 < T: - p_q1 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_k1 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_g1 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn1 = tl.load(g + i_tc1 * H * K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) - # [BC, BC] - b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) - b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) - - if i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn2 = tl.load(g + i_tc2 * H * K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) - b_qg2 = b_q2 * b_gqn2 - b_kg2 = b_k2 * b_gqn2 - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) - b_Aqk20 += tl.dot(b_qg2, b_kgt) - b_Akk20 += tl.dot(b_kg2, b_kgt) - # [BC, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) - # [BC, BC] - b_Aqk21 += tl.dot(b_qg2, b_kgt) - b_Akk21 += tl.dot(b_kg2, b_kgt) - - if i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn3 = tl.load(g + i_tc3 * H * K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) - b_qg3 = b_q3 * b_gqn3 - b_kg3 = b_k3 * b_gqn3 - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) - # [BC, BC] - b_Aqk30 += tl.dot(b_qg3, b_kgt) - b_Akk30 += tl.dot(b_kg3, b_kgt) - # [BK, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) - # [BC, BC] - b_Aqk31 += tl.dot(b_qg3, b_kgt) - b_Akk31 += tl.dot(b_kg3, b_kgt) - # [BK, BC] - b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) - # [BC, BC] - b_Aqk32 += tl.dot(b_qg3, b_kgt) - b_Akk32 += tl.dot(b_kg3, b_kgt) - - ################################################################################ - # save off-diagonal Aqk blocks and prepare Akk - ################################################################################ - if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b1 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,)) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) - b_Akk10 = b_Akk10 * b_b1[:, None] - if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b2 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,)) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) - b_Akk20 = b_Akk20 * b_b2[:, None] - b_Akk21 = b_Akk21 * b_b2[:, None] - if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) - tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b3 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,)) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) - b_Akk30 = b_Akk30 * b_b3[:, None] - b_Akk31 = b_Akk31 * b_b3[:, None] - b_Akk32 = b_Akk32 * b_b3[:, None] - - p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) - - ################################################################################ - # forward substitution on diagonals - ################################################################################ - - if not USE_SAFE_GATE: - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - b_Ai00 = -tl.where(m_A, b_Ai00, 0) - b_Ai11 = -tl.where(m_A, b_Ai11, 0) - b_Ai22 = -tl.where(m_A, b_Ai22, 0) - b_Ai33 = -tl.where(m_A, b_Ai33, 0) - - for i in range(2, min(BC, T - i_tc0)): - b_a00 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a00 = tl.where(o_i < i, b_a00, 0.0) - b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) - b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) - for i in range(BC + 2, min(2 * BC, T - i_tc0)): - b_a11 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a11 = tl.where(o_i < i - BC, b_a11, 0.0) - b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) - b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) - for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)): - b_a22 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0) - b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) - b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22) - for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)): - b_a33 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0) - b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) - b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33) - - b_Ai00 += m_I - b_Ai11 += m_I - b_Ai22 += m_I - b_Ai33 += m_I - - ################################################################################ - # compute merged inverse using off-diagonals - ################################################################################ - - # we used tf32 to maintain matrix inverse's precision whenever possible. - b_Ai10 = -tl.dot( - tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai21 = -tl.dot( - tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai32 = -tl.dot( - tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai22, input_precision=SOLVE_TRIL_DOT_PRECISION - ) - - b_Ai20 = -tl.dot( - b_Ai22, - tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai31 = -tl.dot( - b_Ai33, - tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai30 = -tl.dot( - b_Ai33, - tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - - ################################################################################ - # store full Akk_inv to Akk - ################################################################################ - - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)) - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)) - - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - @triton.heuristics( { @@ -360,7 +35,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=["BK", "NC", "BT"], + key=["BK", "NC", "BT", "HV"], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=["B", "T"]) @@ -383,6 +58,7 @@ def chunk_kda_bwd_kernel_intra( B, T, H: tl.constexpr, + HV: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, @@ -393,7 +69,8 @@ def chunk_kda_bwd_kernel_intra( USE_GATHER: tl.constexpr, ): i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) i_k, i_i = i_kc // NC, i_kc % NC all = B * T @@ -413,36 +90,36 @@ def chunk_kda_bwd_kernel_intra( q += (bos * H + i_h) * K k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - beta += bos * H + i_h - - dAqk += (bos * H + i_h) * BT - dAkk += (bos * H + i_h) * BT - dq += (bos * H + i_h) * K - dq2 += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dk2 += (bos * H + i_h) * K - dg += (bos * H + i_h) * K - dg2 += (bos * H + i_h) * K - db += (i_k * all + bos) * H + i_h - - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + g += (bos * HV + i_hv) * K + beta += bos * HV + i_hv + + dAqk += (bos * HV + i_hv) * BT + dAkk += (bos * HV + i_hv) * BT + dq += (bos * HV + i_hv) * K + dq2 += (bos * HV + i_hv) * K + dk += (bos * HV + i_hv) * K + dk2 += (bos * HV + i_hv) * K + dg += (bos * HV + i_hv) * K + dg2 += (bos * HV + i_hv) * K + db += (i_k * all + bos) * HV + i_hv + + p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) if i_i > 0: - p_gn = g + i_ti * H * K + o_k + p_gn = g + i_ti * HV * K + o_k # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(0, i_i): p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (HV * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (HV * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) # [BC, BK] b_k = tl.load(p_k, boundary_check=(0, 1)) b_gk = tl.load(p_gk, boundary_check=(0, 1)) @@ -459,9 +136,9 @@ def chunk_kda_bwd_kernel_intra( o_i = tl.arange(0, BC) m_dA = (i_ti + o_i) < T - o_dA = (i_ti + o_i) * H * BT + i_i * BC + o_dA = (i_ti + o_i) * HV * BT + i_i * BC p_kj = k + i_ti * H * K + o_k - p_gkj = g + i_ti * H * K + o_k + p_gkj = g + i_ti * HV * K + o_k p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) @@ -472,11 +149,11 @@ def chunk_kda_bwd_kernel_intra( if USE_GATHER: b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) @@ -508,14 +185,14 @@ def chunk_kda_bwd_kernel_intra( b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * b_gqk, 0.0) p_kj += H * K - p_gkj += H * K + p_gkj += HV * K b_db = tl.sum(b_dk2 * b_k, 1) b_dk2 *= b_b[:, None] - p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dq2 = tl.make_block_ptr(dq2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T,), (H,), (i_ti,), (BC,), (0,)) + p_dq = tl.make_block_ptr(dq, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dq2 = tl.make_block_ptr(dq2, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (HV,), (i_ti,), (BC,), (0,)) b_dg2 = b_q * b_dq2 b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) @@ -527,16 +204,16 @@ def chunk_kda_bwd_kernel_intra( NC = min(NC, tl.cdiv(T - i_t * BT, BC)) if i_i < NC - 1: - p_gn = g + (min(i_ti + BC, T) - 1) * H * K + o_k + p_gn = g + (min(i_ti + BC, T) - 1) * HV * K + o_k # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(i_i + 1, NC): p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT + i_j * BC,), (BC,), (0,)) + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, HV * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, HV * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) # [BC] b_b = tl.load(p_b, boundary_check=(0,)) # [BC, BK] @@ -558,25 +235,25 @@ def chunk_kda_bwd_kernel_intra( b_dkt += tl.dot(b_dAqk, b_qg) b_dkt += tl.dot(b_dAkk, b_kbg) b_dkt *= exp2(b_gn - b_g) - o_dA = i_ti * H * BT + i_i * BC + o_i + o_dA = i_ti * HV * BT + i_i * BC + o_i p_qj = q + i_ti * H * K + o_k p_kj = k + i_ti * H * K + o_k - p_gkj = g + i_ti * H * K + o_k - p_bj = beta + i_ti * H + p_gkj = g + i_ti * HV * K + o_k + p_bj = beta + i_ti * HV if SAFE_GATE: if USE_GATHER: b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, HV * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, HV * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) @@ -598,8 +275,8 @@ def chunk_kda_bwd_kernel_intra( else: for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC,] - b_dAqk = tl.load(dAqk + o_dA + j * H * BT) - b_dAkk = tl.load(dAkk + o_dA + j * H * BT) + b_dAqk = tl.load(dAqk + o_dA + j * HV * BT) + b_dAkk = tl.load(dAkk + o_dA + j * HV * BT) # [BK,] b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) @@ -612,12 +289,12 @@ def chunk_kda_bwd_kernel_intra( p_qj += H * K p_kj += H * K - p_gkj += H * K - p_bj += H - p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_gkj += HV * K + p_bj += HV + p_dk = tl.make_block_ptr(dk, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dk2 = tl.make_block_ptr(dk2, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dg2 = tl.make_block_ptr(dg2, (T, K), (HV * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) @@ -627,122 +304,6 @@ def chunk_kda_bwd_kernel_intra( tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) -@triton.heuristics( - { - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - } -) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] - ], - key=["BT", "BC"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_kda_fwd_kernel_intra_sub_chunk( - q, - k, - g, - beta, - Aqk, - Akk, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_GATHER: tl.constexpr, -): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - i_ti = i_t * BT + i_i * BC - if i_ti >= T: - return - - o_c = i_ti + tl.arange(0, BC) - m_c = o_c < T - - q = q + (bos * H + i_h) * K - k = k + (bos * H + i_h) * K - g = g + (bos * H + i_h) * K - beta = beta + bos * H + i_h - Aqk = Aqk + (bos * H + i_h) * BT - Akk = Akk + (bos * H + i_h) * BC - - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - # caculate offset - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + tl.arange(0, BK) - b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) - b_gn = b_gn[None, :] - - # current block, keep numerical stability by subtracting the left boundary - # less than 85 to avoid overflow in exp2 - b_gm = (b_g - b_gn).to(tl.float32) - - b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0) - b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0) - - b_kgt = tl.trans(b_k * b_gk) - - b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale - b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] - - o_i = tl.arange(0, BC) - m_Aqk = o_i[:, None] >= o_i[None, :] - m_Akk = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) - b_Akk = tl.where(m_Akk, b_Akk, 0.0) - - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (H * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - - ################################################################################ - # forward substitution - ################################################################################ - - b_Ai = -b_Akk - for i in range(2, min(BC, T - i_ti)): - b_a = -tl.load(Akk + (i_ti + i) * H * BC + o_i) - b_a = tl.where(o_i < i, b_a, 0.0) - b_a += tl.sum(b_a[:, None] * b_Ai, 0) - b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) - b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - def chunk_kda_fwd_intra( q: torch.Tensor, k: torch.Tensor, @@ -759,12 +320,10 @@ def chunk_kda_fwd_intra( unified_gref: bool = False, # Set True for ~5% extra perf (slightly lower precision) ): assert safe_gate, "Only safe_gate=True is supported in chunk_kda_fwd_intra for now" - B, T, H_QK, K = k.shape + B, T, H, K = k.shape # GVA: g/beta/v live in h_v head space; q/k live in h_qk head space. - H_V = v.size(2) - assert H_QK > 0 and H_V > 0 and H_V % H_QK == 0, ( - f"HV ({H_V}) must be a positive multiple of HQK ({H_QK})" - ) + HV = v.size(2) + assert H > 0 and HV > 0 and HV % H == 0, f"HV ({HV}) must be a positive multiple of HQK ({H})" BT = chunk_size if cu_seqlens is None: @@ -779,8 +338,8 @@ def chunk_kda_fwd_intra( ) # Aqk and Akk are produced per v-head by the intra kernel. - Aqk = torch.empty(B, T, H_V, BT, device=k.device, dtype=k.dtype) - Akk = torch.empty(B, T, H_V, BT, device=k.device, dtype=k.dtype) + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + Akk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) tile_counter = torch.zeros(1, dtype=torch.int32, device=q.device) cula_cuda.chunk_kda_fwd_intra_cuda( @@ -790,8 +349,8 @@ def chunk_kda_fwd_intra( # w, u, kg, qg all live in h_v head space. w = torch.empty_like(v) u = torch.empty_like(v) - qg = torch.empty(B, T, H_V, K, device=q.device, dtype=q.dtype) if disable_recompute else None - kg = torch.empty(B, T, H_V, K, device=k.device, dtype=k.dtype) if gk is not None else None + qg = torch.empty(B, T, HV, K, device=q.device, dtype=q.dtype) if disable_recompute else None + kg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) if gk is not None else None cula_cuda.recompute_w_u_cuda( k, v, beta, Akk, gk, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q if disable_recompute else None, qg @@ -816,7 +375,7 @@ def chunk_kda_bwd_intra( chunk_size: int = 64, safe_gate: bool = False, ): - B, T, H, K = k.shape + B, T, H, K, HV = *k.shape, g.shape[2] BT = chunk_size BC = min(16, BT) BK = min(32, triton.next_power_of_2(K)) @@ -827,11 +386,11 @@ def chunk_kda_bwd_intra( NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) - dq2 = torch.empty_like(q) - dk2 = torch.empty_like(k) + dq2 = torch.empty_like(dq) + dk2 = torch.empty_like(dk) db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) dg2 = torch.empty_like(dg, dtype=torch.float) - grid = (NK * NC, NT, B * H) + grid = (NK * NC, NT, B * HV) chunk_kda_bwd_kernel_intra[grid]( q=q, k=k, @@ -851,6 +410,7 @@ def chunk_kda_bwd_intra( B=B, T=T, H=H, + HV=HV, K=K, BT=BT, BC=BC, @@ -864,4 +424,4 @@ def chunk_kda_bwd_intra( db = db2.sum(0).add_(db) dg = dg2 - return dq, dk, db, dg \ No newline at end of file + return dq, dk, db, dg diff --git a/cula/ops/chunk_delta_h_sm100.py b/cula/ops/chunk_delta_h_sm100.py index b2a61de3..a34f7d84 100644 --- a/cula/ops/chunk_delta_h_sm100.py +++ b/cula/ops/chunk_delta_h_sm100.py @@ -192,7 +192,7 @@ def _plan_tmem_offsets(tiled_mma_wh, tile_wh, tiled_mma_kv, tile_kv, state_tmem_ ) return wh_off, state_off, vnew_off, kv_off, total - def _compute_grid(self, B, H, V): + def _compute_grid(self, B, HV, V): num_v_tiles = (V + self.BV - 1) // self.BV if self.is_varlen: if self.persistent: @@ -202,26 +202,26 @@ def _compute_grid(self, B, H, V): return (sm_count, 1, 1) else: # Non-persistent: one CTA per work unit, free HW scheduling - total_work_units = num_v_tiles * H * B + total_work_units = num_v_tiles * HV * B return (total_work_units, 1, 1) - return (num_v_tiles, H, B) + return (num_v_tiles, HV, B) @cute.jit def __call__( self, k_in: cute.Tensor, # [B, T, H, K] or [T_total, H, K] - w_in: cute.Tensor, # [B, T, H, K] or [T_total, H, K] - u_in: cute.Tensor, # [B, T, H, V] or [T_total, H, V] - g_in: cute.Tensor, # [B, T, H] or [T_total, H] (fp32, unused currently) - gk_in: cute.Tensor, # [B, T, H, K] or [T_total, H, K] (fp32) - h_out_in: cute.Tensor, # [B, NT, H, K, V] or [NT_total, H, K, V] - v_new_in: cute.Tensor, # [B, T, H, V] or [T_total, H, V] - h0_in: cute.Tensor, # [B, H, K, V] (fp32) - ht_in: cute.Tensor, # [B, H, K, V] + w_in: cute.Tensor, # [B, T, HV, K] or [T_total, HV, K] + u_in: cute.Tensor, # [B, T, HV, V] or [T_total, HV, V] + g_in: cute.Tensor, # [B, T, HV] or [T_total, HV] (fp32, unused currently) + gk_in: cute.Tensor, # [B, T, HV, K] or [T_total, HV, K] (fp32) + h_out_in: cute.Tensor, # [B, NT, HV, K, V] or [NT_total, HV, K, V] + v_new_in: cute.Tensor, # [B, T, HV, V] or [T_total, HV, V] + h0_in: cute.Tensor, # [B, HV, K, V] (fp32) + ht_in: cute.Tensor, # [B, HV, K, V] cu_seqlens_in: cute.Tensor, # [N+1] int32 chunk_offsets_in: cute.Tensor, # [N+1] int32 workspace_in: cute.Tensor, # workspace buffer - problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + problem_size: tuple[Int32, Int32, Int32, Int32, Int32, Int32], total_nt: Int32, use_g: Int32, use_gk: Int32, @@ -243,7 +243,7 @@ def __call__( chunk_offsets_ptr = chunk_offsets_in.iterator workspace_ptr = workspace_in.iterator - B, T, H, K, V = problem_size + B, T, H, HV, K, V = problem_size # For varlen: B=num_seqs, T=total_tokens, data tensors use data_B=1. # For non-varlen: data_B=B, NT=ceil(T/BT). @@ -259,34 +259,34 @@ def __call__( kt_layout = cute.make_layout((K, T, (H, data_B)), stride=(1, H * K, (K, T * H * K))) kt = cute.make_tensor(k_ptr, kt_layout) - w_layout = cute.make_layout((T, K, (H, data_B)), stride=(H * K, 1, (K, T * H * K))) + w_layout = cute.make_layout((T, K, (HV, data_B)), stride=(HV * K, 1, (K, T * HV * K))) w = cute.make_tensor(w_ptr, w_layout) - u_layout = cute.make_layout((T, V, (H, data_B)), stride=(H * V, 1, (V, T * H * V))) + u_layout = cute.make_layout((T, V, (HV, data_B)), stride=(HV * V, 1, (V, T * HV * V))) u = cute.make_tensor(u_ptr, u_layout) v_new = cute.make_tensor(v_new_ptr, u_layout) # h_out: for varlen, NT=total_chunks and data_B=1; for non-varlen, NT=per-seq chunks and data_B=B h_out_T_layout = cute.make_layout( - (V, K, (NT, H, data_B)), - stride=(1, V, (H * K * V, K * V, NT * H * K * V)), + (V, K, (NT, HV, data_B)), + stride=(1, V, (HV * K * V, K * V, NT * HV * K * V)), ) h_out_T = cute.make_tensor(h_out_ptr, h_out_T_layout) # h0/ht always use B=num_seqs (same for both varlen and non-varlen) - h0_layout = cute.make_layout((K, V, (H, B)), stride=(V, 1, (K * V, H * K * V))) + h0_layout = cute.make_layout((K, V, (HV, B)), stride=(V, 1, (K * V, HV * K * V))) h0 = cute.make_tensor(h0_ptr, h0_layout) - ht_T_layout = cute.make_layout((V, K, (H, B)), stride=(1, V, (K * V, H * K * V))) + ht_T_layout = cute.make_layout((V, K, (HV, B)), stride=(1, V, (K * V, HV * K * V))) ht_T = cute.make_tensor(ht_ptr, ht_T_layout) # gk K-first view for TMA: (K, T, (H, data_B)) with K contiguous - gk_K_layout = cute.make_layout((K, T, (H, data_B)), stride=(1, H * K, (K, T * H * K))) + gk_K_layout = cute.make_layout((K, T, (HV, data_B)), stride=(1, HV * K, (K, T * HV * K))) gk_K = cute.make_tensor(gk_ptr, gk_K_layout) # Transposed U view: (V, T, (H, data_B)) to match WH acc shape (M=BV, N=BT) - u_T_layout = cute.make_layout((V, T, (H, data_B)), stride=(1, H * V, (V, T * H * V))) + u_T_layout = cute.make_layout((V, T, (HV, data_B)), stride=(1, HV * V, (V, T * HV * V))) u_T = cute.make_tensor(u_ptr, u_T_layout) self.k_dtype = kt.element_type @@ -432,8 +432,8 @@ def __call__( # v_new transposed GMEM view: (V, T, (H, data_B)) for TMA store v_new_T_layout = cute.make_layout( - (V, T, (H, data_B)), - stride=(1, H * V, (V, T * H * V)), + (V, T, (HV, data_B)), + stride=(1, HV * V, (V, T * HV * V)), ) v_new_T = cute.make_tensor(v_new_ptr, v_new_T_layout) @@ -557,7 +557,7 @@ class SharedStorage: sched_consumed_mbar: cute.struct.MemRange[Int64, 2] self.shared_storage = SharedStorage - self.grid = self._compute_grid(B, H, V) + self.grid = self._compute_grid(B, HV, V) self.kernel( wh_tiled_mma, @@ -642,7 +642,7 @@ def kernel( cu_seqlens: cute.Tensor, chunk_offsets: cute.Tensor, workspace_iter: cute.Pointer, - problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + problem_size: tuple[Int32, Int32, Int32, Int32, Int32, Int32], use_gk: Int32, use_initial_state: Int32, store_final_state: Int32, @@ -819,7 +819,7 @@ def kernel( tCtAccKV = cute.make_tensor(tmem_ptr + self.tmem_kv_off, tCtAccKV_fake.layout) # ===================== Block indices ===================== - B, T, H, K, V = problem_size + B, T, H, HV, K, V = problem_size BT = self.BT if cutlass.const_expr(self.is_varlen): @@ -828,7 +828,7 @@ def kernel( block_idx_x = cute.arch.block_idx()[0] grid_dim_x = cute.arch.grid_dim()[0] num_v_tiles = (V + self.BV - 1) // self.BV - total_work_units = num_v_tiles * H * B + total_work_units = num_v_tiles * HV * B if cutlass.const_expr(self.persistent): # Dynamic scheduling: while loop uses work_idx < total_work_units num_iters = Int32(0) # not used, while loop controls iteration @@ -838,6 +838,7 @@ def kernel( work_idx = Int32(0) v_tile_idx = Int32(0) hidx = Int32(0) + i_h = Int32(0) bidx = Int32(0) tok_offset = Int32(0) seq_len = Int32(0) @@ -846,6 +847,7 @@ def kernel( chunk_off = Int32(0) else: (v_tile_idx, hidx, bidx) = cute.arch.block_idx() + i_h = hidx // (HV // H) tok_offset = Int32(0) seq_len = T NT = (T + BT - 1) // BT @@ -907,8 +909,9 @@ def kernel( work_idx = block_idx_x + wu_iter * grid_dim_x v_tile_idx = work_idx % num_v_tiles temp_work = work_idx // num_v_tiles - hidx = temp_work % H - bidx = temp_work // H + hidx = temp_work % HV + bidx = temp_work // HV + i_h = hidx // (HV // H) tok_offset = cu_seqlens[bidx] seq_len = cu_seqlens[bidx + 1] - tok_offset NT = (seq_len + BT - 1) // BT @@ -946,7 +949,7 @@ def kernel( self.kv_mma_tiler, kv_tiled_mma, data_bidx, - hidx, + i_h, ) # U TMA load partition (non-MMA, epilog-style) @@ -1087,7 +1090,7 @@ def kernel( if cutlass.const_expr(self.is_varlen): if cutlass.const_expr(not self.persistent): work_idx = block_idx_x + wu_iter * grid_dim_x - bidx_mma = (work_idx // num_v_tiles) // H + bidx_mma = (work_idx // num_v_tiles) // HV tok_off_mma = cu_seqlens[bidx_mma] NT = (cu_seqlens[bidx_mma + 1] - tok_off_mma + BT - 1) // BT if cutlass.const_expr(PRINT_DEBUG): @@ -1280,8 +1283,9 @@ def kernel( work_idx = block_idx_x + wu_iter * grid_dim_x v_tile_idx = work_idx % num_v_tiles temp_work = work_idx // num_v_tiles - hidx = temp_work % H - bidx = temp_work // H + hidx = temp_work % HV + bidx = temp_work // HV + i_h = hidx // (HV // H) tok_offset = cu_seqlens[bidx] seq_len = cu_seqlens[bidx + 1] - tok_offset NT = (seq_len + BT - 1) // BT @@ -1494,8 +1498,9 @@ def kernel( work_idx = block_idx_x + wu_iter * grid_dim_x v_tile_idx = work_idx % num_v_tiles temp_work = work_idx // num_v_tiles - hidx = temp_work % H - bidx = temp_work // H + hidx = temp_work % HV + bidx = temp_work // HV + i_h = hidx // (HV // H) tok_offset = cu_seqlens[bidx] seq_len = cu_seqlens[bidx + 1] - tok_offset NT = (seq_len + BT - 1) // BT @@ -1582,7 +1587,7 @@ def kernel( # Construct GMEM tile for this chunk vnew_chunk_raw = ( v_new_tensor.iterator - + (tok_offset + chunk_idx * BT) * H * V + + (tok_offset + chunk_idx * BT) * HV * V + hidx * V + v_tile_idx * self.BV ) @@ -1593,7 +1598,7 @@ def kernel( assumed_align=16, ) vnew_stride_t = cute.assume( - H * V, + HV * V, divby=128 // self.io_dtype.width, ) gVnew_chunk = cute.make_tensor( @@ -1792,11 +1797,11 @@ def reference_bf16_roundtrip(k, w, u, g=None, gk=None, h0=None, chunk_size=64): # Compile cache + TVM-FFI API # --------------------------------------------------------------------------- -# Internal cache: maps (is_varlen, persistent, H, K, V, chunk_size) → compiled_fn +# Internal cache: maps (is_varlen, persistent, H, HV, K, V, chunk_size) → compiled_fn _delta_h_kernel_cache: dict = {} -def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fast_math): +def _compile_delta_h_variant(is_varlen, persistent, H, HV, K, V, chunk_size, use_fast_math): """Compile one ChunkDeltaRuleFwdH kernel variant. Returns the compiled TVM-FFI callable. Uses make_fake_compact_tensor and make_fake_stream for compilation with @@ -1825,7 +1830,7 @@ def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fas sym_ns = cute.sym_int() # num_seqs (varlen h0/ht) or B (non-varlen, == sym_a) if is_varlen: - # varlen: data tensors are [T_total, H, ...] (3D) + # varlen: data tensors are [T_total, H/HV, ...] (3D) k_fake = make_fake_compact_tensor( cutlass.BFloat16, (sym_a, H, K), @@ -1834,42 +1839,42 @@ def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fas ) w_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, H, K), + (sym_a, HV, K), stride_order=(2, 1, 0), assumed_align=128, ) u_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, H, V), + (sym_a, HV, V), stride_order=(2, 1, 0), assumed_align=128, ) g_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_a, H), + (sym_a, HV), stride_order=(1, 0), assumed_align=128, ) gk_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_a, H, K), + (sym_a, HV, K), stride_order=(2, 1, 0), assumed_align=128, ) v_new_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, H, V), + (sym_a, HV, V), stride_order=(2, 1, 0), assumed_align=128, ) h_out_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_nt, H, K, V), + (sym_nt, HV, K, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) else: - # non-varlen: data tensors are [B, T, H, ...] (4D) + # non-varlen: data tensors are [B, T, H/HV, ...] (4D) k_fake = make_fake_compact_tensor( cutlass.BFloat16, (sym_a, sym_b, H, K), @@ -1878,52 +1883,52 @@ def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fas ) w_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_b, H, K), + (sym_a, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128, ) u_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_b, H, V), + (sym_a, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) g_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_a, sym_b, H), + (sym_a, sym_b, HV), stride_order=(2, 1, 0), assumed_align=128, ) gk_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_a, sym_b, H, K), + (sym_a, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128, ) v_new_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_b, H, V), + (sym_a, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) h_out_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_nt, H, K, V), + (sym_a, sym_nt, HV, K, V), stride_order=(4, 3, 2, 1, 0), assumed_align=128, ) - # h0/ht use [B, H, K, V] (non-varlen) or [num_seqs, H, K, V] (varlen) + # h0/ht use [B, HV, K, V] (non-varlen) or [num_seqs, HV, K, V] (varlen) # In varlen mode, num_seqs != T_total, so use a separate sym_ns h0_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_ns, H, K, V), + (sym_ns, HV, K, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) ht_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_ns, H, K, V), + (sym_ns, HV, K, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) @@ -1958,7 +1963,7 @@ def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fas cu_fake, co_fake, ws_fake, - (Int32(1), Int32(1), Int32(H), Int32(K), Int32(V)), + (Int32(1), Int32(1), Int32(H), Int32(HV), Int32(K), Int32(V)), Int32(1), # total_nt dummy Int32(0), # use_g Int32(0), # use_gk @@ -1971,7 +1976,7 @@ def _compile_delta_h_variant(is_varlen, persistent, H, K, V, chunk_size, use_fas return compiled_fn -def _get_compiled_delta_h(is_varlen, persistent, H, K, V, chunk_size): +def _get_compiled_delta_h(is_varlen, persistent, H, HV, K, V, chunk_size): """Get a compiled ChunkDeltaRuleFwdH kernel with on-demand (lazy) compilation. Each variant is compiled exactly once and cached. Compilation is deferred @@ -1980,14 +1985,15 @@ def _get_compiled_delta_h(is_varlen, persistent, H, K, V, chunk_size): where a subsequent cute.compile can invalidate previously compiled but not-yet-executed functions. - Cache key: (is_varlen, persistent, H, K, V, chunk_size, USE_FAST_MATH) + Cache key: (is_varlen, persistent, H, HV, K, V, chunk_size, USE_FAST_MATH) """ - key = (is_varlen, persistent, H, K, V, chunk_size, USE_FAST_MATH) + key = (is_varlen, persistent, H, HV, K, V, chunk_size, USE_FAST_MATH) if key not in _delta_h_kernel_cache: _delta_h_kernel_cache[key] = _compile_delta_h_variant( is_varlen, persistent, H, + HV, K, V, chunk_size, @@ -2016,13 +2022,17 @@ def chunk_gated_delta_rule_fwd_h( Interface aligned with FLA's chunk_gated_delta_rule_fwd_h for fair benchmarking. Allocates output tensors internally and returns (h, v_new, final_state). + GVA (Gated Value Attention): k uses H (QK) heads; w, u, g, gk, h, h0, ht + use HV (value) heads. HV is inferred from u.shape[2]. When H == HV this + reduces to standard (non-GVA) behavior. + Args: - k: key tensor [B, T, H, K] bf16 - w: decay weight tensor [B, T, H, K] bf16 - u: value tensor [B, T, H, V] bf16 - g: scalar gate [B, T, H] fp32, or None - gk: key gate [B, T, H, K] fp32, or None - initial_state: h0 [N, H, K, V] fp32, or None + k: key tensor [B, T, H, K] bf16 + w: decay weight tensor [B, T, HV, K] bf16 + u: value tensor [B, T, HV, V] bf16 + g: scalar gate [B, T, HV] fp32, or None + gk: key gate [B, T, HV, K] fp32, or None + initial_state: h0 [N, HV, K, V] fp32, or None output_final_state: whether to return final_state chunk_size: chunk size (default 64) save_new_value: whether to return v_new @@ -2032,15 +2042,18 @@ def chunk_gated_delta_rule_fwd_h( Returns: (h, v_new, final_state) — same as FLA - h: [B, NT, H, K, V] bf16 (or [1, NT_total, H, K, V] for varlen) - v_new: [B, T, H, V] bf16 (or None if save_new_value=False) - final_state: [N, H, K, V] fp32 (or None if output_final_state=False) + h: [B, NT, HV, K, V] bf16 (or [1, NT_total, HV, K, V] for varlen) + v_new: [B, T, HV, V] bf16 (or None if save_new_value=False) + final_state: [N, HV, K, V] fp32 (or None if output_final_state=False) """ B, T, H, K_dim = k.shape + HV = u.shape[2] V_dim = u.shape[3] BT = chunk_size is_varlen = cu_seqlens is not None + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) # N: the actual number of sequences in the batch with either equal or variable lengths @@ -2069,29 +2082,29 @@ def chunk_gated_delta_rule_fwd_h( w_kern = w[0] u_kern = u[0] # Use torch.empty for dummies the kernel won't read (flag-gated) - g_kern = g[0] if g is not None else torch.empty(T, H, device=k.device, dtype=torch.float32) - gk_kern = gk[0] if gk is not None else torch.empty(T, H, K_dim, device=k.device, dtype=torch.float32) + g_kern = g[0] if g is not None else torch.empty(T, HV, device=k.device, dtype=torch.float32) + gk_kern = gk[0] if gk is not None else torch.empty(T, HV, K_dim, device=k.device, dtype=torch.float32) # Allocate outputs (3D for kernel) - h_out_kern = k.new_empty(total_nt, H, K_dim, V_dim) # bf16 + h_out_kern = k.new_empty(total_nt, HV, K_dim, V_dim) # bf16 v_new_kern = torch.empty_like(u_kern) # always allocate; kernel checks save_v_new flag h0_kern = ( initial_state if initial_state is not None - else torch.empty(N, H, K_dim, V_dim, device=k.device, dtype=torch.float32) + else torch.empty(N, HV, K_dim, V_dim, device=k.device, dtype=torch.float32) ) # ht is purely an output (kernel writes all elements when store_final_state=1); # use empty instead of zeros to skip the zero-fill kernel launch. # NOTE: Ensure final output is zeros # vLLM will use padding for CUDA Graph - ht_kern = torch.zeros(N, H, K_dim, V_dim, device=k.device, dtype=torch.float32) + ht_kern = torch.zeros(N, HV, K_dim, V_dim, device=k.device, dtype=torch.float32) # Workspace: first 4 bytes used as atomic counter for dynamic scheduling workspace = torch.zeros(max(N * 128, 4), dtype=torch.uint8, device=k.device) - ps = (Int32(N), Int32(T), Int32(H), Int32(K_dim), Int32(V_dim)) + ps = (Int32(N), Int32(T), Int32(H), Int32(HV), Int32(K_dim), Int32(V_dim)) - compiled_fn = _get_compiled_delta_h(True, persistent, H, K_dim, V_dim, chunk_size) + compiled_fn = _get_compiled_delta_h(True, persistent, H, HV, K_dim, V_dim, chunk_size) compiled_fn( k_kern, w_kern, @@ -2125,29 +2138,29 @@ def chunk_gated_delta_rule_fwd_h( N = B # Allocate outputs - h = k.new_empty(B, NT, H, K_dim, V_dim) # bf16 + h = k.new_empty(B, NT, HV, K_dim, V_dim) # bf16 v_new_out = torch.empty_like(u) # always allocate; kernel checks save_v_new flag # Use torch.empty for dummies the kernel won't read (flag-gated) h0 = ( initial_state if initial_state is not None - else torch.empty(B, H, K_dim, V_dim, device=k.device, dtype=torch.float32) + else torch.empty(B, HV, K_dim, V_dim, device=k.device, dtype=torch.float32) ) # ht must share sym_ns (first dim) with h0, so always use B - ht = k.new_zeros(B, H, K_dim, V_dim, dtype=torch.float32) + ht = k.new_zeros(B, HV, K_dim, V_dim, dtype=torch.float32) # Dummy tensors for unused optional gate inputs (kernel checks flags) - g_kern = g if g is not None else torch.empty(B, T, H, device=k.device, dtype=torch.float32) - gk_kern = gk if gk is not None else torch.empty(B, T, H, K_dim, device=k.device, dtype=torch.float32) + g_kern = g if g is not None else torch.empty(B, T, HV, device=k.device, dtype=torch.float32) + gk_kern = gk if gk is not None else torch.empty(B, T, HV, K_dim, device=k.device, dtype=torch.float32) # Dummy cu_seqlens / chunk_offsets / workspace (kernel requires them) cu_dummy = torch.empty(2, dtype=torch.int32, device=k.device) co_dummy = torch.empty(2, dtype=torch.int32, device=k.device) ws_dummy = torch.empty(128, dtype=torch.uint8, device=k.device) - ps = (Int32(B), Int32(T), Int32(H), Int32(K_dim), Int32(V_dim)) + ps = (Int32(B), Int32(T), Int32(H), Int32(HV), Int32(K_dim), Int32(V_dim)) - compiled_fn = _get_compiled_delta_h(False, persistent, H, K_dim, V_dim, chunk_size) + compiled_fn = _get_compiled_delta_h(False, persistent, H, HV, K_dim, V_dim, chunk_size) compiled_fn( k, w, @@ -2181,21 +2194,25 @@ def main(): parser.add_argument("--batch_size", type=int, default=1) parser.add_argument("--seq_len", type=int, default=256) parser.add_argument("--num_heads", type=int, default=1) + parser.add_argument( + "--num_v_heads", type=int, default=None, help="Number of value heads (default: num_heads, i.e. no GVA)" + ) parser.add_argument("--head_dim_k", type=int, default=128) parser.add_argument("--head_dim_v", type=int, default=128) parser.add_argument("--chunk_size", type=int, default=64) args = parser.parse_args() B, T, H, K, V = args.batch_size, args.seq_len, args.num_heads, args.head_dim_k, args.head_dim_v + HV = args.num_v_heads if args.num_v_heads is not None else H BT = args.chunk_size NT = (T + BT - 1) // BT - print(f"V2 Test: B={B}, T={T}, H={H}, K={K}, V={V}, BT={BT}, NT={NT}") + print(f"V2 Test: B={B}, T={T}, H={H}, HV={HV}, K={K}, V={V}, BT={BT}, NT={NT}") torch.manual_seed(42) k = torch.randn(B, T, H, K, device="cuda", dtype=torch.bfloat16) * 0.1 - w = torch.randn(B, T, H, K, device="cuda", dtype=torch.bfloat16) * 0.1 - u = torch.randn(B, T, H, V, device="cuda", dtype=torch.bfloat16) * 0.1 + w = torch.randn(B, T, HV, K, device="cuda", dtype=torch.bfloat16) * 0.1 + u = torch.randn(B, T, HV, V, device="cuda", dtype=torch.bfloat16) * 0.1 def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, store_ht, do_save_vnew=0): h_out, v_new, ht = chunk_gated_delta_rule_fwd_h( @@ -2212,24 +2229,28 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st torch.cuda.synchronize() # Ensure consistent return shapes for backward compat with manual tests if h_out is None: - h_out = torch.zeros(B, NT, H, K, V, device="cuda", dtype=torch.bfloat16) + h_out = torch.zeros(B, NT, HV, K, V, device="cuda", dtype=torch.bfloat16) if v_new is None: - v_new = torch.zeros(B, T, H, V, device="cuda", dtype=torch.bfloat16) + v_new = torch.zeros(B, T, HV, V, device="cuda", dtype=torch.bfloat16) if ht is None: - ht = torch.zeros(B, H, K, V, device="cuda", dtype=torch.float32) + ht = torch.zeros(B, HV, K, V, device="cuda", dtype=torch.float32) return h_out, v_new, ht all_pass = True + # For GVA (H != HV), expand k to HV heads for reference comparison + G = HV // H + k_ref = k.repeat_interleave(G, dim=2) if G > 1 else k + # ===== Test 1: No gating, no h0 ===== print("\n" + "=" * 60) print("Test 1: No gating, no h0") - g_z = torch.zeros(B, T, H, device="cuda", dtype=torch.float32) - gk_z = torch.zeros(B, T, H, K, device="cuda", dtype=torch.float32) - h0_z = torch.zeros(B, H, K, V, device="cuda", dtype=torch.float32) + g_z = torch.zeros(B, T, HV, device="cuda", dtype=torch.float32) + gk_z = torch.zeros(B, T, HV, K, device="cuda", dtype=torch.float32) + h0_z = torch.zeros(B, HV, K, V, device="cuda", dtype=torch.float32) h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_z, h0_z, 0, 0, 0, 0) - _, h_ref_bf16 = reference_bf16_roundtrip(k, w, u, h0=None, chunk_size=BT) + _, h_ref_bf16 = reference_bf16_roundtrip(k_ref, w, u, h0=None, chunk_size=BT) max_diff = 0.0 for t in range(min(NT - 1, len(h_ref_bf16))): @@ -2243,15 +2264,14 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st # ===== Test 2: With gk + h0 ===== print("\n" + "=" * 60) print("Test 2: With gk + h0") - gk_val = torch.randn(B, T, H, K, device="cuda", dtype=torch.float32) * 0.1 + gk_val = torch.randn(B, T, HV, K, device="cuda", dtype=torch.float32) * 0.1 gk_val = -torch.abs(gk_val) gk_val = gk_val.cumsum(dim=1) - # Pre-scale by RCP_LN2 to match KDA convention (kernel does exp2 directly) gk_val = gk_val * INV_LN2 - h0_val = torch.randn(B, H, K, V, device="cuda", dtype=torch.float32) * 0.01 + h0_val = torch.randn(B, HV, K, V, device="cuda", dtype=torch.float32) * 0.01 h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_val, h0_val, 0, 1, 1, 0) - _, h_ref_bf16 = reference_bf16_roundtrip(k, w, u, gk=gk_val, h0=h0_val, chunk_size=BT) + _, h_ref_bf16 = reference_bf16_roundtrip(k_ref, w, u, gk=gk_val, h0=h0_val, chunk_size=BT) max_diff = 0.0 for t in range(min(NT - 1, len(h_ref_bf16))): @@ -2265,14 +2285,13 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st # ===== Test 3: With gk gating ===== print("\n" + "=" * 60) print("Test 3: With gk gating") - gk_val = torch.randn(B, T, H, K, device="cuda", dtype=torch.float32) * 0.1 + gk_val = torch.randn(B, T, HV, K, device="cuda", dtype=torch.float32) * 0.1 gk_val = -torch.abs(gk_val) gk_val = gk_val.cumsum(dim=1) - # Pre-scale by RCP_LN2 to match KDA convention (kernel does exp2 directly) gk_val = gk_val * INV_LN2 h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_val, h0_z, 0, 1, 0, 0) - _, h_ref_bf16 = reference_bf16_roundtrip(k, w, u, gk=gk_val, h0=None, chunk_size=BT) + _, h_ref_bf16 = reference_bf16_roundtrip(k_ref, w, u, gk=gk_val, h0=None, chunk_size=BT) max_diff = 0.0 for t in range(min(NT - 1, len(h_ref_bf16))): @@ -2286,10 +2305,10 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st # ===== Test 4: With h0 initial state ===== print("\n" + "=" * 60) print("Test 4: With h0 initial state") - h0_val = torch.randn(B, H, K, V, device="cuda", dtype=torch.float32) * 0.01 + h0_val = torch.randn(B, HV, K, V, device="cuda", dtype=torch.float32) * 0.01 h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_z, h0_val, 0, 0, 1, 0) - _, h_ref_bf16 = reference_bf16_roundtrip(k, w, u, h0=h0_val, chunk_size=BT) + _, h_ref_bf16 = reference_bf16_roundtrip(k_ref, w, u, h0=h0_val, chunk_size=BT) # h_out[0] should be h0 (bf16 rounded) h0_bf16 = h0_val.to(torch.bfloat16) @@ -2310,12 +2329,9 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st print("Test 5: store_final_state") h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_z, h0_z, 0, 0, 0, 1) - _, h_ref_bf16 = reference_bf16_roundtrip(k, w, u, h0=None, chunk_size=BT) + _, h_ref_bf16 = reference_bf16_roundtrip(k_ref, w, u, h0=None, chunk_size=BT) - # ht should match the last h_ref (after all chunks) - ht_ref = h_ref_bf16[-1] # last chunk's state - # ht layout: (B, H, K, V) but kernel writes in transposed (V, K) format - # Compare ht[0, 0] with ht_ref + ht_ref = h_ref_bf16[-1] d_ht = (ht[0, 0].float() - ht_ref.float()).abs().max().item() print(f" ht vs ref: {d_ht:.6f}") t5_pass = d_ht < 0.5 @@ -2327,7 +2343,7 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st print("Test 6: gk + h0 + ht (all features)") h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_val, h0_val, 0, 1, 1, 1) - _, h_ref_bf16 = reference_bf16_roundtrip(k, w, u, gk=gk_val, h0=h0_val, chunk_size=BT) + _, h_ref_bf16 = reference_bf16_roundtrip(k_ref, w, u, gk=gk_val, h0=h0_val, chunk_size=BT) max_diff = 0.0 for t in range(min(NT - 1, len(h_ref_bf16))): @@ -2375,7 +2391,7 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st print("Test 8: v_new output (no gating)") h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_z, h0_z, 0, 0, 0, 0, do_save_vnew=1) - vnew_ref, _ = reference_bf16_roundtrip(k, w, u, h0=None, chunk_size=BT) + vnew_ref, _ = reference_bf16_roundtrip(k_ref, w, u, h0=None, chunk_size=BT) d_vnew = (v_new.float() - vnew_ref.float()).abs().max().item() print(f" v_new max diff: {d_vnew:.6f}") @@ -2388,7 +2404,7 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st print("Test 9: v_new output (with gk gating)") h_out, v_new, ht = run_kernel(k, w, u, g_z, gk_val, h0_z, 0, 1, 0, 0, do_save_vnew=1) - vnew_ref, _ = reference_bf16_roundtrip(k, w, u, gk=gk_val, h0=None, chunk_size=BT) + vnew_ref, _ = reference_bf16_roundtrip(k_ref, w, u, gk=gk_val, h0=None, chunk_size=BT) d_vnew = (v_new.float() - vnew_ref.float()).abs().max().item() print(f" v_new max diff: {d_vnew:.6f}") @@ -2418,12 +2434,13 @@ def run_kernel(k_t, w_t, u_t, g_t, gk_t, h0_t, use_g_val, use_gk_val, use_h0, st # ===== Benchmark ===== print("\n" + "=" * 60) - print("Benchmark: B=4, T=4096, H=64, K=128, V=128") - Bb, Tb, Hb = 4, 4096, 64 + hv_tag = f"/{HV}" if HV != H else "" + print(f"Benchmark: B=4, T=4096, H={H}{hv_tag}, K=128, V=128") + Bb, Tb = 4, 4096 torch.manual_seed(999) - kb = torch.randn(Bb, Tb, Hb, K, device="cuda", dtype=torch.bfloat16) * 0.1 - wb = torch.randn(Bb, Tb, Hb, K, device="cuda", dtype=torch.bfloat16) * 0.1 - ub = torch.randn(Bb, Tb, Hb, V, device="cuda", dtype=torch.bfloat16) * 0.1 + kb = torch.randn(Bb, Tb, H, K, device="cuda", dtype=torch.bfloat16) * 0.1 + wb = torch.randn(Bb, Tb, HV, K, device="cuda", dtype=torch.bfloat16) * 0.1 + ub = torch.randn(Bb, Tb, HV, V, device="cuda", dtype=torch.bfloat16) * 0.1 def run_bench(): chunk_gated_delta_rule_fwd_h( diff --git a/cula/ops/fwd_o_sm100.py b/cula/ops/fwd_o_sm100.py index 2c820aad..9d4bcb46 100644 --- a/cula/ops/fwd_o_sm100.py +++ b/cula/ops/fwd_o_sm100.py @@ -198,7 +198,7 @@ def __init__( ) self.buffer_align_bytes = 1024 - def _compute_grid(self, B, T, H, V, total_nt=None): + def _compute_grid(self, B, T, HV, V, total_nt=None): """Compute grid dimensions for kernel launch.""" num_v_tiles = (V + self.BV - 1) // self.BV if self.persistent: @@ -210,10 +210,10 @@ def _compute_grid(self, B, T, H, V, total_nt=None): return (sm_count, 1, 1) elif self.is_varlen: # Non-persistent varlen: one CTA per work unit. - total_work_units = num_v_tiles * total_nt * H + total_work_units = num_v_tiles * total_nt * HV return (total_work_units, 1, 1) NT = (T + self.BT - 1) // self.BT - return (num_v_tiles, NT, B * H) + return (num_v_tiles, NT, B * HV) @staticmethod def _plan_tmem_offsets( @@ -260,14 +260,14 @@ def _plan_tmem_offsets( def __call__( self, q_in: cute.Tensor, # [B, T, H, K] (B=1 for varlen) - v_in: cute.Tensor, # [B, T, H, V] (B=1 for varlen) - g_in: cute.Tensor, # [B, T, H, K] fp32 (B=1 for varlen) - h_in: cute.Tensor, # [B, NT, H, K, V] (B=1 for varlen) - o_in: cute.Tensor, # [B, T, H, V] (B=1 for varlen) - A_in: cute.Tensor, # [B, T, H, BT] (B=1 for varlen) + v_in: cute.Tensor, # [B, T, HV, V] (B=1 for varlen) + g_in: cute.Tensor, # [B, T, HV, K] fp32 (B=1 for varlen) + h_in: cute.Tensor, # [B, NT, HV, K, V] (B=1 for varlen) + o_in: cute.Tensor, # [B, T, HV, V] (B=1 for varlen) + A_in: cute.Tensor, # [B, T, HV, BT] (B=1 for varlen) cu_seqlens_in: cute.Tensor, # [N+1] int32 chunk_indices_in: cute.Tensor, # [NT, 2] int32 - problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + problem_size: tuple[Int32, Int32, Int32, Int32, Int32, Int32], total_nt: Int32, # total chunks across all seqs (varlen) stream, ): @@ -281,7 +281,7 @@ def __call__( cu_seqlens_ptr = cu_seqlens_in.iterator chunk_indices_ptr = chunk_indices_in.iterator - B, T, H, K, V = problem_size + B, T, H, HV, K, V = problem_size BT = self.BT # For varlen: B=num_seqs, T=max_seqlen (or total_tokens), data_B=1 @@ -303,17 +303,17 @@ def __call__( ) q = cute.make_tensor(q_ptr, q_layout) - # g layout: token-indexed (T, K, (H, data_B)) — fp32 (separate from q) + # g layout: token-indexed (T, K, (HV, data_B)) — fp32 g_layout = cute.make_layout( - (T, K, (H, data_B)), - stride=(H * K, 1, (K, T * H * K)), + (T, K, (HV, data_B)), + stride=(HV * K, 1, (K, T * HV * K)), ) g = cute.make_tensor(g_ptr, g_layout) - # o: row-major (T, V, (H, data_B)) — token-indexed for direct GMEM write (varlen) + # o: row-major (T, V, (HV, data_B)) — token-indexed for direct GMEM write (varlen) o_layout = cute.make_layout( - (T, V, (H, data_B)), - stride=(H * V, 1, (V, T * H * V)), + (T, V, (HV, data_B)), + stride=(HV * V, 1, (V, T * HV * V)), ) o = cute.make_tensor(o_ptr, o_layout) @@ -323,8 +323,8 @@ def __call__( # TMA descriptor collapses the degenerate H dim; keeping batch # at coord-2 guarantees it always maps to an existing TMA dim. v_T_layout = cute.make_layout( - (V, T, (data_B, H)), - stride=(1, H * V, (T * H * V, V)), + (V, T, (data_B, HV)), + stride=(1, HV * V, (T * HV * V, V)), ) v_T = cute.make_tensor(v_ptr, v_T_layout) @@ -337,15 +337,15 @@ def __call__( h_nt_total = B * NT # NOTE: Mode 2 uses (batch, H) order — see v_T comment above. h_T_layout = cute.make_layout( - (V, K, (h_nt_total, H)), - stride=(1, V, (H * K * V, K * V)), + (V, K, (h_nt_total, HV)), + stride=(1, V, (HV * K * V, K * V)), ) h_T = cute.make_tensor(h_ptr, h_T_layout) - # A layout: token-indexed (T, BT, (H, data_B)) + # A layout: token-indexed (T, BT, (HV, data_B)) a_layout = cute.make_layout( - (T, BT, (H, data_B)), - stride=(H * BT, 1, (BT, T * H * BT)), + (T, BT, (HV, data_B)), + stride=(HV * BT, 1, (BT, T * HV * BT)), ) A = cute.make_tensor(A_ptr, a_layout) @@ -570,7 +570,7 @@ class SharedStorage: ) # ===================== Grid ===================== - grid = self._compute_grid(B, T, H, V, total_nt=total_nt) + grid = self._compute_grid(B, T, HV, V, total_nt=total_nt) # ===================== cu_seqlens / chunk_indices tensors ===================== cu_seqlens = cute.make_tensor(cu_seqlens_ptr, cute.make_layout((B + 1,))) @@ -683,7 +683,7 @@ def kernel( problem_size, total_nt, ): - B, T, H, K, V = problem_size + B, T, H, HV, K, V = problem_size BT = self.BT # ===================== Work decode ===================== @@ -693,12 +693,13 @@ def kernel( # Persistent kernel: 1D grid, work decoded inside each warp's loop block_idx_x = cute.arch.block_idx()[0] grid_dim_x = cute.arch.grid_dim()[0] - total_work_units = num_v_tiles * total_nt * H + total_work_units = num_v_tiles * total_nt * HV num_iters = (total_work_units - block_idx_x + grid_dim_x - 1) // grid_dim_x # Pre-initialize persistent loop variables (CuTe DSL requirement) i_v = Int32(0) chunk_global_idx = Int32(0) i_h = Int32(0) + i_qh = Int32(0) i_b = Int32(0) i_t = Int32(0) tok_offset = Int32(0) @@ -713,8 +714,9 @@ def kernel( i_v = cute.arch.block_idx()[0] i_t = cute.arch.block_idx()[1] i_bh = cute.arch.block_idx()[2] - i_b = i_bh // H - i_h = i_bh % H + i_b = i_bh // HV + i_h = i_bh % HV + i_qh = i_h // (HV // H) tok_offset = i_b * T seq_len = T data_bidx = i_b @@ -874,6 +876,7 @@ def kernel( temp_work = work_idx // num_v_tiles chunk_flat = temp_work % total_nt i_h = temp_work // total_nt + i_qh = i_h // (HV // H) if cutlass.const_expr(self.is_varlen): i_b = chunk_indices[(chunk_flat, 0)] i_t = chunk_indices[(chunk_flat, 1)] @@ -901,7 +904,7 @@ def kernel( # --- Unconditional TMA partitions --- bSG_sQ, bSG_gQ = self._epilog_partition_varlen( tma_atom_q, - tma_q_v[None, None, (i_h, data_bidx)], + tma_q_v[None, None, (i_qh, data_bidx)], (self.BT, self.BK), sQ_epi, ) @@ -1001,7 +1004,7 @@ def kernel( # Bulk prefetch: SMEM → registers (all 256 bf16 at once) cute.autovec_copy(tOsO, tOrO) - o_chunk_raw = o_tensor.iterator + (tok_offset + i_t * BT) * H * V + i_h * V + i_v * self.BV + o_chunk_raw = o_tensor.iterator + (tok_offset + i_t * BT) * HV * V + i_h * V + i_v * self.BV o_chunk_ptr = cute.make_ptr( self.io_dtype, o_chunk_raw.toint(), @@ -1009,7 +1012,7 @@ def kernel( assumed_align=16, ) o_stride_bt = cute.assume( - H * V, + HV * V, divby=128 // self.io_dtype.width, ) gO_chunk = cute.make_tensor( @@ -1580,7 +1583,7 @@ def reference_chunk_gla_fwd_o(q, v, g, h, A, scale, chunk_size=64): # Compile cache + TVM-FFI API # --------------------------------------------------------------------------- -# Internal cache: maps (is_varlen, persistent, H, K, V, scale, chunk_size) → compiled_fn +# Internal cache: maps (is_varlen, persistent, H, HV, K, V, scale, chunk_size) → compiled_fn _fwd_o_kernel_cache: dict = {} # Pre-allocated dummy tensors for non-varlen path (avoid per-call torch.zeros) @@ -1588,7 +1591,7 @@ def reference_chunk_gla_fwd_o(q, v, g, h, A, scale, chunk_size=64): _fwd_o_dummy_chunk_indices: torch.Tensor = None -def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, use_fast_math): +def _compile_fwd_o_variant(is_varlen, persistent, H, HV, K, V, scale, chunk_size, use_fast_math): """Compile one ChunkGlaFwdO kernel variant. Returns the compiled TVM-FFI callable. Uses make_fake_compact_tensor and make_fake_stream for compilation with @@ -1615,8 +1618,8 @@ def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, us BT = chunk_size if is_varlen: - # varlen: tensors are [1, T_total, H, ...] (4D with B=1) - # This avoids squeeze(0) CPU overhead at the call site. + # varlen: tensors are [1, T_total, H/HV, ...] (4D with B=1) + # q uses H (QK heads), g/v/o/A use HV (value heads) q_fake = make_fake_compact_tensor( cutlass.BFloat16, (1, sym_b, H, K), @@ -1625,30 +1628,31 @@ def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, us ) v_fake = make_fake_compact_tensor( cutlass.BFloat16, - (1, sym_b, H, V), + (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) g_fake = make_fake_compact_tensor( cutlass.Float32, - (1, sym_b, H, K), + (1, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128, ) o_fake = make_fake_compact_tensor( cutlass.BFloat16, - (1, sym_b, H, V), + (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) A_fake = make_fake_compact_tensor( cutlass.BFloat16, - (1, sym_b, H, BT), + (1, sym_b, HV, BT), stride_order=(3, 2, 1, 0), assumed_align=128, ) else: - # non-varlen: tensors are [B, T, H, ...] (4D) + # non-varlen: tensors are [B, T, H/HV, ...] (4D) + # q uses H (QK heads), g/v/o/A use HV (value heads) q_fake = make_fake_compact_tensor( cutlass.BFloat16, (sym_a, sym_b, H, K), @@ -1657,42 +1661,42 @@ def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, us ) v_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_b, H, V), + (sym_a, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) g_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_a, sym_b, H, K), + (sym_a, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128, ) o_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_b, H, V), + (sym_a, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128, ) A_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_b, H, BT), + (sym_a, sym_b, HV, BT), stride_order=(3, 2, 1, 0), assumed_align=128, ) if is_varlen: - # varlen: h is [1, NT_total, H, K, V] (5D with B=1) + # varlen: h is [1, NT_total, HV, K, V] (5D with B=1) h_fake = make_fake_compact_tensor( cutlass.BFloat16, - (1, sym_nt, H, K, V), + (1, sym_nt, HV, K, V), stride_order=(4, 3, 2, 1, 0), assumed_align=128, ) else: - # non-varlen: h is [B, NT, H, K, V] (5D) + # non-varlen: h is [B, NT, HV, K, V] (5D) h_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_a, sym_nt, H, K, V), + (sym_a, sym_nt, HV, K, V), stride_order=(4, 3, 2, 1, 0), assumed_align=128, ) @@ -1720,7 +1724,7 @@ def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, us A_fake, cu_fake, ci_fake, - (Int32(1), Int32(1), Int32(H), Int32(K), Int32(V)), + (Int32(1), Int32(1), Int32(H), Int32(HV), Int32(K), Int32(V)), Int32(1), stream_fake, options=COMPILE_OPTIONS, @@ -1728,7 +1732,7 @@ def _compile_fwd_o_variant(is_varlen, persistent, H, K, V, scale, chunk_size, us return compiled_fn -def _get_compiled_fwd_o(is_varlen, persistent, H, K, V, scale, chunk_size): +def _get_compiled_fwd_o(is_varlen, persistent, H, HV, K, V, scale, chunk_size): """Get a compiled ChunkGlaFwdO kernel with on-demand (lazy) compilation. Each variant is compiled exactly once and cached. Compilation is deferred @@ -1737,14 +1741,15 @@ def _get_compiled_fwd_o(is_varlen, persistent, H, K, V, scale, chunk_size): where a subsequent cute.compile can invalidate previously compiled but not-yet-executed functions. - Cache key: (is_varlen, persistent, H, K, V, scale, chunk_size, USE_FAST_MATH) + Cache key: (is_varlen, persistent, H, HV, K, V, scale, chunk_size, USE_FAST_MATH) """ - key = (is_varlen, persistent, H, K, V, scale, chunk_size, USE_FAST_MATH) + key = (is_varlen, persistent, H, HV, K, V, scale, chunk_size, USE_FAST_MATH) if key not in _fwd_o_kernel_cache: _fwd_o_kernel_cache[key] = _compile_fwd_o_variant( is_varlen, persistent, H, + HV, K, V, scale, @@ -1778,15 +1783,15 @@ def chunk_gla_fwd_o( sym_int() is used for B, T, NT so a single compilation handles all batch-size / sequence-length combinations. - Cache key: (is_varlen, persistent, H, K, V, scale, chunk_size) + Cache key: (is_varlen, persistent, H, HV, K, V, scale, chunk_size) Args: - q: query tensor — [B, T, H, K] bf16 (both non-varlen and varlen with B=1) - v: value tensor — [B, T, H, V] bf16 (both non-varlen and varlen with B=1) - g: gate tensor — [B, T, H, K] fp32 (both non-varlen and varlen with B=1) - h: state tensor — [B, NT, H, K, V] bf16 (B=1 for varlen) - o: output tensor (pre-allocated) — same shape as q but with V dim - A: attention matrix — [B, T, H, BT] bf16 (both non-varlen and varlen with B=1) + q: query tensor — [B, T, H, K] bf16 (H = QK heads) + v: value tensor — [B, T, HV, V] bf16 (HV = value heads, HV >= H) + g: gate tensor — [B, T, HV, K] fp32 + h: state tensor — [B, NT, HV, K, V] bf16 (B=1 for varlen) + o: output tensor (pre-allocated) — [B, T, HV, V] bf16 + A: attention matrix — [B, T, HV, BT] bf16 scale: attention scale factor chunk_size: chunk size (default: 64) cu_seqlens: cumulative sequence lengths [N+1] int32 (varlen only) @@ -1802,20 +1807,22 @@ def chunk_gla_fwd_o( "cu_seqlens and chunk_indices are required for varlen mode" ) assert q.dim() == 4 and q.shape[0] == 1, f"varlen mode expects [1, T_total, H, K] input, got shape {q.shape}" - assert h.dim() == 5 and h.shape[0] == 1, f"varlen mode expects [1, NT_total, H, K, V] for h, got shape {h.shape}" + assert h.dim() == 5 and h.shape[0] == 1, f"varlen mode expects [1, NT_total, HV, K, V] for h, got shape {h.shape}" T_total = q.shape[1] H = q.shape[2] + HV = v.shape[2] K = q.shape[3] V = v.shape[3] num_seqs = cu_seqlens.shape[0] - 1 total_nt_val = chunk_indices.shape[0] - ps = (Int32(num_seqs), Int32(T_total), Int32(H), Int32(K), Int32(V)) + ps = (Int32(num_seqs), Int32(T_total), Int32(H), Int32(HV), Int32(K), Int32(V)) else: B, T, H, K = q.shape + HV = v.shape[2] V = v.shape[3] NT = (T + chunk_size - 1) // chunk_size total_nt_val = B * NT - ps = (Int32(B), Int32(T), Int32(H), Int32(K), Int32(V)) + ps = (Int32(B), Int32(T), Int32(H), Int32(HV), Int32(K), Int32(V)) if cu_seqlens is None: global _fwd_o_dummy_cu_seqlens if _fwd_o_dummy_cu_seqlens is None or _fwd_o_dummy_cu_seqlens.device != q.device: @@ -1831,6 +1838,7 @@ def chunk_gla_fwd_o( is_varlen, persistent, H, + HV, K, V, scale, @@ -1864,6 +1872,7 @@ def main(): parser.add_argument("--B", type=int, default=2) parser.add_argument("--T", type=int, default=256) parser.add_argument("--H", type=int, default=4) + parser.add_argument("--HV", type=int, default=None, help="Number of value heads (default: same as --H)") parser.add_argument("--K", type=int, default=128) parser.add_argument("--V", type=int, default=128) parser.add_argument("--scale", type=float, default=None) @@ -1873,12 +1882,16 @@ def main(): if args.scale is None: args.scale = args.K**-0.5 B, T, H, K, V = args.B, args.T, args.H, args.K, args.V + HV = args.HV if args.HV is not None else H + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + G = HV // H BT = args.chunk_size scale = args.scale NT = (T + BT - 1) // BT dtype, device = torch.bfloat16, "cuda" - print(f"Config: B={B}, T={T}, H={H}, K={K}, V={V}, BT={BT}, scale={scale:.4f}") + hv_str = f"/{HV}" if HV != H else "" + print(f"Config: B={B}, T={T}, H={H}{hv_str}, K={K}, V={V}, BT={BT}, scale={scale:.4f}") print(f" Chunks per seq: {NT}, Total chunks: {B * NT}") if args.test in ("correctness", "both"): @@ -1888,13 +1901,14 @@ def main(): print("\n=== Non-Varlen Correctness Test ===") torch.manual_seed(42) q_nv = torch.randn(B, T, H, K, dtype=dtype, device=device) - v_nv = torch.randn(B, T, H, V, dtype=dtype, device=device) - g_nv = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - h_nv = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 - A_nv = torch.randn(B, T, H, BT, dtype=dtype, device=device) * 0.1 + v_nv = torch.randn(B, T, HV, V, dtype=dtype, device=device) + g_nv = torch.randn(B, T, HV, K, dtype=torch.float32, device=device) * 0.1 + h_nv = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + A_nv = torch.randn(B, T, HV, BT, dtype=dtype, device=device) * 0.1 - o_ref_nv = reference_chunk_gla_fwd_o(q_nv, v_nv, g_nv, h_nv, A_nv, scale, BT) - o_nv = torch.zeros(B, T, H, V, dtype=dtype, device=device) + q_ref = q_nv.repeat_interleave(G, dim=2) + o_ref_nv = reference_chunk_gla_fwd_o(q_ref, v_nv, g_nv, h_nv, A_nv, scale, BT) + o_nv = torch.zeros(B, T, HV, V, dtype=dtype, device=device) chunk_gla_fwd_o( q=q_nv, @@ -1943,13 +1957,14 @@ def main(): ci_t = build_chunk_indices(seq_lens, BT=BT, device=device) q_flat = torch.randn(1, T_total, H, K, dtype=dtype, device=device) - v_flat = torch.randn(1, T_total, H, V, dtype=dtype, device=device) - g_flat = torch.randn(1, T_total, H, K, dtype=torch.float32, device=device) * 0.1 - h_flat = torch.randn(1, total_nt_val, H, K, V, dtype=dtype, device=device) * 0.01 - A_flat = torch.randn(1, T_total, H, BT, dtype=dtype, device=device) * 0.1 - o_flat = torch.zeros(1, T_total, H, V, dtype=dtype, device=device) + v_flat = torch.randn(1, T_total, HV, V, dtype=dtype, device=device) + g_flat = torch.randn(1, T_total, HV, K, dtype=torch.float32, device=device) * 0.1 + h_flat = torch.randn(1, total_nt_val, HV, K, V, dtype=dtype, device=device) * 0.01 + A_flat = torch.randn(1, T_total, HV, BT, dtype=dtype, device=device) * 0.1 + o_flat = torch.zeros(1, T_total, HV, V, dtype=dtype, device=device) # Reference per-sequence + q_ref_flat = q_flat[:, :, :, :].repeat_interleave(G, dim=2) o_ref_flat = torch.zeros_like(o_flat) for seq_idx, sl in enumerate(seq_lens): s = cu_seqlens_list[seq_idx] @@ -1957,7 +1972,13 @@ def main(): co = chunk_offsets_list[seq_idx] nt_seq = (sl + BT - 1) // BT o_seq = reference_chunk_gla_fwd_o( - q_flat[:, s:e], v_flat[:, s:e], g_flat[:, s:e], h_flat[:, co : co + nt_seq], A_flat[:, s:e], scale, BT + q_ref_flat[:, s:e], + v_flat[:, s:e], + g_flat[:, s:e], + h_flat[:, co : co + nt_seq], + A_flat[:, s:e], + scale, + BT, ) o_ref_flat[:, s:e] = o_seq @@ -1995,12 +2016,13 @@ def main(): for i in range(3): torch.manual_seed(i * 100) q_cr = torch.randn(B, T, H, K, dtype=dtype, device=device) - v_cr = torch.randn(B, T, H, V, dtype=dtype, device=device) - g_cr = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - h_cr = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 - A_cr = torch.randn(B, T, H, BT, dtype=dtype, device=device) * 0.1 - o_cr = torch.zeros(B, T, H, V, dtype=dtype, device=device) - o_ref_cr = reference_chunk_gla_fwd_o(q_cr, v_cr, g_cr, h_cr, A_cr, scale, BT) + v_cr = torch.randn(B, T, HV, V, dtype=dtype, device=device) + g_cr = torch.randn(B, T, HV, K, dtype=torch.float32, device=device) * 0.1 + h_cr = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + A_cr = torch.randn(B, T, HV, BT, dtype=dtype, device=device) * 0.1 + o_cr = torch.zeros(B, T, HV, V, dtype=dtype, device=device) + q_ref_cr = q_cr.repeat_interleave(G, dim=2) + o_ref_cr = reference_chunk_gla_fwd_o(q_ref_cr, v_cr, g_cr, h_cr, A_cr, scale, BT) chunk_gla_fwd_o( q=q_cr, @@ -2027,11 +2049,11 @@ def main(): for bench_T in [1024, 2048, 4096]: bench_NT = (bench_T + BT - 1) // BT q_b = torch.randn(B, bench_T, H, K, dtype=dtype, device=device) - v_b = torch.randn(B, bench_T, H, V, dtype=dtype, device=device) - g_b = torch.randn(B, bench_T, H, K, dtype=torch.float32, device=device) * 0.1 - h_b = torch.randn(B, bench_NT, H, K, V, dtype=dtype, device=device) * 0.01 - A_b = torch.randn(B, bench_T, H, BT, dtype=dtype, device=device) * 0.1 - o_b = torch.zeros(B, bench_T, H, V, dtype=dtype, device=device) + v_b = torch.randn(B, bench_T, HV, V, dtype=dtype, device=device) + g_b = torch.randn(B, bench_T, HV, K, dtype=torch.float32, device=device) * 0.1 + h_b = torch.randn(B, bench_NT, HV, K, V, dtype=dtype, device=device) * 0.01 + A_b = torch.randn(B, bench_T, HV, BT, dtype=dtype, device=device) * 0.1 + o_b = torch.zeros(B, bench_T, HV, V, dtype=dtype, device=device) # Warmup (also triggers lazy compilation if needed) for _ in range(3): diff --git a/tests/test_kda.py b/tests/test_kda.py index fadadbfe..d2b17d32 100644 --- a/tests/test_kda.py +++ b/tests/test_kda.py @@ -35,6 +35,7 @@ "B", "T", "H", + "HV", "D", "gate_logit_normalizer", "mask_p", @@ -46,17 +47,21 @@ [ pytest.param( *test, - id="B{}-T{}-H{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), + id="B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), ) for test in [ - (1, 63, 1, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 500, 3, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 1000, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), - (3, 1024, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 128, 1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (4, 2048, 8, 128, 1, 0, False, True, True, torch.bfloat16), + (1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), + (2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), + (2, 1000, 3, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), + (3, 1024, 4, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), + (4, 1024, 4, 4, 128, 1, 0, False, False, True, torch.bfloat16), + (4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), + (2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), + (4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), + # GVA cases: HV > H + (2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), + (2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), + (2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), ] ], ) @@ -64,6 +69,7 @@ def test_safe_gate_chunk( B: int, T: int, H: int, + HV: int, D: int, gate_logit_normalizer: float, mask_p: float, @@ -77,11 +83,11 @@ def test_safe_gate_chunk( torch.manual_seed(42) q = torch.rand(B, T, H, D, dtype=dtype) k = torch.rand(B, T, H, D, dtype=dtype) - v = torch.rand(B, T, H, D, dtype=dtype) - g = torch.randn(B, T, H, D, dtype=torch.float if not use_gate_in_kernel else dtype) + v = torch.rand(B, T, HV, D, dtype=dtype) + g = torch.randn(B, T, HV, D, dtype=torch.float if not use_gate_in_kernel else dtype) if use_gate_in_kernel: - A_log = torch.randn(H, dtype=torch.float) - dt_bias = torch.randn(H * D, dtype=torch.float) + A_log = torch.randn(HV, dtype=torch.float) + dt_bias = torch.randn(HV * D, dtype=torch.float) else: g = F.logsigmoid(g) / gate_logit_normalizer g = g * (torch.rand_like(g) > mask_p) @@ -94,8 +100,8 @@ def test_safe_gate_chunk( lower_bound = None naive_kda_gate_fn = naive_kda_gate - beta = torch.randn(B, T, H, dtype=torch.float32).sigmoid().to(beta_dtype) - h0 = torch.randn(B, H, D, D, dtype=torch.float32) + beta = torch.randn(B, T, HV, dtype=torch.float32).sigmoid().to(beta_dtype) + h0 = torch.randn(B, HV, D, D, dtype=torch.float32) if use_gate_in_kernel: A_log, dt_bias = map(lambda x: x.to(device).requires_grad_(True), (A_log, dt_bias)) q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) @@ -158,17 +164,18 @@ def test_safe_gate_chunk( @pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) @pytest.mark.parametrize("disable_recompute", [True, False], ids=["no_recomp", "recomp"]) @pytest.mark.parametrize( - ("H", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), + ("H", "HV", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), [ - pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) + pytest.param(*test, id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) for test in [ - (4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), - (4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), + (4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), + (4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), + (4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), + (4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), + (4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), # ======Varlen test with simulated trace======= ( + 32, 32, 128, 0, @@ -177,6 +184,7 @@ def test_safe_gate_chunk( True, ), ( + 32, 32, 128, 0, @@ -185,6 +193,7 @@ def test_safe_gate_chunk( True, ), ( + 32, 32, 128, 0, @@ -193,6 +202,20 @@ def test_safe_gate_chunk( True, ), ( + 32, + 32, + 128, + 0, + [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], + torch.bfloat16, + True, + ), + # ======GVA varlen cases: HV > H======= + (2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), + (4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), + (4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), + ( + 8, 32, 128, 0, @@ -205,6 +228,7 @@ def test_safe_gate_chunk( ) def test_safe_gate_chunk_varlen( H: int, + HV: int, D: int, mask_p: float, cu_seqlens: list[int], @@ -221,15 +245,15 @@ def test_safe_gate_chunk_varlen( q = torch.randn((1, T, H, D), dtype=dtype) k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) - v = torch.randn((1, T, H, D), dtype=dtype) - g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)) + v = torch.randn((1, T, HV, D), dtype=dtype) + g = F.logsigmoid(torch.randn(1, T, HV, D, dtype=torch.float)) mask = torch.rand_like(g) > mask_p g = g * mask + (~mask) * (-1000) if safe_gate: g = g.clamp(-5, 0) - beta = torch.randn(1, T, H, dtype=torch.float32).sigmoid().to(beta_dtype) - h0 = torch.randn((N, H, D, D), dtype=torch.float32) + beta = torch.randn(1, T, HV, dtype=torch.float32).sigmoid().to(beta_dtype) + h0 = torch.randn((N, HV, D, D), dtype=torch.float32) q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, beta, h0)) do = torch.randn_like(v) diff --git a/tests/test_kda_compare_fla.py b/tests/test_kda_compare_fla.py index 9d5e08dd..c88a40af 100644 --- a/tests/test_kda_compare_fla.py +++ b/tests/test_kda_compare_fla.py @@ -34,6 +34,7 @@ "B", "T", "H", + "HV", "D", "gate_logit_normalizer", "mask_p", @@ -45,17 +46,21 @@ [ pytest.param( *test, - id="B{}-T{}-H{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), + id="B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), ) for test in [ - (1, 63, 1, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 500, 3, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 1000, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), - (3, 1024, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 128, 1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (4, 2048, 8, 128, 1, 0, False, True, True, torch.bfloat16), + (1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), + (2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), + (2, 1000, 3, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), + (3, 1024, 4, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), + (4, 1024, 4, 4, 128, 1, 0, False, False, True, torch.bfloat16), + (4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), + (2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), + (4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), + # GVA cases: HV > H + (2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), + (2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), + (2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), ] ], ) @@ -63,6 +68,7 @@ def test_safe_gate_chunk( B: int, T: int, H: int, + HV: int, D: int, gate_logit_normalizer: float, mask_p: float, @@ -76,11 +82,11 @@ def test_safe_gate_chunk( torch.manual_seed(42) q = torch.rand(B, T, H, D, dtype=dtype) k = torch.rand(B, T, H, D, dtype=dtype) - v = torch.rand(B, T, H, D, dtype=dtype) - g = torch.randn(B, T, H, D, dtype=torch.float if not use_gate_in_kernel else dtype) + v = torch.rand(B, T, HV, D, dtype=dtype) + g = torch.randn(B, T, HV, D, dtype=torch.float if not use_gate_in_kernel else dtype) if use_gate_in_kernel: - A_log = torch.randn(H, dtype=torch.float) - dt_bias = torch.randn(H * D, dtype=torch.float) + A_log = torch.randn(HV, dtype=torch.float) + dt_bias = torch.randn(HV * D, dtype=torch.float) else: g = F.logsigmoid(g) / gate_logit_normalizer g = g * (torch.rand_like(g) > mask_p) @@ -91,8 +97,8 @@ def test_safe_gate_chunk( else: lower_bound = None - beta = torch.randn(B, T, H, dtype=torch.float32).sigmoid().to(beta_dtype) - h0 = torch.randn(B, H, D, D, dtype=torch.float32) + beta = torch.randn(B, T, HV, dtype=torch.float32).sigmoid().to(beta_dtype) + h0 = torch.randn(B, HV, D, D, dtype=torch.float32) if use_gate_in_kernel: A_log, dt_bias = map(lambda x: x.to(device).requires_grad_(True), (A_log, dt_bias)) q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) @@ -162,17 +168,18 @@ def test_safe_gate_chunk( @pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) @pytest.mark.parametrize("disable_recompute", [True, False], ids=["no_recomp", "recomp"]) @pytest.mark.parametrize( - ("H", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), + ("H", "HV", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), [ - pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) + pytest.param(*test, id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) for test in [ - (4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), - (4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), + (4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), + (4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), + (4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), + (4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), + (4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), # ======Varlen test with simulated trace======= ( + 32, 32, 128, 0, @@ -181,6 +188,7 @@ def test_safe_gate_chunk( True, ), ( + 32, 32, 128, 0, @@ -189,6 +197,7 @@ def test_safe_gate_chunk( True, ), ( + 32, 32, 128, 0, @@ -197,6 +206,20 @@ def test_safe_gate_chunk( True, ), ( + 32, + 32, + 128, + 0, + [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], + torch.bfloat16, + True, + ), + # ======GVA varlen cases: HV > H======= + (2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), + (4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), + (4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), + ( + 8, 32, 128, 0, @@ -209,6 +232,7 @@ def test_safe_gate_chunk( ) def test_safe_gate_chunk_varlen( H: int, + HV: int, D: int, mask_p: float, cu_seqlens: list[int], @@ -225,15 +249,15 @@ def test_safe_gate_chunk_varlen( q = torch.randn((1, T, H, D), dtype=dtype) k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) - v = torch.randn((1, T, H, D), dtype=dtype) - g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)) + v = torch.randn((1, T, HV, D), dtype=dtype) + g = F.logsigmoid(torch.randn(1, T, HV, D, dtype=torch.float)) mask = torch.rand_like(g) > mask_p g = g * mask + (~mask) * (-1000) if safe_gate: g = g.clamp(-5, 0) - beta = torch.randn(1, T, H, dtype=torch.float32).sigmoid().to(beta_dtype) - h0 = torch.randn((N, H, D, D), dtype=torch.float32) + beta = torch.randn(1, T, HV, dtype=torch.float32).sigmoid().to(beta_dtype) + h0 = torch.randn((N, HV, D, D), dtype=torch.float32) q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, beta, h0)) do = torch.randn_like(v) diff --git a/tests/test_kda_gva_intra_sm100.py b/tests/test_kda_gva_intra_sm100.py deleted file mode 100644 index 8082946e..00000000 --- a/tests/test_kda_gva_intra_sm100.py +++ /dev/null @@ -1,386 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for SM100 KDA GVA (HV > HQK) support in chunk_kda_fwd_intra. - -The SM100 kernels (kda_fwd_intra / kda_fwd_recomp_w_u) now accept: - * q, k with head-dim ``HQK`` - * v, g, beta with head-dim ``HV`` where ``HV = group_size * HQK`` (group_size >= 1) - -This file verifies that the cuLA GVA path produces numerically matching results -compared to the FLA Triton reference, where the FLA reference does not natively -support GVA and therefore receives ``k`` replicated along the head axis to -``HV`` heads. Both uniform-length and varlen layouts are covered, and an -additional degeneracy test asserts that ``HV == HQK`` (group_size == 1) keeps -the non-GVA behaviour untouched. -""" - -from __future__ import annotations - -import pytest -import torch -from einops import rearrange -from fla.modules.l2norm import l2norm_fwd -from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra -from fla.ops.kda.gate import kda_gate_chunk_cumsum -from fla.ops.utils.constant import RCP_LN2 -from fla.ops.utils.index import prepare_chunk_indices -from fla.utils import assert_close, device - -from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra -from cula.utils import prepare_uniform_cu_seqlens - -pytestmark = pytest.mark.sm100_only - - -# ========================================================================= -# Helpers -# ========================================================================= - -def _repeat_head(x: torch.Tensor, group_size: int, head_dim: int = 2) -> torch.Tensor: - """Replicate ``x`` along the head axis by ``group_size``. - - Mirrors GVA's broadcasting semantics: each QK head is paired with - ``group_size`` consecutive V heads, so ``k[..., h_qk, :]`` is used by - ``v[..., h_qk * group_size : (h_qk + 1) * group_size, :]``. - """ - return x.repeat_interleave(group_size, dim=head_dim).contiguous() - - -def _make_gva_inputs( - B: int, - T: int, - HQK: int, - HV: int, - D: int, - chunk_size: int, - cu_seqlens: torch.Tensor | None = None, - dtype: torch.dtype = torch.bfloat16, - seed: int = 42, -): - """Construct inputs for chunk_kda_fwd_intra in GVA layout. - - Returns: - q, k : (B, T, HQK, D) dtype - v : (B, T, HV, D) dtype - g : (B, T, HV, D) float32, after kda_gate_chunk_cumsum - beta : (B, T, HV) float32 in (0, 1) - scale : float - cu_seqlens : (N+1,) int32 or None - chunk_indices: (NT, 2) int32 or None - """ - assert HV % HQK == 0 and HV >= HQK, f"invalid HV/HQK: {HV}/{HQK}" - - torch.manual_seed(seed) - scale = D ** (-0.5) - - # QK are in HQK head space; V / gates / beta live in HV space. - q = torch.randn(B, T, HQK, D, dtype=dtype, device=device) - k = torch.randn(B, T, HQK, D, dtype=dtype, device=device) - v = torch.randn(B, T, HV, D, dtype=dtype, device=device) - g_raw = torch.randn(B, T, HV, D, dtype=dtype, device=device) - beta = torch.randn(B, T, HV, dtype=torch.float, device=device).sigmoid() - - # l2-normalise q/k so that scale/gate ranges match production use. - q, _ = l2norm_fwd(q) - k, _ = l2norm_fwd(k) - - # FLA gate cumsum only supports packed batch (B=1) when cu_seqlens is set. - if B != 1: - q, k, v, g_raw, beta = map( - lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), - (q, k, v, g_raw, beta), - ) - - # Per-HV gate preprocessing (cumsum inside chunks). - A_log = torch.randn(HV, dtype=torch.float, device=device) - dt_bias = torch.randn(HV * D, dtype=torch.float, device=device) - - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None - ) - g = kda_gate_chunk_cumsum( - g=g_raw, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=-5.0, - ) - return q, k, v, g, beta, scale, cu_seqlens, chunk_indices - - -def _run_fla_ref(q, k_hqk, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, group_size, disable_recompute): - """Reference: replicate k along head axis to HV, then call FLA intra. - - FLA's chunk_kda_fwd_intra assumes H == HQK == HV (no GVA), so we construct - the HV-head view of k and q before invoking it. - """ - k_hv = _repeat_head(k_hqk, group_size) - q_hv = _repeat_head(q, group_size) - return fla_chunk_kda_fwd_intra( - q=q_hv, - k=k_hv, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=disable_recompute, - ) - - -def _run_cula_gva(q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute): - return cula_chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=True, - disable_recompute=disable_recompute, - ) - - -def _assert_intra_outputs_match(ref, tri, disable_recompute: bool) -> None: - """Compare cuLA vs FLA on user-visible intra outputs. - - We intentionally skip ``Aqk``: the cuLA SM100 fused kernel does not - materialise every off-diagonal slot that FLA's multi-kernel path writes, - and the FLA reference can contain NaNs in unused ``Aqk`` entries. The - downstream tensors ``w`` / ``u`` / ``kg`` (and ``Akk``) are the meaningful - correctness signals and match the benchmark's comparison strategy. - """ - w_r, u_r, qg_r, kg_r, _Aqk_r, Akk_r = ref - w_c, u_c, qg_c, kg_c, _Aqk_c, Akk_c = tri - - assert Akk_c.shape == Akk_r.shape, (Akk_c.shape, Akk_r.shape) - assert w_c.shape == w_r.shape, (w_c.shape, w_r.shape) - assert u_c.shape == u_r.shape, (u_c.shape, u_r.shape) - assert kg_c.shape == kg_r.shape, (kg_c.shape, kg_r.shape) - - assert_close("Akk", Akk_r, Akk_c, 0.008) - assert_close("w", w_r, w_c, 0.008) - assert_close("u", u_r, u_c, 0.008) - assert_close("kg", kg_r, kg_c, 0.005) - - if disable_recompute: - assert qg_c is not None and qg_r is not None - assert qg_c.shape == qg_r.shape, (qg_c.shape, qg_r.shape) - assert_close("qg", qg_r, qg_c, 0.005) - else: - assert qg_c is None, "cuLA must not materialise qg when disable_recompute=False" - - -# ========================================================================= -# Uniform-length tests -# ========================================================================= - -@pytest.mark.parametrize("disable_recompute", [False, True], ids=["recomp", "no_recomp"]) -@pytest.mark.parametrize( - ("B", "T", "HQK", "group_size", "D"), - [ - pytest.param(*cfg, id="B{}-T{}-HQK{}-gs{}-D{}".format(*cfg)) - for cfg in [ - # group_size == 2: classic GVA 2:1 - (1, 256, 2, 2, 128), - (2, 512, 4, 2, 128), - # group_size == 4: wider grouping - (1, 1024, 2, 4, 128), - (2, 1024, 4, 4, 128), - # Non-multiple-of-BT sequence length to stress boundary handling. - (1, 500, 2, 2, 128), - (1, 1000, 4, 2, 128), - ] - ], -) -def test_gva_intra_uniform(B, T, HQK, group_size, D, disable_recompute): - """cuLA GVA path must match FLA(k-replicated-to-HV) for uniform seqlens.""" - HV = HQK * group_size - chunk_size = 64 - - cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) - q, k, v, g, beta, scale, cu_seqlens, chunk_indices = _make_gva_inputs( - B=B, T=T, HQK=HQK, HV=HV, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens, - ) - - # cuLA GVA path (k in HQK head space). - w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c = _run_cula_gva( - q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute, - ) - - # FLA reference (k replicated to HV). - w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r = _run_fla_ref( - q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, group_size, disable_recompute, - ) - - _assert_intra_outputs_match( - (w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r), - (w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c), - disable_recompute, - ) - - -# ========================================================================= -# Varlen tests -# ========================================================================= - -@pytest.mark.parametrize("disable_recompute", [False, True], ids=["recomp", "no_recomp"]) -@pytest.mark.parametrize( - ("HQK", "group_size", "D", "cu_seqlens"), - [ - pytest.param(*cfg, id="HQK{}-gs{}-D{}-ns{}".format(cfg[0], cfg[1], cfg[2], len(cfg[3]) - 1)) - for cfg in [ - (2, 2, 128, [0, 256, 500, 1000]), - (4, 2, 128, [0, 100, 300, 1200, 2000]), - (2, 4, 128, [0, 15, 100, 300, 1200, 2048]), - # Simulated realistic trace. - ( - 4, 2, 128, - [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096], - ), - ] - ], -) -def test_gva_intra_varlen(HQK, group_size, D, cu_seqlens, disable_recompute): - """GVA correctness under variable-length (packed) inputs.""" - HV = HQK * group_size - chunk_size = 64 - - cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) - T = int(cu_seqlens_t[-1].item()) - # Packed layout uses B=1 and a flat time axis. - q, k, v, g, beta, scale, cu_seqlens_t, chunk_indices = _make_gva_inputs( - B=1, T=T, HQK=HQK, HV=HV, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens_t, - ) - - w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c = _run_cula_gva( - q, k, v, g, beta, scale, cu_seqlens_t, chunk_indices, chunk_size, disable_recompute, - ) - w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r = _run_fla_ref( - q, k, v, g, beta, scale, cu_seqlens_t, chunk_indices, chunk_size, group_size, disable_recompute, - ) - - _assert_intra_outputs_match( - (w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r), - (w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c), - disable_recompute, - ) - - -# ========================================================================= -# Degeneracy: HV == HQK must match the non-GVA (same-shape) reference -# ========================================================================= - -@pytest.mark.parametrize("disable_recompute", [False, True], ids=["recomp", "no_recomp"]) -@pytest.mark.parametrize( - ("B", "T", "H", "D"), - [ - pytest.param(*cfg, id="B{}-T{}-H{}-D{}".format(*cfg)) - for cfg in [ - (1, 512, 4, 128), - (2, 1024, 4, 128), - ] - ], -) -def test_gva_intra_degenerate_equals_non_gva(B, T, H, D, disable_recompute): - """When HV == HQK, the GVA code path must be byte-for-byte equivalent - to the non-GVA path that existed before this change. - - We do not have a separate "non-GVA" entrypoint, but we can assert the - cuLA path matches FLA with *no* head replication (group_size=1), which - exercises the ``HV == HQK`` fast-path inside the new kernels. - """ - chunk_size = 64 - cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) - q, k, v, g, beta, scale, cu_seqlens, chunk_indices = _make_gva_inputs( - B=B, T=T, HQK=H, HV=H, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens, - ) - - w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c = _run_cula_gva( - q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute, - ) - # group_size=1 → no replication; identical input shape to cuLA. - w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r = fla_chunk_kda_fwd_intra( - q=q, k=k, v=v, gk=g, beta=beta, scale=scale, - cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - safe_gate=True, disable_recompute=disable_recompute, - ) - - _assert_intra_outputs_match( - (w_r, u_r, qg_r, kg_r, Aqk_r, Akk_r), - (w_c, u_c, qg_c, kg_c, Aqk_c, Akk_c), - disable_recompute, - ) - - -# ========================================================================= -# Shape / contract sanity checks (run even without a reference) -# ========================================================================= - -@pytest.mark.parametrize("group_size", [1, 2, 4]) -def test_gva_intra_output_shapes(group_size): - """All outputs of chunk_kda_fwd_intra must live in HV-head space.""" - B, T, HQK, D = 1, 256, 2, 128 - HV = HQK * group_size - chunk_size = 64 - cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) - q, k, v, g, beta, scale, cu_seqlens, chunk_indices = _make_gva_inputs( - B=B, T=T, HQK=HQK, HV=HV, D=D, chunk_size=chunk_size, cu_seqlens=cu_seqlens, - ) - w, u, qg, kg, Aqk, Akk = _run_cula_gva( - q, k, v, g, beta, scale, cu_seqlens, chunk_indices, chunk_size, disable_recompute=True, - ) - - assert Aqk.shape == (B, T, HV, chunk_size), Aqk.shape - assert Akk.shape == (B, T, HV, chunk_size), Akk.shape - assert w.shape == (B, T, HV, D), w.shape - assert u.shape == (B, T, HV, D), u.shape - assert kg.shape == (B, T, HV, D), kg.shape - assert qg is not None and qg.shape == (B, T, HV, D), (None if qg is None else qg.shape) - - -# ========================================================================= -# Negative / assertion tests -# ========================================================================= - -def test_gva_intra_rejects_non_multiple_ratio(): - """HV must be a positive integer multiple of HQK.""" - B, T, HQK, HV, D = 1, 128, 3, 5, 128 # 5 % 3 != 0 - chunk_size = 64 - cu_seqlens = prepare_uniform_cu_seqlens(B, T, torch.device(device), torch.int32) - # We intentionally do not use _make_gva_inputs because the assert fires - # before kernel launch on the python side. - dtype = torch.bfloat16 - q = torch.randn(B, T, HQK, D, dtype=dtype, device=device) - k = torch.randn(B, T, HQK, D, dtype=dtype, device=device) - v = torch.randn(B, T, HV, D, dtype=dtype, device=device) - g = torch.randn(B, T, HV, D, dtype=torch.float, device=device) - beta = torch.randn(B, T, HV, dtype=torch.float, device=device).sigmoid() - - with pytest.raises((AssertionError, RuntimeError), match=r"multiple|h_v"): - cula_chunk_kda_fwd_intra( - q=q, k=k, v=v, gk=g, beta=beta, scale=D ** -0.5, - cu_seqlens=cu_seqlens, chunk_size=chunk_size, - safe_gate=True, disable_recompute=False, - ) From 99d4bbc22627174cc0fd67061a780beaa7adf3d4 Mon Sep 17 00:00:00 2001 From: Kevinzz <2538015266@qq.com> Date: Tue, 26 May 2026 12:45:38 +0800 Subject: [PATCH 20/34] [KDA] add backward chunk_wy_dqkg kernel for SM10X (#74) * init sm100 bwd wy dqkg * integrate and pass test * change kdk compute order, better perf * change dgk compute order, 1% perf * increase A stage to 2, 12% latency reduction * tune wg sync, 2.7% latency reduction * move dg store to aux warp, 2.8% latency reduction * add tma store for non-tail chunk * store dq to tmem to reduce reg spill, 6.8% latency reduction * add more nan/inf tests * remove TS ws mode MMA * support GVA for wy_dqkg * use cute.arch.atomic_add and update check * fix db atomic add and store * change to deterministic db reduce * change iters * change to umma pipeline * change h, dh and v pipelines to different consumers * add tma store desc prefetch * skip deter check for ncu mode * modify deter check * fix * code lint * add nan and inf check * fix * add copyright * refactor wy and intrinsic * delete line-info * delete * refactor bench --------- Co-authored-by: boyu.zbw --- benchmarks/bench_kda_bwd_wy_dqkg_sm100.py | 454 +++ benchmarks/utils.py | 101 + cula/kda/chunk_bwd.py | 5 +- cula/ops/chunk_wy_dqkg_sm100.py | 3194 +++++++++++++++++++++ cula/ops/intrinsics_sm100.py | 418 +++ cula/ops/ptx_umma_ext.py | 961 +++++++ tests/test_ptx_umma_masked.py | 326 +++ tests/test_ptx_umma_ws.py | 591 ++++ 8 files changed, 6048 insertions(+), 2 deletions(-) create mode 100644 benchmarks/bench_kda_bwd_wy_dqkg_sm100.py create mode 100644 cula/ops/chunk_wy_dqkg_sm100.py create mode 100644 cula/ops/intrinsics_sm100.py create mode 100644 cula/ops/ptx_umma_ext.py create mode 100644 tests/test_ptx_umma_masked.py create mode 100644 tests/test_ptx_umma_ws.py diff --git a/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py b/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py new file mode 100644 index 00000000..77aeb779 --- /dev/null +++ b/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +bench_kda_bwd_wy_dqkg_sm100.py — Benchmark: cuLA CuTe DSL vs FLA Triton baseline + for chunk_kda_bwd_wy_dqkg_fused kernel + +Compares: + - Accuracy: relative_rms_error, relative max diff between cuLA and FLA outputs + - Performance: kernel execution time (ms) with CUDA events + +Modes: + - Fixed-length: B=1,2 with various T + - Varlen: variable-length sequences with different distributions + +Usage: + python bench_kda_bwd_wy_dqkg_sm100.py [--mode fixed|varlen|both] [--ncu] [--heads 32 64] + +With --ncu, warmup=1 and iters=1 for ncu profiling: + ncu --set full -o report python bench_kda_bwd_wy_dqkg_sm100.py --mode fixed --ncu +""" + +import argparse +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch +from fla.ops.kda.chunk_bwd import chunk_kda_bwd_wy_dqkg_fused as fla_chunk_kda_bwd_wy_dqkg_fused + +from benchmarks.utils import ( + SEED, + benchmark_cuda_mode_fn, + build_varlen_configs, + exclusive_cumsum, + prepare_bwd_wy_dqkg_fused_inputs, + relative_rms_error_rel_max_mean_abs, + set_seed, +) +from cula.ops.chunk_wy_dqkg_sm100 import chunk_kda_bwd_wy_dqkg_fused as cula_chunk_kda_bwd_wy_dqkg_fused + +torch.backends.cuda.matmul.allow_tf32 = True + +# ============================================================ +# Constants +# ============================================================ +H_DEFAULT = 32 +K = 128 +V = 128 +BT = 64 +DTYPE = torch.bfloat16 +DEVICE = torch.device("cuda") +WARMUP = 25 +N_ITERS = 100 +NCU_MODE = False + + +def generate_balanced_seqlens(total_tokens, num_seqs): + base = total_tokens // num_seqs + remainder = total_tokens % num_seqs + return [base] * (num_seqs - 1) + [base + remainder] + + +# ============================================================ +# Runners +# ============================================================ +def run_fla_triton(inputs: dict): + """Run the FLA Triton baseline.""" + return fla_chunk_kda_bwd_wy_dqkg_fused( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + v_new=inputs["v_new"], + g=inputs["g"], + beta=inputs["beta"], + A=inputs["A"], + h=inputs["h"], + do=inputs["do"], + dh=inputs["dh"], + dv=inputs["dv"], + scale=inputs["scale"], + cu_seqlens=inputs["cu_seqlens"], + chunk_size=BT, + chunk_indices=inputs["chunk_indices"], + transpose_state_layout=False, + ) + + +def run_cula(inputs: dict): + """Run the CuTe DSL Blackwell kernel.""" + return cula_chunk_kda_bwd_wy_dqkg_fused( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + v_new=inputs["v_new"], + g=inputs["g"], + beta=inputs["beta"], + A=inputs["A"], + h=inputs["h"], + do=inputs["do"], + dh=inputs["dh"], + dv=inputs["dv"], + scale=inputs["scale"], + cu_seqlens=inputs["cu_seqlens"], + chunk_size=BT, + chunk_indices=inputs["chunk_indices"], + ) + + +def check_determinism(H=4, HV=None, total_T=2001, num_seqs=4, iters=1000): + """Verify deterministic outputs across repeated runs.""" + if HV is None: + HV = H + torch.manual_seed(42) + seq_lens = generate_balanced_seqlens(total_T, num_seqs) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=DEVICE) + inputs = prepare_bwd_wy_dqkg_fused_inputs( + B=1, + T=total_T, + H=H, + K=K, + V=V, + HV=HV, + chunk_size=BT, + device=DEVICE, + seed=SEED, + cu_seqlens=cu_seqlens, + ) + + ref_dq, ref_dk, ref_dv, ref_db, ref_dg, ref_dA = run_cula(inputs) + for i in range(iters): + dq_out, dk_out, dv_out, db_out, dg_out, dA_out = run_cula(inputs) + assert torch.isnan(dq_out).sum() == 0, f"dq contains NaNs at iter {i}" + assert torch.isnan(dk_out).sum() == 0, f"dk contains NaNs at iter {i}" + assert torch.isnan(dv_out).sum() == 0, f"dv contains NaNs at iter {i}" + assert torch.isnan(db_out).sum() == 0, f"db contains NaNs at iter {i}" + assert torch.isnan(dg_out).sum() == 0, f"dg contains NaNs at iter {i}" + assert torch.isnan(dA_out).sum() == 0, f"dA contains NaNs at iter {i}" + assert torch.isfinite(dq_out).all(), f"dq contains infs at iter {i}" + assert torch.isfinite(dk_out).all(), f"dk contains infs at iter {i}" + assert torch.isfinite(dv_out).all(), f"dv contains infs at iter {i}" + assert torch.isfinite(db_out).all(), f"db contains infs at iter {i}" + assert torch.isfinite(dg_out).all(), f"dg contains infs at iter {i}" + assert torch.isfinite(dA_out).all(), f"dA contains infs at iter {i}" + assert torch.equal(dq_out, ref_dq), f"dq mismatch at iter {i}" + assert torch.equal(dk_out, ref_dk), f"dk mismatch at iter {i}" + assert torch.equal(dv_out, ref_dv), f"dv mismatch at iter {i}" + assert torch.equal(dg_out, ref_dg), f"dg mismatch at iter {i}" + assert torch.equal(dA_out, ref_dA), f"dA mismatch at iter {i}" + assert torch.equal(db_out, ref_db), f"db mismatch at iter {i}" + return True + + +# ============================================================ +# Fixed-length benchmark +# ============================================================ +def bench_fixed(configs, H: int, HV: int | None = None): + if HV is None: + HV = H + print("\n" + "=" * 120) + print(f" Fixed-Length Benchmark: cuLA CuTe DSL vs FLA Triton (H={H}, HV={HV}, K={K}, V={V}, BT={BT})") + print("=" * 120) + results = [] + + for B, T in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + seq_lens = [T] * B + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=DEVICE) + + inputs = prepare_bwd_wy_dqkg_fused_inputs( + B=B, + T=T, + H=H, + K=K, + V=V, + HV=HV, + chunk_size=BT, + device=DEVICE, + seed=SEED, + cu_seqlens=cu_seqlens, + ) + + # Accuracy + ref = run_fla_triton(inputs) # (dq, dk, dv, db, dg, dA) + out = run_cula(inputs) # (dq, dk, dv, db, dg, dA) + torch.cuda.synchronize() + + acc = {} + names = ["dq", "dk", "dv", "db", "dg", "dA"] + for name, r, o in zip(names, ref, out): + rel_rmse, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(r, o) + acc[name] = {"rel_rmse": rel_rmse, "rel_max": rel_max, "mean_diff": mean_diff} + + # Performance + ms_fla = benchmark_cuda_mode_fn( + lambda: run_fla_triton(inputs), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + lambda: run_cula(inputs), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + ) + speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") + + r = { + "B": B, + "T": T, + "accuracy": acc, + "ms_fla": ms_fla, + "ms_cula": ms_cula, + "speedup": speedup, + } + results.append(r) + + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Varlen benchmark +# ============================================================ +def bench_varlen(configs, H: int, HV: int | None = None): + if HV is None: + HV = H + print("\n" + "=" * 120) + print(f" Varlen Benchmark: cuLA CuTe DSL vs FLA Triton (H={H}, HV={HV}, K={K}, V={V}, BT={BT})") + print("=" * 120) + results = [] + + for seq_lens, total_len, dist in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + T = total_len + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=DEVICE) + + inputs = prepare_bwd_wy_dqkg_fused_inputs( + B=1, + T=T, + H=H, + K=K, + V=V, + HV=HV, + chunk_size=BT, + device=DEVICE, + seed=SEED, + cu_seqlens=cu_seqlens, + ) + + # Accuracy + ref = run_fla_triton(inputs) + out = run_cula(inputs) + torch.cuda.synchronize() + + acc = {} + names = ["dq", "dk", "dv", "db", "dg", "dA"] + for name, r, o in zip(names, ref, out): + rel_rmse, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(r, o) + acc[name] = {"rel_rmse": rel_rmse, "rel_max": rel_max, "mean_diff": mean_diff} + + # Performance + ms_fla = benchmark_cuda_mode_fn( + lambda: run_fla_triton(inputs), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + ) + ms_cula = benchmark_cuda_mode_fn( + lambda: run_cula(inputs), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + ) + speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") + + n_seqs = len(seq_lens) + min_l, max_l = min(seq_lens), max(seq_lens) + avg_l = T // n_seqs + tag = f"{dist:>7s} {n_seqs:>2d}seqs T={T} [{min_l}..{max_l}] avg={avg_l}" + + r = { + "tag": tag, + "dist": dist, + "T_total": T, + "n_seqs": n_seqs, + "accuracy": acc, + "ms_fla": ms_fla, + "ms_cula": ms_cula, + "speedup": speedup, + } + results.append(r) + + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Report +# ============================================================ +def print_report(fixed_results, varlen_results, H: int): + sep = "=" * 130 + print(f"\n\n{sep}") + print(" BENCHMARK REPORT: chunk_kda_bwd_wy_dqkg_fused") + print(" cuLA CuTe DSL vs FLA Triton") + wu = 1 if NCU_MODE else WARMUP + ni = 1 if NCU_MODE else N_ITERS + mode_tag = " [NCU mode]" if NCU_MODE else "" + print(f" H={H} K={K} V={V} BT={BT} dtype=bf16{mode_tag}") + print(f" Warmup={wu} Iters={ni}") + print(sep) + + acc_keys = ["dq", "dk", "dv", "db", "dg", "dA"] + acc_header = " ".join(f"{k:>10s}" for k in acc_keys) + + if fixed_results: + print("\n [Fixed-Length]") + print(f" {'─' * 125}") + print(f" {'B':>3s} {'T':>5s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>9s} {'Speedup':>8s} │ {'':>10s}{acc_header}") + print(f" {'─' * 125}") + + for r in fixed_results: + rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in acc_keys) + rel_rmse_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_rmse', 0.0):10.6f}" for k in acc_keys) + print( + f" {r['B']:3d} {r['T']:5d} │ " + f"{r['ms_fla']:9.4f} {r['ms_cula']:9.4f} {r['speedup']:7.2f}x │ " + f"{'rel_max:':>10s}{rel_max_vals}" + ) + print(f" {'':3s} {'':5s} │ {'':9s} {'':9s} {'':8s} │ {'rel_rmse:':>10s}{rel_rmse_vals}") + print(f" {'─' * 125}") + + if varlen_results: + print("\n [Varlen]") + print(f" {'─' * 140}") + print(f" {'Config':>45s} │ {'FLA(ms)':>9s} {'cuLA(ms)':>9s} {'Speedup':>8s} │ {'':>10s}{acc_header}") + print(f" {'─' * 140}") + + for r in varlen_results: + rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in acc_keys) + rel_rmse_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_rmse', 0.0):10.6f}" for k in acc_keys) + print( + f" {r['tag']:>45s} │ " + f"{r['ms_fla']:9.4f} {r['ms_cula']:9.4f} {r['speedup']:7.2f}x │ " + f"{'rel_max:':>10s}{rel_max_vals}" + ) + print(f" {'':>45s} │ {'':9s} {'':9s} {'':8s} │ {'rel_rmse:':>10s}{rel_rmse_vals}") + print(f" {'─' * 140}") + + print(f"\n{sep}\n") + + +# ============================================================ +# Main +# ============================================================ +def main(): + global NCU_MODE + + parser = argparse.ArgumentParser(description="Benchmark chunk_kda_bwd_wy_dqkg_fused: cuLA CuTe DSL vs FLA Triton") + parser.add_argument( + "--mode", + type=str, + default="both", + choices=["fixed", "varlen", "both"], + help="Which benchmark mode to run (default: both)", + ) + parser.add_argument( + "--heads", + nargs="+", + type=int, + default=[H_DEFAULT], + help=f"Head counts to benchmark (default: [{H_DEFAULT}])", + ) + parser.add_argument( + "--hv", + type=int, + default=None, + help="Number of value heads HV (default: same as H, i.e. no GVA). Must be a multiple of H.", + ) + parser.add_argument("--ncu", action="store_true", help="NCU profiling mode: warmup=1, iters=1") + args = parser.parse_args() + + if args.ncu: + NCU_MODE = True + print("[NCU mode] warmup=1, iters=1") + + gpu_name = torch.cuda.get_device_name(0) + print(f"GPU: {gpu_name}") + wu = 1 if NCU_MODE else WARMUP + ni = 1 if NCU_MODE else N_ITERS + print(f"K={K}, V={V}, BT={BT}, dtype={DTYPE}, warmup={wu}, rep={ni}") + + fixed_configs = [ + (1, 256), + (1, 512), + (1, 1024), + (1, 2048), + (1, 4096), + (1, 8192), + (2, 512), + (2, 1024), + (2, 2048), + (2, 4096), + (2, 8192), + ] + + varlen_configs = build_varlen_configs( + num_seqs_list=(10, 20), + total_lens=(4096, 8192, 16384), + dists=("uniform", "random", "skewed"), + ) + + for H in args.heads: + HV = args.hv if args.hv is not None else H + if not args.ncu: + check_determinism(H=H, HV=HV, iters=10000) + + fixed_res, varlen_res = [], [] + + if args.mode in ("fixed", "both"): + fixed_res = bench_fixed(fixed_configs, H, HV) + + if args.mode in ("varlen", "both"): + varlen_res = bench_varlen(varlen_configs, H, HV) + + print_report(fixed_res, varlen_res, H) + + print(f"\n{'=' * 130}") + print(" All benchmarks done.") + print(f"{'=' * 130}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 406093a6..55602490 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -536,3 +536,104 @@ def prepare_intra_inputs(batch_size, T, H, D, device, cu_seqlens=None, chunk_siz ) return q, k, v, g, beta, scale, cu_seqlens, chunk_indices + + +def prepare_bwd_wy_dqkg_fused_inputs( + B: int, + T: int, + H: int, + K: int, + V: int, + HV: int | None = None, + chunk_size: int = CHUNK_SIZE, + device: torch.device | str = "cuda", + seed: int = SEED, + cu_seqlens: torch.Tensor | None = None, + dtype: torch.dtype = torch.bfloat16, +) -> dict: + """Prepare all inputs needed by the bwd_wy_dqkg_fused benchmark runners. + + Generates the full set of tensors consumed by both the FLA Triton and CuTe DSL + chunk_kda_bwd_wy_dqkg_fused kernels. Follows the same flattening convention + used in other prepare_* helpers (B=1 with cu_seqlens for varlen mode). + + HV: number of value heads (default: H). Set HV > H for GVA (grouped value attention). + q/k always have H heads; all other tensors use HV heads. + + Returns a dict with keys used directly by ``run_fla_triton`` and ``run_cutedsl`` + in ``bench_bwd_wy_dqkg_fused.py``. + """ + if HV is None: + HV = H + BT = chunk_size + scale = K**-0.5 + + set_seed(seed) + + # ---- primary token-indexed tensors ---- + q = torch.randn(B, T, H, K, dtype=dtype, device=device) + k = torch.randn(B, T, H, K, dtype=dtype, device=device) + v = torch.randn(B, T, HV, V, dtype=dtype, device=device) + g_raw = torch.randn(B, T, HV, K, dtype=dtype, device=device) + beta = torch.randn(B, T, HV, dtype=torch.float, device=device).sigmoid() + + # l2norm q, k + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + + # gate preprocessing + A_log = torch.randn(HV, dtype=torch.float, device=device) + dt_bias = torch.randn(HV * K, dtype=torch.float, device=device) + + v_new = torch.randn(B, T, HV, V, dtype=dtype, device=device) + do = torch.randn(B, T, HV, V, dtype=dtype, device=device) + dv = torch.randn(B, T, HV, V, dtype=dtype, device=device) + A = torch.randn(B, T, HV, BT, dtype=dtype, device=device) * 0.1 + + # ---- chunk-indexed state tensors ---- + if cu_seqlens is not None: + cu_seqlens = cu_seqlens.int() + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = chunk_indices.shape[0] + else: + NT = (B * T + BT - 1) // BT + chunk_indices = None + + # h/dh: both FLA Triton and CuTe DSL use bf16 [B, NT, HV, K, V] + h = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + dh = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + + # flatten to batch_size=1 for cu_seqlens compatibility + if B != 1: + q, k = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (q, k)) + v, g_raw, beta = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (v, g_raw, beta)) + v_new, do, dv, A = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (v_new, do, dv, A)) + h, dh = map(lambda x: rearrange(x, "b nt ... -> 1 (b nt) ..."), (h, dh)) + + g = kda_gate_chunk_cumsum( + g=g_raw, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=-5.0, + ) + + return dict( + q=q, + k=k, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=A, + h=h, + dh=dh, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) diff --git a/cula/kda/chunk_bwd.py b/cula/kda/chunk_bwd.py index 738b1589..d4cf6ee7 100644 --- a/cula/kda/chunk_bwd.py +++ b/cula/kda/chunk_bwd.py @@ -37,6 +37,7 @@ import cula.cudac as cula_cuda from cula.kda.chunk_intra import chunk_kda_bwd_intra +from cula.ops.chunk_wy_dqkg_sm100 import chunk_kda_bwd_wy_dqkg_fused as chunk_kda_bwd_wy_dqkg_fused_cutedsl from cula.utils import prepare_uniform_cu_seqlens _delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") @@ -567,7 +568,7 @@ def chunk_kda_bwd( transpose_state_layout=transpose_state_layout, ) - dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused( + dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused_cutedsl( q=q, k=k, v=v, @@ -583,7 +584,7 @@ def chunk_kda_bwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - transpose_state_layout=transpose_state_layout, + # transpose_state_layout=transpose_state_layout, ) dq, dk, db, dg = chunk_kda_bwd_intra( diff --git a/cula/ops/chunk_wy_dqkg_sm100.py b/cula/ops/chunk_wy_dqkg_sm100.py new file mode 100644 index 00000000..cafbb247 --- /dev/null +++ b/cula/ops/chunk_wy_dqkg_sm100.py @@ -0,0 +1,3194 @@ +import argparse + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass.cute.arch import ( + elect_one, + mbarrier_arrive, + mbarrier_arrive_and_expect_tx, + mbarrier_init, + mbarrier_init_fence, + mbarrier_wait, +) +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05 import ( + make_umma_smem_desc, + smem_descriptor_to_int, +) +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream +from cutlass.cute.tensor import TensorSSA +from cutlass.cute.typing import BFloat16, Float32, Int32, Int64 +from fla.ops.utils import prepare_chunk_indices + +from cula.ops.intrinsics_sm100 import ( + reinterpret_cast, + store_256b, + subvec, + tcgen05_fence_after, + tcgen05_fence_before, + tcgen05_ld_32x32b, + tcgen05_st_32x32b, + umma_arrive, +) +from cula.ops.ptx_umma_ext import ( + Tcgen05SmemDescriptor, + tcgen05mma_ws_ss_f16, +) +from cula.utils import USE_FAST_MATH, assert_blackwell, prepare_uniform_cu_seqlens + +PRINT_DEBUG = False + +LN2 = 0.6931471805599453 +RCP_LN2 = 1.4426950408889634 + +COMPILE_OPTIONS = "--enable-tvm-ffi" + +# Mapping from torch dtype to cutlass dtype (for beta_dtype conversion) +_torch_to_cutlass_dtype = { + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, +} + + +def make_thread_cooperative_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +def _exclusive_cumsum(a: list[int]): + r = [0] + for v in a: + r.append(r[-1] + v) + return r + + +# ── TMEM column offset constants (cta_group::1, M=64, .ws Layout E) ── +TMEM_DA_ACC_OFF = 0 # [0,32) 32 cols dA fp32 acc; Phase 3: [0,16) overwritten by dA_bf16 +TMEM_DQ_ACC_OFF = 32 # [32,96) 64 cols dq fp32 acc; Phase 3: step2/step3 result [32,64) +TMEM_DK_ACC_OFF = 96 # [96,160) 64 cols dk fp32 acc +TMEM_DW_ACC_OFF = 160 # [160,224] 64 cols dw fp32 acc +TMEM_FLEX_OFF = 224 # [224,256) 32 cols dvb time-shared +TMEM_A_BF16_OFF = 256 # [256,272) 16 cols A_bf16 TS opA (persistent) (not used currently) +TMEM_DKGB_ACC_OFF = 272 # [272,336) 64 cols, dkgb fp32 acc +TMEM_DA2_ACC_OFF = 336 # [336,368) 32 cols dA fp32 acc, used for dA=dA@A and dA=A@dA +TMEM_DQ_SCALED_OFF = 368 # [368,432) 64 cols dq_scaled (stored for dg) +TMEM_TOTAL = 512 + +# Instruction descriptor for M=64, N=64, BF16, dense, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=8 at [17:22], TransposeB at [16], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N64_K_MN = (4 << 24) | (8 << 17) | (1 << 16) | (1 << 10) | (1 << 7) | (1 << 4) + +# Instruction descriptor for M=64, N=128, BF16, dense, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=16 at [17:22], TransposeB at [16], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N128_K_MN = (4 << 24) | (16 << 17) | (1 << 16) | (1 << 10) | (1 << 7) | (1 << 4) + +# Instruction descriptor for M=64, N=128, BF16, dense +# Bits: M>>4=4 at [24:28], N>>3=16 at [17:22], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N128_K_K = (4 << 24) | (16 << 17) | (1 << 10) | (1 << 7) | (1 << 4) + +# Instruction descriptor for M=64, N=128, BF16, dense, TransposeA=1, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=16 at [17:22], +# TransposeB at [16], TransposeA at [15], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N128_MN_MN = (4 << 24) | (16 << 17) | (1 << 16) | (1 << 15) | (1 << 10) | (1 << 7) | (1 << 4) + +# Instruction descriptor for M=64, N=64, BF16, dense, TransposeA=1, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=8 at [17:22], +# TransposeB at [16], TransposeA at [15], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N64_MN_MN = (4 << 24) | (8 << 17) | (1 << 16) | (1 << 15) | (1 << 10) | (1 << 7) | (1 << 4) + +# Instruction descriptor for M=64, N=64, BF16, dense +# Bits: M>>4=4 at [24:28], N>>3=8 at [17:22], +# TransposeB at [16], TransposeA at [15], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N64_K_K = (4 << 24) | (8 << 17) | (1 << 10) | (1 << 7) | (1 << 4) + +ELEM_BYTES_BF16 = BFloat16.width // 8 + + +@cute.jit +def smem_load_bf16x8_sw128(raw_ptr: cute.Pointer, row: Int32, col_base: Int32): + """ + Load 8 consecutive bfloat16 from SMEM with Swizzle<3,4,3> layout. + raw_ptr: BFloat16 SMEM base pointer (NOT recast_ptr — raw buffer start) + row: row index in [0, T_TILE=64) + col_base: 8-aligned column index in [0, K_TILE=128) + Logical layout: [BT=64, BV=128] K-major, with the BV=128 dim split into + two halves of 64 elements (high half offset by 4096 elements). + Swizzle<3,4,3> on bf16: phys_elem = elem ^ ((row & 7) << 3) within a half. + Returns an 8-element rmem fragment (bf16). + """ + half = col_base >> Int32(6) + k_inner = col_base & Int32(63) + swizzled = k_inner ^ ((row & Int32(7)) << Int32(3)) + elem_off = half * Int32(4096) + row * Int32(64) + swizzled + aligned_ptr = cute.make_ptr( + BFloat16, + (raw_ptr + elem_off).toint(), + cute.AddressSpace.smem, + assumed_align=16, + ) + smem_t = cute.make_tensor(aligned_ptr, cute.make_layout((8,), stride=(1,))) + rmem_t = cute.make_fragment_like(smem_t) + cute.autovec_copy(smem_t, rmem_t) + return rmem_t + + +@cute.jit +def smem_store_bf16x8_sw128(raw_ptr: cute.Pointer, row: Int32, col_base: Int32, data: cute.Tensor): + """ + Store 8 consecutive bfloat16 to SMEM with Swizzle<3,4,3> layout. + raw_ptr: BFloat16 SMEM base pointer (NOT recast_ptr — raw buffer start) + row: row index in [0, T_TILE=64) + col_base: 8-aligned column index in [0, K_TILE=128) + data: 8-element rmem fragment (bf16) to store. + + NOTE: For the K-major→MN-major dv re-swizzle, source layout + `(BT,BV) K-major Swizzle<3,4,3>` and destination layout + `(BV,BT) MN-major Swizzle<3,4,3>` produce **identical** physical + addresses for the same (row=t, col=v). So this helper uses the same + address formula as the load helper, and the caller passes (row=t, col=v) + for both load (src K-maj) and store (dst MN-maj), implicitly transposing. + """ + half = col_base >> Int32(6) + k_inner = col_base & Int32(63) + swizzled = k_inner ^ ((row & Int32(7)) << Int32(3)) + elem_off = half * Int32(4096) + row * Int32(64) + swizzled + smem_ptr = cute.make_ptr( + BFloat16, + (raw_ptr + elem_off).toint(), + cute.AddressSpace.smem, + assumed_align=16, + ) + smem_t = cute.make_tensor(smem_ptr, cute.make_layout((8,), stride=(1,))) + cute.autovec_copy(data, smem_t) + + +@cute.jit +def smem_load_f32x4_sw128(raw_ptr: cute.Pointer, row: Int32, col_base: Int32): + """ + Load 4 consecutive float32 from SMEM with K_SW128 layout. + Logical layout: [BT=64, BK=128] ROW_MAJOR, tiled over a Float32 K_SW128 atom. + The atom provides a 32-element row stride. The 128-element column is broken + into 4 blocks of 32 elements. + PyCutlass tiles this such that outer blocks stride by 2048 elements: + elem_idx = row * 32 + (col_base % 32) + (col_base / 32) * 2048 + + The TMA hardware performs a 128B Swizzle on physical byte addresses: + byte_idx = elem_idx * 4 + swizzled_byte = byte_idx ^ (((byte_idx >> 7) & 7) << 4) + Dividing by 4 yields the element XOR offset: + elem_xor = ((elem_idx >> 5) & 7) << 2 + Because (elem_idx >> 5) simplifies to 'row + (col_outer * 64)', + the XOR offset simplifies exactly to ((row & 7) << 2). + This only affects the inner 32-element column block. + """ + c_inner = col_base & Int32(31) + c_outer = col_base >> Int32(5) + swizzled_inner = c_inner ^ ((row & Int32(7)) << Int32(2)) + + elem_offset = row * Int32(32) + swizzled_inner + c_outer * Int32(2048) + + aligned_ptr = cute.make_ptr( + Float32, + (raw_ptr + elem_offset).toint(), + cute.AddressSpace.smem, + assumed_align=16, + ) + t = cute.make_tensor(aligned_ptr, cute.make_layout((4,), stride=(1,))) + vals = t.load() + return (vals[0], vals[1], vals[2], vals[3]) + + +@cute.jit +def smem_store_f32x4_sw128(raw_ptr: cute.Pointer, row: Int32, col_base: Int32, data: cute.Tensor): + """ + Store 4 consecutive float32 to SMEM with K_SW128 layout. + Inverse of smem_load_f32x4_sw128 — same address formula, write path. + raw_ptr: Float32 SMEM base pointer (raw buffer start) + row: row index in [0, BT) + col_base: 4-aligned column index (multiples of 4) + data: 4-element rmem fragment (f32) to store. + """ + c_inner = col_base & Int32(31) + c_outer = col_base >> Int32(5) + swizzled_inner = c_inner ^ ((row & Int32(7)) << Int32(2)) + elem_offset = row * Int32(32) + swizzled_inner + c_outer * Int32(2048) + smem_ptr = cute.make_ptr( + Float32, + (raw_ptr + elem_offset).toint(), + cute.AddressSpace.smem, + assumed_align=16, + ) + smem_t = cute.make_tensor(smem_ptr, cute.make_layout((4,), stride=(1,))) + cute.autovec_copy(data, smem_t) + + +@cute.jit +def mma_ws_ss_m64n128_call( + a_smem_layout: cute.Layout, + desc_a_base: Tcgen05SmemDescriptor, + b_smem_layout: cute.Layout, + desc_b_base: Tcgen05SmemDescriptor, + tmem_c: Int32, + K: Int32, + is_accum: bool = False, +): + with elect_one(): + a_outer = a_smem_layout.outer + b_outer = b_smem_layout.outer + scale = 0 if not is_accum else 1 + for ks in cutlass.range_constexpr(K // 16): + a_off = cute.crd2idx(((0, 0), 0, ks, 0), a_outer) * ELEM_BYTES_BF16 + b_off = cute.crd2idx(((0, 0), 0, ks, 0), b_outer) * ELEM_BYTES_BF16 + desc_a = desc_a_base + a_off + desc_b = desc_b_base + b_off + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, IDESC_F16_M64_N128_K_MN, scale) + scale = 1 + + +@cute.jit +def mma_ws_ss_m64n128_k_k_call( + a_smem_layout: cute.Layout, + desc_a_base: Tcgen05SmemDescriptor, + b_smem_layout: cute.Layout, + desc_b_base: Tcgen05SmemDescriptor, + tmem_c: Int32, + K: Int32, + is_accum: bool = False, +): + with elect_one(): + a_outer = a_smem_layout.outer + b_outer = b_smem_layout.outer + scale = 0 if not is_accum else 1 + for ks in cutlass.range_constexpr(K // 16): + a_off = cute.crd2idx(((0, 0), 0, ks, 0), a_outer) * ELEM_BYTES_BF16 + b_off = cute.crd2idx(((0, 0), 0, ks, 0), b_outer) * ELEM_BYTES_BF16 + desc_a = desc_a_base + a_off + desc_b = desc_b_base + b_off + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, IDESC_F16_M64_N128_K_K, scale) + scale = 1 + + +@cute.jit +def mma_ws_ss_m64n128_mn_mn_call( + a_smem_layout: cute.Layout, + desc_a_base: Tcgen05SmemDescriptor, + b_smem_layout: cute.Layout, + desc_b_base: Tcgen05SmemDescriptor, + tmem_c: Int32, + K: Int32, + is_accum: bool = False, +): + with elect_one(): + a_outer = a_smem_layout.outer + b_outer = b_smem_layout.outer + scale = 0 if not is_accum else 1 + for ks in cutlass.range_constexpr(K // 16): + a_off = cute.crd2idx(((0, 0), 0, ks, 0), a_outer) * ELEM_BYTES_BF16 + b_off = cute.crd2idx(((0, 0), 0, ks, 0), b_outer) * ELEM_BYTES_BF16 + desc_a = desc_a_base + a_off + desc_b = desc_b_base + b_off + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, IDESC_F16_M64_N128_MN_MN, scale) + scale = 1 + + +@cute.jit +def mma_ws_ss_m64n64_k_k_call( + a_smem_layout: cute.Layout, + desc_a_base: Tcgen05SmemDescriptor, + b_smem_layout: cute.Layout, + desc_b_base: Tcgen05SmemDescriptor, + tmem_c: Int32, + K: Int32, + is_accum: bool = False, +): + with elect_one(): + a_outer = a_smem_layout.outer + b_outer = b_smem_layout.outer + scale = 0 if not is_accum else 1 + for ks in cutlass.range_constexpr(K // 16): + a_off = cute.crd2idx(((0, 0), 0, ks, 0), a_outer) * ELEM_BYTES_BF16 + b_off = cute.crd2idx(((0, 0), 0, ks, 0), b_outer) * ELEM_BYTES_BF16 + desc_a = desc_a_base + a_off + desc_b = desc_b_base + b_off + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, IDESC_F16_M64_N64_K_K, scale) + scale = 1 + + +@cute.jit +def mma_ws_ss_m64n64_mn_mn_call( + a_smem_layout: cute.Layout, + desc_a_base: Tcgen05SmemDescriptor, + b_smem_layout: cute.Layout, + desc_b_base: Tcgen05SmemDescriptor, + tmem_c: Int32, + K: Int32, + is_accum: bool = False, +): + with elect_one(): + a_outer = a_smem_layout.outer + b_outer = b_smem_layout.outer + scale = 0 if not is_accum else 1 + for ks in cutlass.range_constexpr(K // 16): + a_off = cute.crd2idx(((0, 0), 0, ks, 0), a_outer) * ELEM_BYTES_BF16 + b_off = cute.crd2idx(((0, 0), 0, ks, 0), b_outer) * ELEM_BYTES_BF16 + desc_a = desc_a_base + a_off + desc_b = desc_b_base + b_off + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, IDESC_F16_M64_N64_MN_MN, scale) + scale = 1 + + +class ChunkKdaBwdWyDqkgFused: + """ + CuTe DSL kernel for chunk_kda_bwd_kernel_wy_dqkg_fused. + + Computes backward gradients dq, dk, dv2, dg, db, dA for the KDA + chunkwise delta-rule backward pass. + + Architecture: 1 CudaCore WG + 1 MMA warp + TMA/Aux warps. + """ + + def __init__( + self, + chunk_size: int = 64, + head_dim_k: int = 128, + head_dim_v: int = 128, + acc_dtype: type[cutlass.Numeric] = cutlass.Float32, + io_dtype: type[cutlass.Numeric] = cutlass.BFloat16, + g_dtype: type[cutlass.Numeric] = cutlass.Float32, + beta_dtype: type[cutlass.Numeric] = cutlass.Float32, + scale: float = 1.0, + min_occupancy: int = 1, + use_fast_math: bool = True, + ): + assert chunk_size == 64, "chunk_size must be 64" + assert head_dim_k == 128 and head_dim_v == 128, ( + f"head_dim_k and head_dim_v must both be 128, got head_dim_k={head_dim_k}, head_dim_v={head_dim_v}" + ) + assert_blackwell() + + self.use_fast_math = use_fast_math + self.chunk_size = chunk_size + self.head_dim_k = head_dim_k + self.head_dim_v = head_dim_v + self.acc_dtype = acc_dtype + self.io_dtype = io_dtype + self.g_dtype = g_dtype + self.beta_dtype = beta_dtype + self.scale = scale + + # Tile sizes + self.BT = chunk_size # 64 + self.BK = 128 # K tiling for V-loop GEMM (single K tile) + self.BV = 64 # V tiling for V-loop GEMM (single V tile) + + # Warp layout: WG0/WG1 (8 CudaCore warps) + WG2 (MMA/Load/Aux/Store) + self.threads_per_warp = 32 + self.cuda_warp_ids = (0, 1, 2, 3) # WG0: CudaCore + Store + self.cuda2_warp_ids = (4, 5, 6, 7) # WG1: CudaCore + Store + self.mma_warp_id = 8 # WG2: MMA dispatch + self.load_warp_id = 9 # WG2: TMA Load + self.aux_warp_ids = (10, 11) # WG2: Aux/Load/Store Aux + self.threads_per_cta = self.threads_per_warp * 12 # 384 threads (3 WGs) + + self.num_regs_cuda = 208 + self.num_regs_others = 88 + self.min_occupancy = min_occupancy + + self.cluster_shape_mnk = (1, 1, 1) + self.cta_group = tcgen05.CtaGroup.ONE + + # Number of K/V tiles + self.num_k_tiles = (head_dim_k + self.BK - 1) // self.BK # 128/128 = 1 + self.num_v_tiles = (head_dim_v + self.BV - 1) // self.BV # 128/64 = 2 + + # ── Pipeline stages ── + # V-loop TMA: 2-stage double buffer + self.vloop_stage = 2 + self.kloop_stage = 1 + self.a_stage = 2 + self.mma_stage = 1 + + # ── MMA tiler shapes ── + # V-loop GEMMs: [BT, BV] × [BV, BK] → [BT, BK] + # dq = do @ h : (BT, BK, BV) — M=BT, N=BK, K=BV + # dk = v_new @ dh : (BT, BK, BV) + # dw = dv @ h : (BT, BK, BV) + self.vloop_gemm_tiler = (self.BT, self.BK, self.BV) + + # V-loop i_k==0 GEMMs: [BT, BV] × [BV, BT] → [BT, BT] + # dA = dv @ v^T : (BT, BT, BV) + self.dA_vloop_tiler = (self.BT, self.BT, self.BV) + + # V-loop i_k==0: A @ dv : [BT, BT] × [BT, BV] → [BT, BV] + self.dvb_tiler = (self.BT, self.BV, self.BT) + + # K-loop GEMMs: + # dA += dw @ kg^T : [BT, BK] × [BK, BT] → [BT, BT] → (BT, BT, BK) + self.kloop_dA_tiler = (self.BT, self.BT, self.BK) + # dkgb = A @ dw : [BT, BT] × [BT, BK] → [BT, BK] → (BT, BK, BT) + self.kloop_dkgb_tiler = (self.BT, self.BK, self.BT) + + # dA-post GEMMs: + # dA @ A : [BT, BT] × [BT, BT] → [BT, BT] → (BT, BT, BT) + # A @ dA : same + self.dApost_tiler = (self.BT, self.BT, self.BT) + + # Named barriers + self.tmem_dealloc_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_cta, + ) + self.cuda_wg_sync_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * 8, + ) + self.buffer_align_bytes = 1024 + + # Persistent scheduling + self.persistent = True + hardware_info = cutlass.utils.HardwareInfo() + self.num_sm = hardware_info.get_device_multiprocessor_count() + + def _compute_grid(self, B, T, HV, total_nt=None): + """Compute grid dimensions for persistent kernel launch. + + Grid: (min(num_sm * min_occupancy, total_tiles), 1, 1) + Each CTA handles multiple tiles via stride-by-gridDim.x loop. + """ + assert total_nt is not None + total_tiles = total_nt * HV + grid_x = cutlass.min(Int32(self.num_sm * self.min_occupancy), total_tiles) + return (grid_x, Int32(1), Int32(1)) + + @cute.jit + def __call__( + self, + # ── Inputs ── + q_in: cute.Tensor, # [B, T, H, K] bf16 + k_in: cute.Tensor, # [B, T, H, K] bf16 + v_in: cute.Tensor, # [B, T, HV, V] bf16 + v_new_in: cute.Tensor, # [B, T, HV, V] bf16 + g_in: cute.Tensor, # [B, T, HV, K] fp32 + beta_in: cute.Tensor, # [B, T, HV] fp32 + A_in: cute.Tensor, # [B, T, HV, BT] bf16 + h_in: cute.Tensor, # [B, NT, HV, K, V] bf16 + do_in: cute.Tensor, # [B, T, HV, V] bf16 + dh_in: cute.Tensor, # [B, NT, HV, K, V] bf16 + dv_in: cute.Tensor, # [B, T, HV, V] bf16 + # ── Outputs ── + dq_in: cute.Tensor, # [B, T, HV, K] fp32 + dk_in: cute.Tensor, # [B, T, HV, K] fp32 + dv2_in: cute.Tensor, # [B, T, HV, V] bf16 + dg_in: cute.Tensor, # [B, T, HV, K] fp32 + db_in: cute.Tensor, # [B, T, HV] fp32 + dA_in: cute.Tensor, # [B, T, HV, BT] fp32 + # ── Metadata ── + cu_seqlens_in: cute.Tensor, # [N+1] int32 + chunk_indices_in: cute.Tensor, # [NT, 2] int32 + problem_size: tuple[Int32, Int32, Int32, Int32, Int32, Int32], # (B, T, H, HV, K, V) + total_nt: Int32, + stream, + ): + # ── Extract pointers ── + q_ptr = q_in.iterator + k_ptr = k_in.iterator + v_ptr = v_in.iterator + v_new_ptr = v_new_in.iterator + g_ptr = g_in.iterator + beta_ptr = beta_in.iterator + A_ptr = A_in.iterator + h_ptr = h_in.iterator + do_ptr = do_in.iterator + dh_ptr = dh_in.iterator + dv_ptr = dv_in.iterator + dq_ptr = dq_in.iterator + dk_ptr = dk_in.iterator + dv2_ptr = dv2_in.iterator + dg_ptr = dg_in.iterator + db_ptr = db_in.iterator + dA_ptr = dA_in.iterator + cu_seqlens_ptr = cu_seqlens_in.iterator + chunk_indices_ptr = chunk_indices_in.iterator + + B, T, H, HV, K, V = problem_size + BT = self.BT + + data_B = Int32(1) + NT = total_nt + + # ===================== GMEM layouts ===================== + # Token-indexed tensors: (T, dim, (H, data_B)) + # q, k: (T, K, (H, data_B)) bf16 + qk_layout = cute.make_layout( + (T, K, (H, data_B)), + stride=(H * K, 1, (K, T * H * K)), + ) + q = cute.make_tensor(q_ptr, qk_layout) + k = cute.make_tensor(k_ptr, qk_layout) + + # v, v_new, do, dv, dv2: (T, V, (HV, data_B)) bf16 + tv_layout = cute.make_layout( + (T, V, (HV, data_B)), + stride=(HV * V, 1, (V, T * HV * V)), + ) + v = cute.make_tensor(v_ptr, tv_layout) + v_new = cute.make_tensor(v_new_ptr, tv_layout) + do = cute.make_tensor(do_ptr, tv_layout) + dv = cute.make_tensor(dv_ptr, tv_layout) + dv2 = cute.make_tensor(dv2_ptr, tv_layout) + + # g: (T, K, (HV, data_B)) fp32 + g_layout = cute.make_layout( + (T, K, (HV, data_B)), + stride=(HV * K, 1, (K, T * HV * K)), + ) + g = cute.make_tensor(g_ptr, g_layout) + + # beta: (T, (HV, data_B)) fp32 + beta_layout = cute.make_layout( + (T, (HV, data_B)), + stride=(HV, (1, T * HV)), + ) + beta = cute.make_tensor(beta_ptr, beta_layout) + + # A: (T, BT, (HV, data_B)) bf16 + # NOTE: for A as operand A, A is loaded as transposed view to do MMA + a_t_layout = cute.make_layout( + (BT, T, (HV, data_B)), + stride=(1, HV * BT, (BT, T * HV * BT)), + ) + A_T = cute.make_tensor(A_ptr, a_t_layout) + + # dq, dk: (T, K, (HV, data_B)) fp32 + dqk_layout = cute.make_layout( + (T, K, (HV, data_B)), + stride=(HV * K, 1, (K, T * HV * K)), + ) + dq = cute.make_tensor(dq_ptr, dqk_layout) + dk = cute.make_tensor(dk_ptr, dqk_layout) + + # dg: (T, K, (HV, data_B)) fp32 + dg = cute.make_tensor(dg_ptr, dqk_layout) + + # db: (T, (HV, data_B)) fp32 + db = cute.make_tensor(db_ptr, beta_layout) + + # dA: (T, BT, (HV, data_B)) fp32 + dA_layout = cute.make_layout( + (T, BT, (HV, data_B)), + stride=(HV * BT, 1, (BT, T * HV * BT)), + ) + dA_out = cute.make_tensor(dA_ptr, dA_layout) + + h_nt_total = NT + + # h row-major: (K, V, (h_nt_total, HV)) as operand B + h_layout = cute.make_layout( + (K, V, (h_nt_total, HV)), + stride=(V, 1, (HV * K * V, K * V)), + ) + h = cute.make_tensor(h_ptr, h_layout) + dh = cute.make_tensor(dh_ptr, h_layout) + + # ===================== MMA setup (4 objects) ===================== + # All use tcgen05.mma.ws (Layout E, M=64, cta_group::1). + # 1. vloop_tiled_mma: SS K,K (64,128) — dq, dk, dw + # dq += do @ h, dk += vnew @ dh, dw += dv @ h + vloop_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, # A: K-major + tcgen05.OperandMajorMode.K, # B: K-major + self.acc_dtype, + self.cta_group, + self.vloop_gemm_tiler[:2], # (64, 128) + # default a_source=OperandSource.SMEM → SS mode + ) + + # 2. dA_vloop_tiled_mma: SS K,K (64,64) — dA vloop + kpost_dA + # dA += dv @ v^T, dA += dw @ kg^T + dA_vloop_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + self.cta_group, + self.dA_vloop_tiler[:2], # (64, 64) + # default a_source=OperandSource.SMEM → SS mode + ) + + # 3. dvb_tiled_mma: SS MN,MN (64,64) — dvb + dkgb + # dvb = A @ dv, dkgb = A @ dw + dvb_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.MN, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + self.cta_group, + self.dvb_tiler[:2], # (64, 64) + ) + + # dkgb_tiled_mma: SS MN,MN (64,128) - dkgb + dkgb_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.MN, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + self.cta_group, + self.kloop_dkgb_tiler[:2], # (64, 128) + ) + + # dA_kloop_tiled_mma: SS K,K (64, 64) + # dA += dw @ kg^T + dA_kloop_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + self.cta_group, + self.kloop_dA_tiler[:2], # (64, 64) + ) + + # dA2post_tiled_mma: SS K,K (64,64) + # dA = dA @ A + dA2post_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + self.cta_group, + self.dApost_tiler[:2], # (64, 64) + # tcgen05.OperandSource.SMEM, # SS mode + ) + + # dA3post_tiled_mma: SS MN,MN (64,64) + # dA = A @ dA + dA3post_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.MN, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + self.cta_group, + self.dApost_tiler[:2], # (64, 64) + # tcgen05.OperandSource.SMEM, # SS mode + ) + + # ===================== SMEM layouts ===================== + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + # SS opA layout: do/vnew/dv [BT,BV]=[64,64] K-major + vloop_opA_smem = sm100_utils.make_smem_layout_a( + vloop_tiled_mma, + self.vloop_gemm_tiler, + self.io_dtype, + self.vloop_stage, + ) + + # SS opB layout: h/dh [BK,BV]=[128,64] K-major + vloop_opB_smem = sm100_utils.make_smem_layout_b( + vloop_tiled_mma, + self.vloop_gemm_tiler, + self.io_dtype, + self.vloop_stage, + ) + + # SS opB layout: v [BV,BT]=[128,64] K-major (dA vloop) + v_opB_smem = sm100_utils.make_smem_layout_b( + dA_vloop_tiled_mma, + self.dA_vloop_tiler, + self.io_dtype, + self.vloop_stage, + ) + + # SS opA layout: A MN-major [BT,BT]=[64,64] + A_mn_opA_smem = sm100_utils.make_smem_layout_a( + dvb_tiled_mma, + self.dvb_tiler, + self.io_dtype, + self.a_stage, + ) + + # opB: dv MN-major [BV,BT]=[64,64] + dv_mn_opB_smem = sm100_utils.make_smem_layout_b( + dvb_tiled_mma, + self.dvb_tiler, + self.io_dtype, + self.vloop_stage, + ) + + # opA: dw K-major [BT,BK]=[64,128] + dw_k_opA_smem = sm100_utils.make_smem_layout_a( + dA_vloop_tiled_mma, + self.kloop_dA_tiler, + self.io_dtype, + self.kloop_stage, + ) + + # opB: dw MN-major [BK,BT] + dw_mn_opB_smem = sm100_utils.make_smem_layout_b( + dkgb_tiled_mma, + self.kloop_dkgb_tiler, + self.io_dtype, + self.kloop_stage, + ) + + # opB: kg^T K-major [BT, BK] + kg_k_opB_smem = sm100_utils.make_smem_layout_b( + dA_kloop_tiled_mma, + self.kloop_dA_tiler, + self.io_dtype, + self.kloop_stage, + ) + + # opA: dA K-major [BT,BT] + dA_k_opA_smem = sm100_utils.make_smem_layout_a( + dA2post_tiled_mma, + self.dApost_tiler, + self.io_dtype, + self.mma_stage, + ) + + # opB: A K-major [BT,BT] + A_k_opB_smem = sm100_utils.make_smem_layout_b( + dA2post_tiled_mma, + self.dApost_tiler, + self.io_dtype, + self.a_stage, + ) + + # opB: dA MN-major [BT,BT] + dA_mn_opB_smem = sm100_utils.make_smem_layout_b( + dA3post_tiled_mma, + self.dApost_tiler, + self.io_dtype, + self.mma_stage, + ) + + # --- Epilogue (non-MMA) layouts --- + g_epi_smem_layout = sm100_utils.make_smem_layout_epi( + self.g_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BK), + self.kloop_stage, + ) + + k_epi_smem_layout = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BK), + self.kloop_stage, + ) + + q_epi_smem_layout = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BK), + 1, + ) + + dg_epi_smem_layout = sm100_utils.make_smem_layout_epi( + self.g_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BK), + self.kloop_stage, + ) + + # ===================== Cluster layout ===================== + cluster_layout = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (vloop_tiled_mma.thr_id.shape,), + ) + + # ===================== TMA descriptors ===================== + # Strip stage dimension for TMA atom creation (expects 3 modes, not 4) + vloop_opA_smem_no_stage = cute.select(vloop_opA_smem, mode=[0, 1, 2]) + vloop_opB_smem_no_stage = cute.select(vloop_opB_smem, mode=[0, 1, 2]) + v_opB_smem_no_stage = cute.select(v_opB_smem, mode=[0, 1, 2]) + A_mn_opA_smem_no_stage = cute.select(A_mn_opA_smem, mode=[0, 1, 2]) + + tma_atom_dv, tma_tensor_dv = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + dv, + vloop_opA_smem_no_stage, + self.vloop_gemm_tiler, + vloop_tiled_mma, + cluster_layout.shape, + ) + + tma_atom_A, tma_tensor_A = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + A_T, + A_mn_opA_smem_no_stage, + self.dvb_tiler, + dvb_tiled_mma, + cluster_layout.shape, + ) + + tma_atom_h, tma_tensor_h = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + h, + vloop_opB_smem_no_stage, + self.vloop_gemm_tiler, + vloop_tiled_mma, + cluster_layout.shape, + ) + + tma_atom_dh, tma_tensor_dh = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + dh, + vloop_opB_smem_no_stage, + self.vloop_gemm_tiler, + vloop_tiled_mma, + cluster_layout.shape, + ) + + tma_atom_do, tma_tensor_do = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + do, + vloop_opA_smem_no_stage, + self.vloop_gemm_tiler, + vloop_tiled_mma, + cluster_layout.shape, + ) + + tma_atom_vnew, tma_tensor_vnew = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + v_new, + vloop_opA_smem_no_stage, + self.vloop_gemm_tiler, + vloop_tiled_mma, + cluster_layout.shape, + ) + + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + v, + v_opB_smem_no_stage, + self.dA_vloop_tiler, + dA_vloop_tiled_mma, + cluster_layout.shape, + ) + + g_epi_smem_no_stage = cute.select(g_epi_smem_layout, mode=[0, 1]) + tma_atom_g, tma_tensor_g = cpasync.make_tiled_tma_atom( + tma_load_op, + g, + g_epi_smem_no_stage, + (self.BT, self.BK), + ) + + k_epi_smem_no_stage = cute.select(k_epi_smem_layout, mode=[0, 1]) + tma_atom_k, tma_tensor_k = cpasync.make_tiled_tma_atom( + tma_load_op, + k, + k_epi_smem_no_stage, + (self.BT, self.BK), + ) + + q_epi_smem_no_stage = cute.select(q_epi_smem_layout, mode=[0, 1]) + tma_atom_q, tma_tensor_q = cpasync.make_tiled_tma_atom( + tma_load_op, + q, + q_epi_smem_no_stage, + (self.BT, self.BK), + ) + + dg_epi_smem_no_stage = cute.select(dg_epi_smem_layout, mode=[0, 1]) + tma_atom_dg, tma_tensor_dg = cpasync.make_tiled_tma_atom( + tma_store_op, + dg, + dg_epi_smem_no_stage, + (self.BT, self.BK), + ) + + # ===================== TMA byte counts ===================== + self.tma_bytes_A = cute.size_in_bytes(self.io_dtype, A_mn_opA_smem_no_stage) + self.tma_bytes_dv = cute.size_in_bytes(self.io_dtype, vloop_opA_smem_no_stage) + self.tma_bytes_h = cute.size_in_bytes(self.io_dtype, vloop_opB_smem_no_stage) + self.tma_bytes_dh = cute.size_in_bytes(self.io_dtype, vloop_opB_smem_no_stage) + self.tma_bytes_do = cute.size_in_bytes(self.io_dtype, vloop_opA_smem_no_stage) + self.tma_bytes_vnew = cute.size_in_bytes(self.io_dtype, vloop_opA_smem_no_stage) + self.tma_bytes_g = cute.size_in_bytes(self.g_dtype, g_epi_smem_no_stage) + self.tma_bytes_v = cute.size_in_bytes(self.io_dtype, v_opB_smem_no_stage) + self.tma_bytes_k = cute.size_in_bytes(self.io_dtype, k_epi_smem_no_stage) + self.tma_bytes_q = cute.size_in_bytes(self.io_dtype, q_epi_smem_no_stage) + + # ===================== SharedStorage ===================== + @cute.struct + class SharedStorage: + # ======= mbarrier ======= + bar_load_A: cute.struct.MemRange[Int64, self.a_stage * 2] + bar_load_dv: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_mma_dvb: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_load_beta: cute.struct.MemRange[Int64, 1 * 2] + bar_tma_h: cute.struct.MemRange[Int64, self.vloop_stage] + bar_mma_cuda_h: cute.struct.MemRange[Int64, self.vloop_stage] + bar_tma_dh: cute.struct.MemRange[Int64, self.vloop_stage] + bar_mma_cuda_dh: cute.struct.MemRange[Int64, self.vloop_stage] + bar_tma_v: cute.struct.MemRange[Int64, self.vloop_stage] + bar_mma_cuda_v: cute.struct.MemRange[Int64, self.vloop_stage] + bar_load_do: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_load_g: cute.struct.MemRange[Int64, self.kloop_stage * 2] + bar_load_vnew: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_load_q: cute.struct.MemRange[Int64, self.kloop_stage * 2] + bar_load_k: cute.struct.MemRange[Int64, self.kloop_stage * 2] + bar_mma_dq: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_dw: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_dk: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_dkgb: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_dA: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_dA2: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_dA3: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_mma_done_vloop: cute.struct.MemRange[Int64, self.mma_stage] + bar_prologue_dw: cute.struct.MemRange[Int64, self.kloop_stage * 2] + bar_prologue_kg: cute.struct.MemRange[Int64, self.kloop_stage * 2] + bar_prologue_dA2: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_prologue_dA3: cute.struct.MemRange[Int64, self.mma_stage * 2] + bar_store_dg: cute.struct.MemRange[Int64, self.kloop_stage * 2] + # TMEM holding buffer + tmem_holding_buf: Int32 + # A, stage=2, [BT,BT], 16KB + buf_A: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(A_mn_opA_smem)], + self.buffer_align_bytes, + ] + # k, stage=1, [BT,BK], 16KB + buf_k: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(k_epi_smem_layout)], + self.buffer_align_bytes, + ] + # g, stage=1, [BT,BK], 32KB + buf_g: cute.struct.Align[ + cute.struct.MemRange[self.g_dtype, cute.cosize(g_epi_smem_layout)], + self.buffer_align_bytes, + ] + # q, stage=1, [BT,BK], 16KB + buf_q: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(q_epi_smem_layout)], + self.buffer_align_bytes, + ] + # V-loop buffers, stage=2 + # h, dh, [BK,BV] 32KB*2 + buf_h: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(vloop_opB_smem)], + self.buffer_align_bytes, + ] + buf_dh: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(vloop_opB_smem)], + self.buffer_align_bytes, + ] + # do, dv, v_new, v, [BT,BV] 16KB*4 + buf_do: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(vloop_opA_smem)], + self.buffer_align_bytes, + ] + buf_dv: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(vloop_opA_smem)], + self.buffer_align_bytes, + ] + buf_vnew: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(vloop_opA_smem)], + self.buffer_align_bytes, + ] + buf_v: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(v_opB_smem)], + self.buffer_align_bytes, + ] + + # dw, stage=1, [BT,BK] 16KB + buf_dw: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(dw_k_opA_smem)], + self.buffer_align_bytes, + ] + # Scalars + s_beta: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BT], + 128, + ] + # 2 slots per row, one per warpgroup, for deterministic db reduction + # (avoids cross-wg fp32 atomicAdd on shared memory). + s_db: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BT * 2], + 128, + ] + s_gn: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BK], + 128, + ] + s_dgk: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BK], + 128, + ] + + self.shared_storage = SharedStorage + + # ===================== cu_seqlens / chunk_indices tensors ===================== + cu_seqlens = cute.make_tensor(cu_seqlens_ptr, cute.make_layout((B + 1,))) + chunk_indices = cute.make_tensor(chunk_indices_ptr, cute.make_layout((total_nt, 2), stride=(2, 1))) + + # ===================== Grid ===================== + grid = self._compute_grid(B, T, HV, total_nt=total_nt) + + # ===================== Launch kernel ===================== + self.kernel( + # MMA objects (4) + vloop_tiled_mma, + dA_vloop_tiled_mma, + dvb_tiled_mma, + dA_kloop_tiled_mma, + dA2post_tiled_mma, + dA3post_tiled_mma, + # TMA atoms + tma_atom_dv, + tma_tensor_dv, + tma_atom_A, + tma_tensor_A, + tma_atom_h, + tma_tensor_h, + tma_atom_dh, + tma_tensor_dh, + tma_atom_do, + tma_tensor_do, + tma_atom_g, + tma_tensor_g, + tma_atom_v, + tma_tensor_v, + tma_atom_k, + tma_tensor_k, + tma_atom_vnew, + tma_tensor_vnew, + tma_atom_q, + tma_tensor_q, + tma_atom_dg, + tma_tensor_dg, + # SMEM layouts + vloop_opA_smem, + vloop_opB_smem, + v_opB_smem, + A_mn_opA_smem, + dv_mn_opB_smem, + dw_k_opA_smem, + dw_mn_opB_smem, + kg_k_opB_smem, + A_k_opB_smem, + dA_k_opA_smem, + dA_mn_opB_smem, + g_epi_smem_layout, + k_epi_smem_layout, + q_epi_smem_layout, + # GMEM tensors + q, + k, + g, + beta, + dq, + dk, + dv2, + dg, + db, + dA_out, + # Metadata + cu_seqlens, + chunk_indices, + problem_size, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=self.min_occupancy, + ) + + @cute.kernel + def kernel( + self, + # MMA objects (4) + vloop_tiled_mma: cute.TiledMma, + dA_vloop_tiled_mma: cute.TiledMma, + dvb_tiled_mma: cute.TiledMma, + dA_kloop_tiled_mma: cute.TiledMma, + dA2post_tiled_mma: cute.TiledMma, + dA3post_tiled_mma: cute.TiledMma, + # TMA atoms + tensors + tma_atom_dv: cute.CopyAtom, + tma_tensor_dv: cute.Tensor, + tma_atom_A: cute.CopyAtom, + tma_tensor_A: cute.Tensor, + tma_atom_h: cute.CopyAtom, + tma_tensor_h: cute.Tensor, + tma_atom_dh: cute.CopyAtom, + tma_tensor_dh: cute.Tensor, + tma_atom_do: cute.CopyAtom, + tma_tensor_do: cute.Tensor, + tma_atom_g: cute.CopyAtom, + tma_tensor_g: cute.Tensor, + tma_atom_v: cute.CopyAtom, + tma_tensor_v: cute.Tensor, + tma_atom_k: cute.CopyAtom, + tma_tensor_k: cute.Tensor, + tma_atom_vnew: cute.CopyAtom, + tma_tensor_vnew: cute.Tensor, + tma_atom_q: cute.CopyAtom, + tma_tensor_q: cute.Tensor, + tma_atom_dg: cute.CopyAtom, + tma_tensor_dg: cute.Tensor, + # SMEM layouts + vloop_opA_smem: cute.ComposedLayout, + vloop_opB_smem: cute.ComposedLayout, + v_opB_smem: cute.ComposedLayout, + A_mn_opA_smem: cute.ComposedLayout, + dv_mn_opB_smem: cute.ComposedLayout, + dw_k_opA_smem: cute.ComposedLayout, + dw_mn_opB_smem: cute.ComposedLayout, + kg_k_opB_smem: cute.ComposedLayout, + A_k_opB_smem: cute.ComposedLayout, + dA_k_opA_smem: cute.ComposedLayout, + dA_mn_opB_smem: cute.ComposedLayout, + g_epi_smem_layout: cute.ComposedLayout, + k_epi_smem_layout: cute.ComposedLayout, + q_epi_smem_layout: cute.ComposedLayout, + # GMEM tensors + q_gmem: cute.Tensor, + k_gmem: cute.Tensor, + g_gmem: cute.Tensor, + beta_gmem: cute.Tensor, + dq_gmem: cute.Tensor, + dk_gmem: cute.Tensor, + dv2_gmem: cute.Tensor, + dg_gmem: cute.Tensor, + db_gmem: cute.Tensor, + dA_gmem: cute.Tensor, + # Metadata + cu_seqlens: cute.Tensor, + chunk_indices: cute.Tensor, + problem_size: tuple[Int32, Int32, Int32, Int32, Int32, Int32], # (B, T, H, HV, K, V) + ): + B, T, H, HV, K, V = problem_size + BT = self.BT + + # ===================== Persistent work decode ===================== + # Grid: (min(num_sm * occ, total_tiles), 1, 1) — persistent + block_idx_x = cute.arch.block_idx()[0] + grid_dim_x = cute.arch.grid_dim()[0] + thread_idx = cute.arch.thread_idx()[0] + lane_idx = thread_idx % 32 + + total_work_units = chunk_indices.layout.shape[0] * HV + num_iters = (total_work_units - block_idx_x + grid_dim_x - 1) // grid_dim_x + + num_cuda_warps_total = len(self.cuda_warp_ids) + len(self.cuda2_warp_ids) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + + if warp_idx == self.load_warp_id: + cpasync.prefetch_descriptor(tma_atom_A) + cpasync.prefetch_descriptor(tma_atom_dv) + cpasync.prefetch_descriptor(tma_atom_h) + cpasync.prefetch_descriptor(tma_atom_dh) + cpasync.prefetch_descriptor(tma_atom_do) + cpasync.prefetch_descriptor(tma_atom_g) + cpasync.prefetch_descriptor(tma_atom_v) + cpasync.prefetch_descriptor(tma_atom_vnew) + cpasync.prefetch_descriptor(tma_atom_k) + cpasync.prefetch_descriptor(tma_atom_q) + cpasync.prefetch_descriptor(tma_atom_dg) + + # ===================== SMEM allocation ===================== + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Barrier Initialization + bar_mma_done_vloop_ptr = storage.bar_mma_done_vloop.data_ptr() + # NOTE: for h, dh and v, consumer contains both MMA and CUDA Core, so we use original mbarrier declaration instead of pipeline utils + bar_tma_h_ptr = storage.bar_tma_h.data_ptr() + bar_mma_cuda_h_ptr = storage.bar_mma_cuda_h.data_ptr() + bar_tma_dh_ptr = storage.bar_tma_dh.data_ptr() + bar_mma_cuda_dh_ptr = storage.bar_mma_cuda_dh.data_ptr() + bar_tma_v_ptr = storage.bar_tma_v.data_ptr() + bar_mma_cuda_v_ptr = storage.bar_mma_cuda_v.data_ptr() + if warp_idx == 0: + with elect_one(): + for i in cutlass.range(self.mma_stage, unroll_full=True): + mbarrier_init(bar_mma_done_vloop_ptr + i, 1) + for i in cutlass.range(self.vloop_stage, unroll_full=True): + mbarrier_init(bar_tma_h_ptr + i, 1) + mbarrier_init(bar_mma_cuda_h_ptr + i, num_cuda_warps_total * 32 + 1) + mbarrier_init(bar_tma_dh_ptr + i, 1) + mbarrier_init(bar_mma_cuda_dh_ptr + i, num_cuda_warps_total * 32 + 1) + mbarrier_init(bar_tma_v_ptr + i, 1) + mbarrier_init(bar_mma_cuda_v_ptr + i, num_cuda_warps_total * 32 + 1) + mbarrier_init_fence() + + # ====== Pipeline Definition ====== + pipeline_load_A = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.bar_load_A.data_ptr(), + num_stages=self.a_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_bytes_A, + ) + pipeline_load_dv = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.bar_load_dv.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_bytes_dv, + ) + pipeline_load_do = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.bar_load_do.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_bytes_do, + ) + pipeline_load_vnew = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.bar_load_vnew.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_bytes_vnew, + ) + pipeline_load_g = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_g.data_ptr(), + num_stages=self.kloop_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total + len(self.aux_warp_ids)), + tx_count=self.tma_bytes_g, + ) + pipeline_load_k = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_k.data_ptr(), + num_stages=self.kloop_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total), + tx_count=self.tma_bytes_k, + ) + pipeline_load_q = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_q.data_ptr(), + num_stages=self.kloop_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total), + tx_count=self.tma_bytes_q, + ) + pipeline_mma_dvb = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dvb.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_mma_dq = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dq.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_mma_dk = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dk.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_mma_dw = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dw.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_mma_dA = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dA.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_mma_dA2 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dA2.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_mma_dA3 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dA3.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_prologue_dw = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.bar_prologue_dw.data_ptr(), + num_stages=self.kloop_stage, + producer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + ) + pipeline_prologue_kg = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.bar_prologue_kg.data_ptr(), + num_stages=self.kloop_stage, + producer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + ) + pipeline_prologue_dA2 = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.bar_prologue_dA2.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + ) + pipeline_prologue_dA3 = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.bar_prologue_dA3.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + ) + pipeline_mma_dkgb = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.bar_mma_dkgb.data_ptr(), + num_stages=self.mma_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_load_beta = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_load_beta.data_ptr(), + num_stages=1, + producer_group=make_thread_cooperative_group(len(self.aux_warp_ids) * 32), + consumer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + ) + pipeline_store_dg = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_store_dg.data_ptr(), + num_stages=self.kloop_stage, + producer_group=make_thread_cooperative_group(num_cuda_warps_total * 32), + consumer_group=make_thread_cooperative_group(len(self.aux_warp_ids) * 32), + ) + + # ===================== TMEM allocation ===================== + tmem_alloc_bar = pipeline.NamedBarrier(barrier_id=1, num_threads=self.threads_per_cta) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_bar, + allocator_warp_id=self.load_warp_id, + ) + # Cluster arrive after barrier init + pipeline.pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) + + vloop_opA_smem_no_stage = cute.select(vloop_opA_smem, mode=[0, 1, 2]) + vloop_opB_smem_no_stage = cute.select(vloop_opB_smem, mode=[0, 1, 2]) + A_mn_opA_smem_no_stage = cute.select(A_mn_opA_smem, mode=[0, 1, 2]) + v_opB_smem_no_stage = cute.select(v_opB_smem, mode=[0, 1, 2]) + + sA = storage.buf_A.get_tensor(A_mn_opA_smem.outer, swizzle=A_mn_opA_smem.inner) + sDv = storage.buf_dv.get_tensor(vloop_opA_smem.outer, swizzle=vloop_opA_smem.inner) + sH = storage.buf_h.get_tensor(vloop_opB_smem.outer, swizzle=vloop_opB_smem.inner) + sDh = storage.buf_dh.get_tensor(vloop_opB_smem.outer, swizzle=vloop_opB_smem.inner) + sDo = storage.buf_do.get_tensor(vloop_opA_smem.outer, swizzle=vloop_opA_smem.inner) + sVnew = storage.buf_vnew.get_tensor(vloop_opA_smem.outer, swizzle=vloop_opA_smem.inner) + sV = storage.buf_v.get_tensor(v_opB_smem.outer, swizzle=v_opB_smem.inner) + + sDv_ptr_base = storage.buf_dv.data_ptr().toint() + vloop_opA_bytes_per_stage = cute.size_in_bytes(self.io_dtype, vloop_opA_smem_no_stage) + sDo_ptr_base = storage.buf_do.data_ptr().toint() + sVnew_ptr_base = storage.buf_vnew.data_ptr().toint() + sV_ptr_base = storage.buf_v.data_ptr().toint() + v_opB_bytes_per_stage = cute.size_in_bytes(self.io_dtype, v_opB_smem_no_stage) + sH_ptr_base = storage.buf_h.data_ptr().toint() + sDh_ptr_base = storage.buf_dh.data_ptr().toint() + vloop_opB_bytes_per_stage = cute.size_in_bytes(self.io_dtype, vloop_opB_smem_no_stage) + sA_ptr_base = storage.buf_A.data_ptr().toint() + A_bytes_per_stage = cute.size_in_bytes(self.io_dtype, A_mn_opA_smem_no_stage) + + # NOTE: make_umma_smem_desc requires the iterator to carry the swizzle + # (and ≥16B alignment). When constructing a tensor over a ComposedLayout + # via make_ptr+make_tensor, the swizzle ends up composed on the layout + # rather than the iterator, which breaks make_umma_smem_desc. Use + # recast_ptr to move the swizzle onto the iterator and pair it with the + # underlying (non-swizzle) outer layout. + sDv_mn = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_dv.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=dv_mn_opB_smem.inner, + dtype=self.io_dtype, + ), + dv_mn_opB_smem.outer, + ) + sDw_mn = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_dw.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=dw_mn_opB_smem.inner, + dtype=self.io_dtype, + ), + dw_mn_opB_smem.outer, + ) + sDw_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_dw.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=dw_k_opA_smem.inner, + dtype=self.io_dtype, + ), + dw_k_opA_smem.outer, + ) + sDv_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_dv.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=vloop_opA_smem.inner, + dtype=self.io_dtype, + ), + vloop_opA_smem.outer, + ) + sV_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_v.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=v_opB_smem.inner, + dtype=self.io_dtype, + ), + v_opB_smem.outer, + ) + sA_mn = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_A.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=A_mn_opA_smem.inner, + dtype=self.io_dtype, + ), + A_mn_opA_smem.outer, + ) + sDo_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr(self.io_dtype, storage.buf_do.data_ptr().toint(), cute.AddressSpace.smem, assumed_align=128), + swizzle_=vloop_opA_smem.inner, + dtype=self.io_dtype, + ), + vloop_opA_smem.outer, + ) + sVnew_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr(self.io_dtype, storage.buf_vnew.data_ptr().toint(), cute.AddressSpace.smem, assumed_align=128), + swizzle_=vloop_opA_smem.inner, + dtype=self.io_dtype, + ), + vloop_opA_smem.outer, + ) + sH_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr(self.io_dtype, storage.buf_h.data_ptr().toint(), cute.AddressSpace.smem, assumed_align=128), + swizzle_=vloop_opB_smem.inner, + dtype=self.io_dtype, + ), + vloop_opB_smem.outer, + ) + sDh_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr(self.io_dtype, storage.buf_dh.data_ptr().toint(), cute.AddressSpace.smem, assumed_align=128), + swizzle_=vloop_opB_smem.inner, + dtype=self.io_dtype, + ), + vloop_opB_smem.outer, + ) + sKG_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr(self.io_dtype, storage.buf_k.data_ptr().toint(), cute.AddressSpace.smem, assumed_align=128), + swizzle_=kg_k_opB_smem.inner, + dtype=self.io_dtype, + ), + kg_k_opB_smem.outer, + ) + sA_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_A.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=A_k_opB_smem.inner, + dtype=self.io_dtype, + ), + A_k_opB_smem.outer, + ) + sDA_mn = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr(self.io_dtype, storage.buf_q.data_ptr().toint(), cute.AddressSpace.smem, assumed_align=128), + swizzle_=dA_mn_opB_smem.inner, + dtype=self.io_dtype, + ), + dA_mn_opB_smem.outer, + ) + sDA_k = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_q.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=dA_k_opA_smem.inner, + dtype=self.io_dtype, + ), + dA_k_opA_smem.outer, + ) + sG_raw = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.g_dtype, + storage.buf_g.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=g_epi_smem_layout.inner, + dtype=self.g_dtype, + ), + g_epi_smem_layout.outer, + ) + sG_raw_ptr = cute.make_ptr(self.g_dtype, storage.buf_g.data_ptr().toint(), cute.AddressSpace.smem) + sK_raw = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_k.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=k_epi_smem_layout.inner, + dtype=self.io_dtype, + ), + k_epi_smem_layout.outer, + ) + sK_raw_ptr = cute.make_ptr(self.io_dtype, storage.buf_k.data_ptr().toint(), cute.AddressSpace.smem) + sDw_raw_ptr = cute.make_ptr(self.io_dtype, storage.buf_dw.data_ptr().toint(), cute.AddressSpace.smem) + sQ_raw = cute.make_tensor( + cute.recast_ptr( + cute.make_ptr( + self.io_dtype, + storage.buf_q.data_ptr().toint(), + cute.AddressSpace.smem, + assumed_align=128, + ), + swizzle_=q_epi_smem_layout.inner, + dtype=self.io_dtype, + ), + q_epi_smem_layout.outer, + ) + sQ_raw_ptr = cute.make_ptr(self.io_dtype, storage.buf_q.data_ptr().toint(), cute.AddressSpace.smem) + + # Scalar SMEM buffers (plain layouts, no swizzle) + sBeta = cute.make_tensor( + cute.make_ptr(Float32, storage.s_beta.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((self.BT,), stride=(1,)), + ) + # sDb layout: (BT, 2). Inner dim = wg_idx slot. Stride (1, BT) so each + # wg's column is contiguous (better for the reduce in Phase 3). + sDb = cute.make_tensor( + cute.make_ptr(Float32, storage.s_db.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((self.BT, 2), stride=(1, self.BT)), + ) + sDgk = cute.make_tensor( + cute.make_ptr(Float32, storage.s_dgk.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((self.BK,), stride=(1,)), + ) + sGn = cute.make_tensor( + cute.make_ptr(Float32, storage.s_gn.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((self.BK,), stride=(1,)), + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline.pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + tmem.allocate(TMEM_TOTAL) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + # ===================== Warp dispatch ===================== + # CUDA Core loop body + if warp_idx in self.cuda_warp_ids or warp_idx in self.cuda2_warp_ids: + cute.arch.setmaxregister_increase(self.num_regs_cuda) + + load_beta_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + load_g_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + mma_dvb_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + mma_dq_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + mma_dw_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + mma_dk_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + load_k_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + prologue_dw_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kloop_stage) + prologue_kg_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kloop_stage) + mma_dgkb_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + load_q_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + mma_dA_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + mma_dA2_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + mma_dA3_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + prologue_dA2_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + prologue_dA3_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + store_dg_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kloop_stage) + + wg_idx = tidx // 128 + local_tidx = tidx % 128 + warp_id = local_tidx // 32 + warp_row_tile = warp_id % 2 + warp_col_tile = warp_id // 2 + row = warp_row_tile * 32 + lane_idx # BT1 + bk_num_cols = self.BK // 2 + bv_num_cols = self.BV // 2 + bk_num_cols_per_wg = bk_num_cols // 2 + bv_num_cols_per_wg = bv_num_cols // 2 + bt_num_cols_per_wg = self.BT // 4 + # ref: https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-e + bv_col_base = warp_col_tile * (self.BV // 2) + wg_idx * bv_num_cols_per_wg + bk_col_base = warp_col_tile * (self.BK // 2) + wg_idx * bk_num_cols_per_wg + bt_col_base = warp_col_tile * (self.BT // 2) + wg_idx * bt_num_cols_per_wg + # 8 fp32 store each time for store_256b + num_stores_f32 = bk_num_cols_per_wg // 8 + + vloop_stage_idx = 0 + vloop_phase = 0 + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + G = HV // H + i_t = work_idx // HV # chunk index (global) + i_hv = work_idx % HV # value-head index + i_h = i_hv // G # q/k head index + # Decode chunk_indices + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + tok_offset = cu_seqlens[(batch_idx,)] + seq_len = cu_seqlens[(batch_idx + 1,)] - tok_offset + sub_seq_len = min(self.BT, seq_len - tile_idx * self.BT) + + # NOTE: must sync before next wu_iter's `sDgk[local_tidx] = 0` + # init, otherwise WG0 of next iter may overwrite sDgk while + # WG1 of this iter (row == sub_seq_len - 1 lane) is still + # reading sDgk[col] above. This was the source of the + # non-deterministic dg accuracy bug. + self.cuda_wg_sync_barrier.arrive_and_wait() + # fill db, dgk to 0. Each wg zeroes its own sDb column. + if local_tidx < self.BT: + sDb[local_tidx, 0] = Float32(0.0) + sDb[local_tidx, 1] = Float32(0.0) + if local_tidx < self.BK: + sDgk[local_tidx] = Float32(0.0) + self.cuda_wg_sync_barrier.arrive_and_wait() + + pipeline_load_beta.consumer_wait(load_beta_consumer_state) + cute.arch.fence_proxy("async.shared", space="cta") + + beta_val = sBeta[(row,)] + db_val = Float32(0.0) + for v_iter in cutlass.range(self.num_v_tiles): + # dgk += sum(h * dh, axis=0) + mbarrier_wait(bar_tma_h_ptr + vloop_stage_idx, vloop_phase) + mbarrier_wait(bar_tma_dh_ptr + vloop_stage_idx, vloop_phase) + + sH_raw_ptr = cute.make_ptr( + self.io_dtype, sH_ptr_base + vloop_stage_idx * vloop_opB_bytes_per_stage, cute.AddressSpace.smem + ) + sDh_raw_ptr = cute.make_ptr( + self.io_dtype, sDh_ptr_base + vloop_stage_idx * vloop_opB_bytes_per_stage, cute.AddressSpace.smem + ) + # each thread in one WG processes one row + self.cuda_wg_sync_barrier.arrive_and_wait() + if wg_idx == 0: + for i in cutlass.range_constexpr(self.BV // 8): + col = i * 8 + h_vals = smem_load_bf16x8_sw128(sH_raw_ptr, local_tidx, col) + dh_vals = smem_load_bf16x8_sw128(sDh_raw_ptr, local_tidx, col) + h_dh_vals = cute.make_rmem_tensor((8,), Float32) + h_dh_vals.store(h_vals.load().to(Float32) * dh_vals.load().to(Float32)) + for j in cutlass.range_constexpr(8): + sDgk[(local_tidx,)] += h_dh_vals[j] + + mbarrier_arrive(bar_mma_cuda_h_ptr + vloop_stage_idx) + mbarrier_arrive(bar_mma_cuda_dh_ptr + vloop_stage_idx) + + pipeline_mma_dvb.consumer_wait(mma_dvb_consumer_state) + tcgen05_fence_after() + dvb_i32 = tcgen05_ld_32x32b(bv_num_cols_per_wg, TMEM_FLEX_OFF + wg_idx * bv_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_mma_dvb.consumer_release(mma_dvb_consumer_state) + mma_dvb_consumer_state.advance() + + dvb_f32 = reinterpret_cast(dvb_i32, Int32, bv_num_cols_per_wg, Float32) + dvb_f32_val = TensorSSA(dvb_f32, (bv_num_cols_per_wg,), Float32) + + # db += sum(dvb * v, axis=1) + mbarrier_wait(bar_tma_v_ptr + vloop_stage_idx, vloop_phase) + rV_bf16 = cute.make_rmem_tensor((bv_num_cols_per_wg,), self.io_dtype) + sV_raw_ptr_cur = cute.make_ptr( + self.io_dtype, sV_ptr_base + vloop_stage_idx * v_opB_bytes_per_stage, cute.AddressSpace.smem + ) + if row < sub_seq_len: + for i in cutlass.range_constexpr(bv_num_cols_per_wg // 8): + col_base = bv_col_base + i * 8 + vals = smem_load_bf16x8_sw128(sV_raw_ptr_cur, row, col_base) + rV_bf16[i * 8 + 0] = vals[0] + rV_bf16[i * 8 + 1] = vals[1] + rV_bf16[i * 8 + 2] = vals[2] + rV_bf16[i * 8 + 3] = vals[3] + rV_bf16[i * 8 + 4] = vals[4] + rV_bf16[i * 8 + 5] = vals[5] + rV_bf16[i * 8 + 6] = vals[6] + rV_bf16[i * 8 + 7] = vals[7] + else: + rV_bf16.fill(BFloat16(0.0)) + rV_fp32 = cute.make_rmem_tensor((bv_num_cols_per_wg,), Float32) + rV_fp32.store(rV_bf16.load().to(Float32)) + rV_fp32.store(rV_fp32.load() * dvb_f32_val) + if row < sub_seq_len: + for i in cutlass.range_constexpr(bv_num_cols_per_wg): + db_val += rV_fp32[i] + + mbarrier_arrive(bar_mma_cuda_v_ptr + vloop_stage_idx) + + # ── dv2 epilogue: dv2 = dvb * beta, cast to bf16, store to gmem ── + dvb_f32_rmem = cute.make_rmem_tensor((bv_num_cols_per_wg,), Float32) + dvb_f32_rmem.store(dvb_f32_val * beta_val) + + dvb_bf16_rmem = cute.make_rmem_tensor((bv_num_cols_per_wg,), self.io_dtype) + dvb_bf16_rmem.store(dvb_f32_rmem.load().to(self.io_dtype)) + + # bf16 vector → i32 vector for store_256b (8 i32 = 16 bf16 = 32 bytes per store). + dvb_bf16_val = dvb_bf16_rmem.load() + dvb_i32_vec = reinterpret_cast(dvb_bf16_val, self.io_dtype, bv_num_cols_per_wg, Int32) + # bv_num_cols bf16 = bv_num_cols // 16 stores of 256b each. + num_stores_per_row = bv_num_cols_per_wg // 16 # = 4 for BV=128 + + base_addr = ( + dv2_gmem.iterator + + (tok_offset + tile_idx * self.BT + row) * HV * V + + i_hv * V + + v_iter * self.BV + + bv_col_base + ).toint() + if row < sub_seq_len: + for s in cutlass.range_constexpr(num_stores_per_row): + chunk = subvec(dvb_i32_vec, s * 8, 8) + store_256b(base_addr + s * 32, chunk) + + vloop_stage_idx = (vloop_stage_idx + 1) % self.vloop_stage + vloop_phase ^= 1 + + # gk_exp = exp2(g) + pipeline_load_g.consumer_wait(load_g_consumer_state) + # write to gn + sGn[local_tidx] = sG_raw[(sub_seq_len - 1, local_tidx, 0)] + + # row-major load, match TMEM layout + rG = cute.make_rmem_tensor((self.BK // 4,), self.g_dtype) + if row < sub_seq_len: + for i in cutlass.range_constexpr(self.BK // 4 // 4): + col_base = bk_col_base + i * 4 + vals = smem_load_f32x4_sw128(sG_raw_ptr, row, col_base) + rG[i * 4 + 0] = vals[0] + rG[i * 4 + 1] = vals[1] + rG[i * 4 + 2] = vals[2] + rG[i * 4 + 3] = vals[3] + else: + rG.fill(Float32(0.0)) + rG_val = rG.load() + rG_exp_val = cute.exp2(rG_val, fastmath=self.use_fast_math) + + # wait for dq, dq=dq*gk_exp*scale, GMEM store + pipeline_mma_dq.consumer_wait(mma_dq_consumer_state) + tcgen05_fence_after() + dq_i32 = tcgen05_ld_32x32b(bk_num_cols_per_wg, TMEM_DQ_ACC_OFF + wg_idx * bk_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_mma_dq.consumer_release(mma_dq_consumer_state) + mma_dq_consumer_state.advance() + + dq_f32 = reinterpret_cast(dq_i32, Int32, bk_num_cols_per_wg, Float32) + dq_f32_val = TensorSSA(dq_f32, (bk_num_cols_per_wg,), Float32) + + rDq = cute.make_rmem_tensor((bk_num_cols_per_wg,), Float32) + rDq.store(dq_f32_val * rG_exp_val * Float32(self.scale)) + + dq_f32_val_store = rDq.load() + dq_i32_vec = reinterpret_cast(dq_f32_val_store, Float32, bk_num_cols_per_wg, Int32) + # store to TMEM first to reduce register usage + tcgen05_st_32x32b(bk_num_cols_per_wg, TMEM_DQ_SCALED_OFF + wg_idx * bk_num_cols_per_wg, dq_i32_vec) + cute.arch.fence_view_async_tmem_store() + dq_base_addr = ( + dq_gmem.iterator + (tok_offset + tile_idx * self.BT + row) * HV * K + i_hv * K + bk_col_base + ).toint() + if row < sub_seq_len: + for s in cutlass.range_constexpr(num_stores_f32): + chunk = subvec(dq_i32_vec, s * 8, 8) + store_256b(dq_base_addr + s * 32, chunk) + + # wait for dw + pipeline_mma_dw.consumer_wait(mma_dw_consumer_state) + tcgen05_fence_after() + dw_i32 = tcgen05_ld_32x32b(bk_num_cols_per_wg, TMEM_DW_ACC_OFF + wg_idx * bk_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_mma_dw.consumer_release(mma_dw_consumer_state) + mma_dw_consumer_state.advance() + + # dw = -dw, convert to bf16, write to smem + dw_f32 = reinterpret_cast(dw_i32, Int32, bk_num_cols_per_wg, Float32) + dw_f32_val = TensorSSA(dw_f32, (bk_num_cols_per_wg,), Float32) + + dw_bf16_rmem = cute.make_rmem_tensor((bk_num_cols_per_wg,), BFloat16) + if row < sub_seq_len: + dw_bf16_rmem.store((-dw_f32_val).to(BFloat16)) + else: + dw_bf16_rmem.fill(BFloat16(0.0)) + + pipeline_prologue_dw.producer_acquire(prologue_dw_producer_state) + # store bf16x8 each time + dw_smem_num_stores = bk_num_cols_per_wg // 8 + for i in cutlass.range_constexpr(dw_smem_num_stores): + col_base = bk_col_base + i * 8 + chunk = cute.local_tile(dw_bf16_rmem, (8,), (i,)) + smem_store_bf16x8_sw128(sDw_raw_ptr, row, col_base, chunk) + + cute.arch.fence_proxy("async.shared", space="cta") + pipeline_prologue_dw.producer_commit(prologue_dw_producer_state) + prologue_dw_producer_state.advance() + + pipeline_load_k.consumer_wait(load_k_consumer_state) + # compute kg = k * gk_exp + rK = cute.make_rmem_tensor((self.BK // 4,), self.io_dtype) + if row < sub_seq_len: + for i in cutlass.range_constexpr(self.BK // 4 // 8): + col_base = bk_col_base + i * 8 + vals = smem_load_bf16x8_sw128(sK_raw_ptr, row, col_base) + rK[i * 8 + 0] = vals[0] + rK[i * 8 + 1] = vals[1] + rK[i * 8 + 2] = vals[2] + rK[i * 8 + 3] = vals[3] + rK[i * 8 + 4] = vals[4] + rK[i * 8 + 5] = vals[5] + rK[i * 8 + 6] = vals[6] + rK[i * 8 + 7] = vals[7] + else: + rK.fill(BFloat16(0.0)) + rK_fp32 = cute.make_rmem_tensor((self.BK // 4,), Float32) + rK_fp32.store(rK.load().to(Float32)) + rK_fp32_val = rK_fp32.load() + rKG_val = rK_fp32_val * rG_exp_val + + # write kg to K smem, + # notify dA += dw @ kg^T + rKG_bf16 = cute.make_rmem_tensor((self.BK // 4,), BFloat16) + rKG_bf16.store(rKG_val.to(BFloat16)) + + pipeline_prologue_kg.producer_acquire(prologue_kg_producer_state) + for i in cutlass.range_constexpr(self.BK // 4 // 8): + col_base = bk_col_base + i * 8 + chunk_kg = cute.local_tile(rKG_bf16, (8,), (i,)) + smem_store_bf16x8_sw128(sK_raw_ptr, row, col_base, chunk_kg) + + cute.arch.fence_proxy("async.shared", space="cta") + pipeline_prologue_kg.producer_commit(prologue_kg_producer_state) + prologue_kg_producer_state.advance() + + # wait for dkgb + pipeline_mma_dkgb.consumer_wait(mma_dgkb_consumer_state) + tcgen05_fence_after() + dkgb_i32 = tcgen05_ld_32x32b(bk_num_cols_per_wg, TMEM_DKGB_ACC_OFF + wg_idx * bk_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_mma_dkgb.consumer_release(mma_dgkb_consumer_state) + mma_dgkb_consumer_state.advance() + + # db += sum(dkgb * kg, axis=1) + dkgb_f32 = reinterpret_cast(dkgb_i32, Int32, bk_num_cols_per_wg, Float32) + dkgb_f32_val = TensorSSA(dkgb_f32, (bk_num_cols_per_wg,), Float32) + rKgb_kg = cute.make_rmem_tensor((bk_num_cols_per_wg,), Float32) + rKgb_kg.store(dkgb_f32_val * rKG_val) + + if row < sub_seq_len: + for i in cutlass.range_constexpr(bk_num_cols_per_wg): + db_val += rKgb_kg[i] + + # Deterministic db reduction without atomicAdd. + # 4 partitions per row come from 4 warps (warp_row_tile in {0,1}, + # warp_col_tile in {0,1}) x 2 wgs. Reduce in a fixed order so + # the result is bitwise reproducible across launches: + # Phase 1: warp_col_tile==0 writes its db_val into + # sDb[row, wg_idx] (single writer per slot) + # Phase 2: warp_col_tile==1 RMW-adds its db_val into the + # same slot (still single writer per slot) + # Phase 3: WG0 sums the 2 wg-slots in fixed order and stores + # to GMEM. + # No race, no atomic, no fp ordering nondeterminism. + if warp_col_tile == 0 and row < sub_seq_len: + sDb[row, wg_idx] = db_val + self.cuda_wg_sync_barrier.arrive_and_wait() + if warp_col_tile == 1 and row < sub_seq_len: + sDb[row, wg_idx] = sDb[row, wg_idx] + db_val + self.cuda_wg_sync_barrier.arrive_and_wait() + # store db to GMEM (WG0 only). Sum order is fixed (slot 0 + slot 1). + if wg_idx == 0 and local_tidx < sub_seq_len: + db_sum = sDb[(local_tidx, 0)] + sDb[(local_tidx, 1)] + db_gmem[(tok_offset + tile_idx * self.BT + local_tidx, (i_hv, Int32(0)))] = db_sum + + # dk = dk * exp2(gn[None, :] - g) + pipeline_mma_dk.consumer_wait(mma_dk_consumer_state) + tcgen05_fence_after() + dk_i32 = tcgen05_ld_32x32b(bk_num_cols_per_wg, TMEM_DK_ACC_OFF + wg_idx * bk_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_mma_dk.consumer_release(mma_dk_consumer_state) + mma_dk_consumer_state.advance() + + dk_f32 = reinterpret_cast(dk_i32, Int32, bk_num_cols_per_wg, Float32) + dk_f32_val = TensorSSA(dk_f32, (bk_num_cols_per_wg,), Float32) + + rDk = cute.make_rmem_tensor((bk_num_cols_per_wg,), Float32) + if row < sub_seq_len: + for i in cutlass.range_constexpr(bk_num_cols_per_wg): + exp_g_gn = cute.exp2(sGn[(bk_col_base + i,)] - rG_val[i], fastmath=self.use_fast_math) + rDk[i] = dk_f32_val[i] * exp_g_gn + else: + rDk.fill(Float32(0.0)) + + # kdk = k * dk + rKdk = cute.make_rmem_tensor((bk_num_cols_per_wg,), Float32) + rKdk.store(rK_fp32.load() * rDk.load()) + + # gb = gk_exp * beta[:, None] + rGb = cute.make_rmem_tensor((bk_num_cols_per_wg,), Float32) + rGb.store(rG_exp_val * beta_val) + + # dk = dk + dkgb * gb + rDk.store(rDk.load() + dkgb_f32_val * rGb.load()) + rDk_val = rDk.load() + dk_i32_vec = reinterpret_cast(rDk_val, Float32, bk_num_cols_per_wg, Int32) + # GMEM store dk + # 8 fp32 store each time for store_256b + dk_base_addr = ( + dk_gmem.iterator + (tok_offset + tile_idx * self.BT + row) * HV * K + i_hv * K + bk_col_base + ).toint() + if row < sub_seq_len: + for s in cutlass.range_constexpr(num_stores_f32): + chunk_dk = subvec(dk_i32_vec, s * 8, 8) + store_256b(dk_base_addr + s * 32, chunk_dk) + + # dgk += sum(kdk, axis=0) + # write kdk to G SMEM then do BT-dim reduce + for i in cutlass.range_constexpr(self.BK // 4 // 4): + col_base = bk_col_base + i * 4 + chunk_kdk = cute.local_tile(rKdk, (4,), (i,)) + smem_store_f32x4_sw128(sG_raw_ptr, row, col_base, chunk_kdk) + self.cuda_wg_sync_barrier.arrive_and_wait() + + # dgk *= exp2(gn) + if wg_idx == 0: + sDgk[(local_tidx,)] *= cute.exp2(sGn[(local_tidx,)], fastmath=self.use_fast_math) + + self.cuda_wg_sync_barrier.arrive_and_wait() + if wg_idx == 0: + sum = Float32(0.0) + for r in cutlass.range(self.BT, unroll_full=True): + sum += sG_raw[(r, local_tidx, 0)] + sDgk[(local_tidx,)] += sum + + # dg1 = kg * dkgb * beta[:, None], can reuse kg RMEM + rDg = cute.make_rmem_tensor((bk_num_cols_per_wg,), Float32) + rDg.store(rKG_val * dkgb_f32_val * beta_val) + + pipeline_load_q.consumer_wait(load_q_consumer_state) + # dg2 = q * dq - kdk + dg1 + rQ = cute.make_rmem_tensor((bk_num_cols_per_wg,), self.io_dtype) + if row < sub_seq_len: + for i in cutlass.range_constexpr(self.BK // 4 // 8): + col_base = bk_col_base + i * 8 + vals = smem_load_bf16x8_sw128(sQ_raw_ptr, row, col_base) + rQ[i * 8 + 0] = vals[0] + rQ[i * 8 + 1] = vals[1] + rQ[i * 8 + 2] = vals[2] + rQ[i * 8 + 3] = vals[3] + rQ[i * 8 + 4] = vals[4] + rQ[i * 8 + 5] = vals[5] + rQ[i * 8 + 6] = vals[6] + rQ[i * 8 + 7] = vals[7] + else: + rQ.fill(BFloat16(0.0)) + dq_scaled_i32 = tcgen05_ld_32x32b(bk_num_cols_per_wg, TMEM_DQ_SCALED_OFF + wg_idx * bk_num_cols_per_wg) + cute.arch.fence_view_async_tmem_load() + dq_scaled_f32 = reinterpret_cast(dq_scaled_i32, Int32, bk_num_cols_per_wg, Float32) + dq_scaled_f32_val = TensorSSA(dq_scaled_f32, (bk_num_cols_per_wg,), Float32) + rDg.store(rQ.load().to(Float32) * dq_scaled_f32_val + rDg.load() - rKdk.load()) + + self.cuda_wg_sync_barrier.arrive_and_wait() + # dg = dg2 + m_last * dgk, GMEM store dg + if row == sub_seq_len - 1: + for i in cutlass.range_constexpr(bk_num_cols_per_wg): + col = bk_col_base + i + rDg[i] += sDgk[(col,)] + + # Stage dg to SMEM first. A dedicated store warp later does + # SMEM -> RMEM -> GMEM with store_256b, keeping GMEM store + # address/vector live ranges out of the high-register CC path. + pipeline_store_dg.producer_acquire(store_dg_producer_state) + if row < sub_seq_len: + for i in cutlass.range_constexpr(bk_num_cols_per_wg // 4): + col_base = bk_col_base + i * 4 + chunk_dg = cute.local_tile(rDg, (4,), (i,)) + smem_store_f32x4_sw128(sG_raw_ptr, row, col_base, chunk_dg) + + cute.arch.fence_proxy("async.shared", space="cta") + pipeline_store_dg.producer_commit(store_dg_producer_state) + store_dg_producer_state.advance() + + pipeline_load_g.consumer_release(load_g_consumer_state) + load_g_consumer_state.advance() + + pipeline_mma_dA.consumer_wait(mma_dA_consumer_state) + tcgen05_fence_after() + dA_i32 = tcgen05_ld_32x32b(bt_num_cols_per_wg, TMEM_DA_ACC_OFF + wg_idx * bt_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_mma_dA.consumer_release(mma_dA_consumer_state) + mma_dA_consumer_state.advance() + # NOTE: only release k smem after dA finished, because kg reuses k smem in dA += dw @ kg^T + pipeline_load_k.consumer_release(load_k_consumer_state) + load_k_consumer_state.advance() + + # dA = dA * beta[None, :], apply strict lower-triangular mask. + # Triton reference multiplies by the column beta (`b_beta[None, :]`) + # and keeps only `row > col`. + dA_f32 = reinterpret_cast(dA_i32, Int32, bt_num_cols_per_wg, Float32) + dA_f32_val = TensorSSA(dA_f32, (bt_num_cols_per_wg,), Float32) + rDA = cute.make_rmem_tensor((bt_num_cols_per_wg,), BFloat16) + for i in cutlass.range_constexpr(bt_num_cols_per_wg): + col = bt_col_base + i + beta_col = sBeta[(col,)] + dA_scaled = (dA_f32_val[i] * beta_col).to(BFloat16) + if col < row: + rDA[i] = dA_scaled + else: + rDA[i] = BFloat16(0.0) + if row >= sub_seq_len: + rDA.fill(BFloat16(0.0)) + + pipeline_prologue_dA2.producer_acquire(prologue_dA2_producer_state) + + for i in cutlass.range_constexpr(bt_num_cols_per_wg // 8): + col_base = bt_col_base + i * 8 + chunk_dA = cute.local_tile(rDA, (8,), (i,)) + smem_store_bf16x8_sw128(sQ_raw_ptr, row, col_base, chunk_dA) + # notify dA2 = dA @ A + cute.arch.fence_proxy("async.shared", space="cta") + pipeline_prologue_dA2.producer_commit(prologue_dA2_producer_state) + prologue_dA2_producer_state.advance() + + pipeline_load_beta.consumer_release(load_beta_consumer_state) + load_beta_consumer_state.advance() + + # wait for dA2 + pipeline_mma_dA2.consumer_wait(mma_dA2_consumer_state) + tcgen05_fence_after() + dA2_i32 = tcgen05_ld_32x32b(bt_num_cols_per_wg, TMEM_DA2_ACC_OFF + wg_idx * bt_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + pipeline_prologue_dA3.producer_acquire(prologue_dA3_producer_state) + # write dA2 to smem notify dA2 = A @ dA2 + dA2_f32 = reinterpret_cast(dA2_i32, Int32, bt_num_cols_per_wg, Float32) + dA2_f32_val = TensorSSA(dA2_f32, (bt_num_cols_per_wg,), Float32) + rDA2 = cute.make_rmem_tensor((bt_num_cols_per_wg,), BFloat16) + if row < sub_seq_len: + rDA2.store(dA2_f32_val.to(BFloat16)) + else: + rDA2.fill(BFloat16(0.0)) + for i in cutlass.range_constexpr(bt_num_cols_per_wg // 8): + col_base = bt_col_base + i * 8 + chunk_dA2 = cute.local_tile(rDA2, (8,), (i,)) + smem_store_bf16x8_sw128(sQ_raw_ptr, row, col_base, chunk_dA2) + + cute.arch.fence_proxy("async.shared", space="cta") + pipeline_prologue_dA3.producer_commit(prologue_dA3_producer_state) + prologue_dA3_producer_state.advance() + + # wait for dA2 + pipeline_mma_dA3.consumer_wait(mma_dA3_consumer_state) + tcgen05_fence_after() + dA3_i32 = tcgen05_ld_32x32b(bt_num_cols_per_wg, TMEM_DA2_ACC_OFF + wg_idx * bt_num_cols_per_wg) + tcgen05_fence_before() + cute.arch.fence_view_async_tmem_load() + + # release mma dA2 after dA3 is finished, protect DA2 TMEM + pipeline_mma_dA2.consumer_release(mma_dA2_consumer_state) + mma_dA2_consumer_state.advance() + pipeline_mma_dA3.consumer_release(mma_dA3_consumer_state) + mma_dA3_consumer_state.advance() + # NOTE: release smem Q because we reuse to store bf16 dA + pipeline_load_q.consumer_release(load_q_consumer_state) + load_q_consumer_state.advance() + + # dA = -dA, apply strict lower-triangular mask + dA3_f32 = reinterpret_cast(dA3_i32, Int32, bt_num_cols_per_wg, Float32) + dA3_f32_val = TensorSSA(dA3_f32, (bt_num_cols_per_wg,), Float32) + rDA3 = cute.make_rmem_tensor((bt_num_cols_per_wg,), Float32) + rDA3.store(-dA3_f32_val) + for i in cutlass.range_constexpr(bt_num_cols_per_wg): + col = bt_col_base + i + if col >= row: + rDA3[i] = Float32(0.0) + rDA3_val = rDA3.load() + dA3_i32_vec = reinterpret_cast(rDA3_val, Float32, bt_num_cols_per_wg, Int32) + # GMEM store dA + num_stores_dA = bt_num_cols_per_wg // 8 + dA_base_addr = ( + dA_gmem.iterator + (tok_offset + tile_idx * self.BT + row) * HV * BT + i_hv * BT + bt_col_base + ).toint() + if row < sub_seq_len: + for s in cutlass.range_constexpr(num_stores_dA): + chunk_dA_store = subvec(dA3_i32_vec, s * 8, 8) + store_256b(dA_base_addr + s * 32, chunk_dA_store) + + # Load loop body + elif warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_others) + + load_A_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.a_stage) + load_dv_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_do_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_vnew_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_g_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kloop_stage) + load_k_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kloop_stage) + load_q_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kloop_stage) + + vloop_stage_idx = 0 + vloop_phase = 1 # init as 1 for producer + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + G = HV // H + i_t = work_idx // HV # chunk index (global) + i_hv = work_idx % HV # value-head index + i_h = i_hv // G # q/k head index + + # Decode chunk_indices + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + tok_offset = cu_seqlens[(batch_idx,)] + seq_len = cu_seqlens[(batch_idx + 1,)] - tok_offset + sub_seq_len = min(self.BT, seq_len - tile_idx * self.BT) + + # Load A + tma_A_v = cute.domain_offset((0, tok_offset, (0, 0)), tma_tensor_A) + tAsA, tAgA = self._tma_partition_A( + tma_atom_A, + tma_A_v, + sA, + self.dvb_tiler, # [BT, BV, BT] + dvb_tiled_mma, + Int32(0), + i_hv, + ) + pipeline_load_A.producer_acquire(load_A_producer_state) + cute.copy( + tma_atom_A, + tAgA[(None, 0, tile_idx)], + tAsA[(None, load_A_producer_state.index)], + tma_bar_ptr=pipeline_load_A.producer_get_barrier(load_A_producer_state), + ) + load_A_producer_state.advance() + + # V-loop + for v_iter in cutlass.range(self.num_v_tiles): + tma_h_v = cute.domain_offset((0, v_iter * self.BV, (0, 0)), tma_tensor_h) + tHsH, tHgH = self._tma_partition_B( + tma_atom_h, + tma_h_v, + sH, + self.vloop_gemm_tiler, # [BT, BK, BV] + vloop_tiled_mma, + i_hv, + i_t, + ) + mbarrier_wait(bar_mma_cuda_h_ptr + vloop_stage_idx, vloop_phase) + with elect_one(): + mbarrier_arrive_and_expect_tx(bar_tma_h_ptr + vloop_stage_idx, self.tma_bytes_h) + cute.copy( + tma_atom_h, + tHgH[(None, 0, 0)], + tHsH[(None, vloop_stage_idx)], + tma_bar_ptr=bar_tma_h_ptr + vloop_stage_idx, + ) + + tma_dh_v = cute.domain_offset((0, v_iter * self.BV, (0, 0)), tma_tensor_dh) + tDHsDH, tDHgDH = self._tma_partition_B( + tma_atom_dh, + tma_dh_v, + sDh, + self.vloop_gemm_tiler, # [BT, BK, BV] + vloop_tiled_mma, + i_hv, + i_t, + ) + mbarrier_wait(bar_mma_cuda_dh_ptr + vloop_stage_idx, vloop_phase) + with elect_one(): + mbarrier_arrive_and_expect_tx(bar_tma_dh_ptr + vloop_stage_idx, self.tma_bytes_dh) + cute.copy( + tma_atom_dh, + tDHgDH[(None, 0, 0)], + tDHsDH[(None, vloop_stage_idx)], + tma_bar_ptr=bar_tma_dh_ptr + vloop_stage_idx, + ) + + tma_do_v = cute.domain_offset((tok_offset, v_iter * self.BV, (0, 0)), tma_tensor_do) + tDOsDo, tDOgDo = self._tma_partition_A( + tma_atom_do, + tma_do_v, + sDo, + self.vloop_gemm_tiler, # [BT, BK, BV] + vloop_tiled_mma, + Int32(0), + i_hv, + ) + pipeline_load_do.producer_acquire(load_do_producer_state) + cute.copy( + tma_atom_do, + tDOgDo[(None, tile_idx, 0)], + tDOsDo[(None, vloop_stage_idx)], + tma_bar_ptr=pipeline_load_do.producer_get_barrier(load_do_producer_state), + ) + load_do_producer_state.advance() + + tma_dv_v = cute.domain_offset((tok_offset, v_iter * self.BV, (0, 0)), tma_tensor_dv) + tDVsDv, tDVgDV = self._tma_partition_A( + tma_atom_dv, + tma_dv_v, + sDv, + self.vloop_gemm_tiler, # [BT, BK, BV] + vloop_tiled_mma, + Int32(0), + i_hv, + ) + pipeline_load_dv.producer_acquire(load_dv_producer_state) + cute.copy( + tma_atom_dv, + tDVgDV[(None, tile_idx, 0)], + tDVsDv[(None, vloop_stage_idx)], + tma_bar_ptr=pipeline_load_dv.producer_get_barrier(load_dv_producer_state), + ) + load_dv_producer_state.advance() + + tma_v_v = cute.domain_offset((tok_offset, v_iter * self.BV, (0, 0)), tma_tensor_v) + tVsV, tVgV = self._tma_partition_B( + tma_atom_v, + tma_v_v, + sV, + self.dA_vloop_tiler, # [BT, BT, BV] + dA_vloop_tiled_mma, + Int32(0), + i_hv, + ) + mbarrier_wait(bar_mma_cuda_v_ptr + vloop_stage_idx, vloop_phase) + with elect_one(): + mbarrier_arrive_and_expect_tx(bar_tma_v_ptr + vloop_stage_idx, self.tma_bytes_v) + cute.copy( + tma_atom_v, + tVgV[(None, tile_idx, 0)], + tVsV[(None, vloop_stage_idx)], + tma_bar_ptr=bar_tma_v_ptr + vloop_stage_idx, + ) + + # load v_new + tma_vnew_v = cute.domain_offset((tok_offset, v_iter * self.BV, (0, 0)), tma_tensor_vnew) + tVnewsVnew, tVnewgVnew = self._tma_partition_A( + tma_atom_vnew, + tma_vnew_v, + sVnew, + self.vloop_gemm_tiler, # [BT, BK, BV] + vloop_tiled_mma, + Int32(0), + i_hv, + ) + pipeline_load_vnew.producer_acquire(load_vnew_producer_state) + cute.copy( + tma_atom_vnew, + tVnewgVnew[(None, tile_idx, 0)], + tVnewsVnew[(None, vloop_stage_idx)], + tma_bar_ptr=pipeline_load_vnew.producer_get_barrier(load_vnew_producer_state), + ) + load_vnew_producer_state.advance() + + vloop_stage_idx = (vloop_stage_idx + 1) % self.vloop_stage + vloop_phase ^= 1 + + # Load g + tma_g_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_g) + tGsG, tGgG = self._epilog_partition_varlen( + tma_atom_g, + tma_g_v[None, None, (i_hv, Int32(0))], + (self.BT, self.BK), + sG_raw, + ) + pipeline_load_g.producer_acquire(load_g_producer_state) + cute.copy( + tma_atom_g, + tGgG[(None, tile_idx, 0)], + tGsG[(None, 0)], # hardcode stage to 0 because kloop_stage is 1 + tma_bar_ptr=pipeline_load_g.producer_get_barrier(load_g_producer_state), + ) + load_g_producer_state.advance() + + # Load k + tma_k_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_k) + tKsK, tKgK = self._epilog_partition_varlen( + tma_atom_k, + tma_k_v[None, None, (i_h, Int32(0))], + (self.BT, self.BK), + sK_raw, + ) + pipeline_load_k.producer_acquire(load_k_producer_state) + cute.copy( + tma_atom_k, + tKgK[(None, tile_idx, 0)], + tKsK[(None, 0)], # hardcode stage to 0 because kloop_stage is 1 + tma_bar_ptr=pipeline_load_k.producer_get_barrier(load_k_producer_state), + ) + load_k_producer_state.advance() + + tma_q_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_q) + tQsQ, tQgQ = self._epilog_partition_varlen( + tma_atom_q, + tma_q_v[None, None, (i_h, Int32(0))], + (self.BT, self.BK), + sQ_raw, + ) + pipeline_load_q.producer_acquire(load_q_producer_state) + cute.copy( + tma_atom_q, + tQgQ[(None, tile_idx, 0)], + tQsQ[(None, 0)], # hardcode stage to 0 because kloop_stage is 1 + tma_bar_ptr=pipeline_load_q.producer_get_barrier(load_q_producer_state), + ) + load_q_producer_state.advance() + + # MMA loop body + elif warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_others) + + load_A_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.a_stage) + load_dv_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + mma_dvb_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + load_do_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + load_vnew_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + mma_dq_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + mma_dk_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + mma_dw_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + prologue_dw_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + prologue_kg_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + mma_dgkb_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + mma_dA_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + mma_dA2_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + mma_dA3_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_stage) + prologue_dA2_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + prologue_dA3_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_stage) + + vloop_stage_idx = 0 + a_stage_idx = 0 + mma_vloop_phase = 0 + vloop_phase = 0 + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + G = HV // H + i_t = work_idx // HV # chunk index (global) + i_hv = work_idx % HV # value-head index (unused in MMA warp) + i_h = i_hv // G # q/k head index (unused in MMA warp) + + # Decode chunk_indices + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + tok_offset = cu_seqlens[(batch_idx,)] + seq_len = cu_seqlens[(batch_idx + 1,)] - tok_offset + sub_seq_len = min(self.BT, seq_len - tile_idx * self.BT) + + zeros8 = cute.make_rmem_tensor((8,), dtype=self.io_dtype) + zeros8.fill(BFloat16(0.0)) + + pipeline_load_A.consumer_wait(load_A_consumer_state) + sA_raw_ptr = cute.make_ptr( + self.io_dtype, + sA_ptr_base + a_stage_idx * A_bytes_per_stage, + cute.AddressSpace.smem, + ) + if sub_seq_len < self.BT: + for i in cutlass.range_constexpr(self.BT // 32): + row = i * 32 + lane_idx + if row >= sub_seq_len: + for col in cutlass.range_constexpr(self.BT // 8): + # A tile is MN_SW128 in shared memory; use raw swizzled + # address stores to avoid layout-coordinate ambiguity. + smem_store_bf16x8_sw128(sA_raw_ptr, row, col * 8, zeros8) + # Make generic-proxy SMEM stores visible to UMMA async-proxy readers. + cute.arch.fence_proxy("async.shared", space="cta") + + for v_iter in cutlass.range(self.num_v_tiles): + is_accum = False if v_iter == 0 else True + mbarrier_wait(bar_tma_h_ptr + vloop_stage_idx, vloop_phase) + pipeline_load_do.consumer_wait(load_do_consumer_state) + sDo_raw_ptr = cute.make_ptr( + self.io_dtype, + sDo_ptr_base + vloop_stage_idx * vloop_opA_bytes_per_stage, + cute.AddressSpace.smem, + ) + if sub_seq_len < self.BT: + for i in cutlass.range_constexpr(self.BT // 32): + row = i * 32 + lane_idx + if row >= sub_seq_len: + for col in cutlass.range_constexpr(self.BV // 8): + # dv tile uses the same Swizzle<3,4,3> physical mapping. + smem_store_bf16x8_sw128(sDo_raw_ptr, row, col * 8, zeros8) + cute.arch.fence_proxy("async.shared", space="cta") + + if v_iter == 0: + pipeline_mma_dq.producer_acquire(mma_dq_producer_state) + + # dq+=do@h + sDo_k_cur = sDo_k[(None, None, None, vloop_stage_idx)] + sH_k_cur = sH_k[(None, None, None, vloop_stage_idx)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDo_k_cur.iterator, sDo_k_cur.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sH_k_cur.iterator, sH_k_cur.layout, "k")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n128_k_k_call( + vloop_opA_smem, desc_a_base, vloop_opB_smem, desc_b_base, TMEM_DQ_ACC_OFF, self.BV, is_accum + ) + + pipeline_load_do.consumer_release(load_do_consumer_state) + load_do_consumer_state.advance() + + if v_iter == self.num_v_tiles - 1: + pipeline_mma_dq.producer_commit(mma_dq_producer_state) + mma_dq_producer_state.advance() + + pipeline_load_dv.consumer_wait(load_dv_consumer_state) + sDv_raw = cute.make_ptr( + self.io_dtype, + sDv_ptr_base + vloop_stage_idx * vloop_opA_bytes_per_stage, + cute.AddressSpace.smem, + ) + if sub_seq_len < self.BT: + for i in cutlass.range_constexpr(self.BT // 32): + row = i * 32 + lane_idx + if row >= sub_seq_len: + for col in cutlass.range_constexpr(self.BV // 8): + # dv tile uses the same Swizzle<3,4,3> physical mapping. + smem_store_bf16x8_sw128(sDv_raw, row, col * 8, zeros8) + cute.arch.fence_proxy("async.shared", space="cta") + + # if lane_idx == 0: + # cute.printf("V_iter", v_iter) + # cute.print_tensor(sDv[None, None, None, vloop_stage_idx]) + pipeline_mma_dvb.producer_acquire(mma_dvb_producer_state) + sDv_mn_cur = sDv_mn[(None, None, None, vloop_stage_idx)] + sA_mn_cur = sA_mn[(None, None, None, a_stage_idx)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sA_mn_cur.iterator, sA_mn_cur.layout, "mn")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDv_mn_cur.iterator, sDv_mn_cur.layout, "mn")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n64_mn_mn_call( + A_mn_opA_smem, desc_a_base, dv_mn_opB_smem, desc_b_base, TMEM_FLEX_OFF, self.BT + ) + + pipeline_mma_dvb.producer_commit(mma_dvb_producer_state) + mma_dvb_producer_state.advance() + + # dw += dv @ h + if v_iter == 0: + pipeline_mma_dw.producer_acquire(mma_dw_producer_state) + + sDv_k_cur = sDv_k[(None, None, None, vloop_stage_idx)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDv_k_cur.iterator, sDv_k_cur.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sH_k_cur.iterator, sH_k_cur.layout, "k")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n128_k_k_call( + vloop_opA_smem, desc_a_base, vloop_opB_smem, desc_b_base, TMEM_DW_ACC_OFF, self.BV, is_accum + ) + + # dA += dv @ v^T + mbarrier_wait(bar_tma_v_ptr + vloop_stage_idx, vloop_phase) + sV_raw = cute.make_ptr( + self.io_dtype, sV_ptr_base + vloop_stage_idx * v_opB_bytes_per_stage, cute.AddressSpace.smem + ) + if sub_seq_len < self.BT: + for i in cutlass.range_constexpr(self.BT // 32): + row = i * 32 + lane_idx + if row >= sub_seq_len: + for col in cutlass.range_constexpr(self.BV // 8): + # dv tile uses the same Swizzle<3,4,3> physical mapping. + smem_store_bf16x8_sw128(sV_raw, row, col * 8, zeros8) + cute.arch.fence_proxy("async.shared", space="cta") + + if v_iter == 0: + pipeline_mma_dA.producer_acquire(mma_dA_producer_state) + + sV_k_cur = sV_k[(None, None, None, vloop_stage_idx)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDv_k_cur.iterator, sDv_k_cur.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sV_k_cur.iterator, sV_k_cur.layout, "k")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n64_k_k_call( + vloop_opA_smem, desc_a_base, v_opB_smem, desc_b_base, TMEM_DA_ACC_OFF, self.BV, is_accum + ) + + # dv pipeline calls tcgen05.commit for dv@h and dv@v^T + pipeline_load_dv.consumer_release(load_dv_consumer_state) + load_dv_consumer_state.advance() + + if v_iter == self.num_v_tiles - 1: + pipeline_mma_dw.producer_commit(mma_dw_producer_state) + mma_dw_producer_state.advance() + + umma_arrive(bar_mma_cuda_h_ptr + vloop_stage_idx) + umma_arrive(bar_mma_cuda_v_ptr + vloop_stage_idx) + + # dk += v_new @ dh + pipeline_load_vnew.consumer_wait(load_vnew_consumer_state) + sDvnew_raw_ptr = cute.make_ptr( + self.io_dtype, + sVnew_ptr_base + vloop_stage_idx * vloop_opA_bytes_per_stage, + cute.AddressSpace.smem, + ) + if sub_seq_len < self.BT: + for i in cutlass.range_constexpr(self.BT // 32): + row = i * 32 + lane_idx + if row >= sub_seq_len: + for col in cutlass.range_constexpr(self.BV // 8): + # dv tile uses the same Swizzle<3,4,3> physical mapping. + smem_store_bf16x8_sw128(sDvnew_raw_ptr, row, col * 8, zeros8) + cute.arch.fence_proxy("async.shared", space="cta") + + mbarrier_wait(bar_tma_dh_ptr + vloop_stage_idx, vloop_phase) + if v_iter == 0: + pipeline_mma_dk.producer_acquire(mma_dk_producer_state) + + sVnew_k_cur = sVnew_k[(None, None, None, vloop_stage_idx)] + sDh_k_cur = sDh_k[(None, None, None, vloop_stage_idx)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sVnew_k_cur.iterator, sVnew_k_cur.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDh_k_cur.iterator, sDh_k_cur.layout, "k")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n128_k_k_call( + vloop_opA_smem, desc_a_base, vloop_opB_smem, desc_b_base, TMEM_DK_ACC_OFF, self.BV, is_accum + ) + + # vnew pipeline calls tcgen05.commit + pipeline_load_vnew.consumer_release(load_vnew_consumer_state) + load_vnew_consumer_state.advance() + + if v_iter == self.num_v_tiles - 1: + pipeline_mma_dk.producer_commit(mma_dk_producer_state) + mma_dk_producer_state.advance() + + umma_arrive(bar_mma_cuda_dh_ptr + vloop_stage_idx) + + # add tcgen05.commit and mbar.wait to make sure dq/dk/dw MMA finished + umma_arrive(bar_mma_done_vloop_ptr + 0) + mbarrier_wait(bar_mma_done_vloop_ptr + 0, mma_vloop_phase) + mma_vloop_phase ^= 1 + + vloop_stage_idx = (vloop_stage_idx + 1) % self.vloop_stage + vloop_phase ^= 1 + + pipeline_prologue_dw.consumer_wait(prologue_dw_consumer_state) + cute.arch.fence_proxy("async.shared", space="cta") + # dkgb = A @ dw + pipeline_mma_dkgb.producer_acquire(mma_dgkb_producer_state) + sA_mn_cur = sA_mn[(None, None, None, a_stage_idx)] + sDw_mn_cur = sDw_mn[(None, None, None, 0)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sA_mn_cur.iterator, sA_mn_cur.layout, "mn")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDw_mn_cur.iterator, sDw_mn_cur.layout, "mn")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n128_mn_mn_call( + A_mn_opA_smem, desc_a_base, dw_mn_opB_smem, desc_b_base, TMEM_DKGB_ACC_OFF, self.BT + ) + + pipeline_mma_dkgb.producer_commit(mma_dgkb_producer_state) + mma_dgkb_producer_state.advance() + + pipeline_prologue_kg.consumer_wait(prologue_kg_consumer_state) + cute.arch.fence_proxy("async.shared", space="cta") + # dA += dw @ kg^T + sDw_k_cur = sDw_k[(None, None, None, 0)] + sKG_k_cur = sKG_k[(None, None, None, 0)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDw_k_cur.iterator, sDw_k_cur.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sKG_k_cur.iterator, sKG_k_cur.layout, "k")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n64_k_k_call( + dw_k_opA_smem, desc_a_base, kg_k_opB_smem, desc_b_base, TMEM_DA_ACC_OFF, self.BK, True + ) + + pipeline_mma_dA.producer_commit(mma_dA_producer_state) + mma_dA_producer_state.advance() + pipeline_prologue_kg.consumer_release(prologue_kg_consumer_state) + prologue_kg_consumer_state.advance() + + pipeline_prologue_dw.consumer_release(prologue_dw_consumer_state) + prologue_dw_consumer_state.advance() + + # dA2 = dA @ A + pipeline_mma_dA2.producer_acquire(mma_dA2_producer_state) + pipeline_prologue_dA2.consumer_wait(prologue_dA2_consumer_state) + cute.arch.fence_proxy("async.shared", space="cta") + + sDA_k_cur = sDA_k[(None, None, None, 0)] + sA_k_cur = sA_k[(None, None, None, a_stage_idx)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDA_k_cur.iterator, sDA_k_cur.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sA_k_cur.iterator, sA_k_cur.layout, "k")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n64_k_k_call(dA_k_opA_smem, desc_a_base, A_k_opB_smem, desc_b_base, TMEM_DA2_ACC_OFF, self.BT) + + pipeline_mma_dA2.producer_commit(mma_dA2_producer_state) + mma_dA2_producer_state.advance() + pipeline_prologue_dA2.consumer_release(prologue_dA2_consumer_state) + prologue_dA2_consumer_state.advance() + + # dA3 = A @ dA2 + pipeline_mma_dA3.producer_acquire(mma_dA3_producer_state) + pipeline_prologue_dA3.consumer_wait(prologue_dA3_consumer_state) + cute.arch.fence_proxy("async.shared", space="cta") + + sA_mn_cur = sA_mn[(None, None, None, a_stage_idx)] + sDA_mn_cur = sDA_mn[(None, None, None, 0)] + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(sA_mn_cur.iterator, sA_mn_cur.layout, "mn")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(sDA_mn_cur.iterator, sDA_mn_cur.layout, "mn")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + mma_ws_ss_m64n64_mn_mn_call(A_mn_opA_smem, desc_a_base, dA_mn_opB_smem, desc_b_base, TMEM_DA2_ACC_OFF, self.BT) + + pipeline_mma_dA3.producer_commit(mma_dA3_producer_state) + mma_dA3_producer_state.advance() + pipeline_prologue_dA3.consumer_release(prologue_dA3_consumer_state) + prologue_dA3_consumer_state.advance() + + pipeline_load_A.consumer_release(load_A_consumer_state) + load_A_consumer_state.advance() + + a_stage_idx = (a_stage_idx + 1) % self.a_stage + + # Load aux loop body + elif warp_idx in self.aux_warp_ids: + cute.arch.setmaxregister_decrease(self.num_regs_others) + tidx = thread_idx - (self.threads_per_cta - 64) + + load_beta_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + load_g_store_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + store_dg_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kloop_stage) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + G = HV // H + i_t = work_idx // HV # chunk index (global) + i_hv = work_idx % HV # value-head index + i_h = i_hv // G # q/k head index (unused in aux warp) + + # Decode chunk_indices + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + tok_offset = cu_seqlens[(batch_idx,)] + seq_len = cu_seqlens[(batch_idx + 1,)] - tok_offset + sub_seq_len = min(self.BT, seq_len - tile_idx * self.BT) + + pipeline_load_beta.producer_acquire(load_beta_producer_state) + beta_f32 = Float32(0.0) + if tidx < sub_seq_len: + beta_f32 = Float32(beta_gmem[(tok_offset + tile_idx * self.BT + tidx, (i_hv, Int32(0)))]) + sBeta[(tidx,)] = beta_f32 + + cute.arch.fence_proxy("async.shared", space="cta") + pipeline_load_beta.producer_commit(load_beta_producer_state) + load_beta_producer_state.advance() + + pipeline_load_g.consumer_wait(load_g_store_consumer_state) + pipeline_store_dg.consumer_wait(store_dg_consumer_state) + + tma_dg_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_dg) + tDGsDG, tDGgDG = self._epilog_partition_varlen( + tma_atom_dg, + tma_dg_v[None, None, (i_hv, Int32(0))], + (self.BT, self.BK), + sG_raw, + ) + if sub_seq_len < self.BT: + # Tail chunk, direct store + store_lane_row = tidx >> Int32(4) # 0..3 + store_col_base = (tidx & Int32(15)) * Int32(8) # 0,8,...,120 + for row_quad in cutlass.range_constexpr(self.BT // 4): + store_row = row_quad * 4 + store_lane_row + if store_row < sub_seq_len: + vals0 = smem_load_f32x4_sw128(sG_raw_ptr, store_row, store_col_base) + vals1 = smem_load_f32x4_sw128(sG_raw_ptr, store_row, store_col_base + Int32(4)) + dg_store_rmem = cute.make_rmem_tensor((8,), Float32) + dg_store_rmem[0] = vals0[0] + dg_store_rmem[1] = vals0[1] + dg_store_rmem[2] = vals0[2] + dg_store_rmem[3] = vals0[3] + dg_store_rmem[4] = vals1[0] + dg_store_rmem[5] = vals1[1] + dg_store_rmem[6] = vals1[2] + dg_store_rmem[7] = vals1[3] + dg_store_i32_vec = reinterpret_cast(dg_store_rmem.load(), Float32, 8, Int32) + dg_base_addr = ( + dg_gmem.iterator + + (tok_offset + tile_idx * self.BT + store_row) * HV * K + + i_hv * K + + store_col_base + ).toint() + store_256b(dg_base_addr, dg_store_i32_vec) + else: + # Non-tail chunk, TMA store + cute.arch.fence_proxy("async.shared", space="cta") + cute.copy( + tma_atom_dg, + tDGsDG[(None, 0)], # hardcode stage to 0 because kloop_stage is 1 + tDGgDG[(None, tile_idx, 0)], + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + pipeline_store_dg.consumer_release(store_dg_consumer_state) + store_dg_consumer_state.advance() + pipeline_load_g.consumer_release(load_g_store_consumer_state) + load_g_store_consumer_state.advance() + + # ===================== TMEM cleanup ===================== + tmem.relinquish_alloc_permit() + self.tmem_dealloc_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr, TMEM_TOTAL) + + @cute.jit + def _tma_partition_A(self, tma_atom, tma_tensor, smem, tile_shape, tiled_mma, batch_idx, hidx): + """Partition a TMA tensor as MMA A-operand (M,K dims). + + ``tma_tensor`` should already have domain_offset applied for varlen. + + For tile_shape = (BT, BK, BV) = (M, N, K): + coord = (None, 0, None) — slices out the N-tile axis (mode 1) at 0, + leaving mode 0 (M=BT) and mode 2 (K=BV) free for TMA to iterate. + + Returns (tXsX, tXgX) — SMEM partition and GMEM coordinate partition. + """ + coord = (None, 0, None) + gX = cute.local_tile(tma_tensor, cute.slice_(tile_shape, coord), (None, None, (hidx, batch_idx))) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_A(gX) + tXsX, tXgX = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 3), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + @cute.jit + def _tma_partition_B(self, tma_atom, tma_tensor, smem, tile_shape, tiled_mma, batch_idx, hidx): + """Partition a TMA tensor as MMA B-operand (N,K dims). + + Mirrors the identical helper in recompute_wu.py / fwd_o.py. + ``tma_tensor`` should already have domain_offset applied for varlen. + + For tile_shape = (BT, BK, BV) = (M, N, K): + coord = (0, None, None) — slices out the M-tile axis (mode 0) at 0, + leaving mode 1 (N=BK) and mode 2 (K=BV) free for TMA to iterate. + + Returns (tXsX, tXgX) — SMEM partition and GMEM coordinate partition. + """ + coord = (0, None, None) + gX = cute.local_tile(tma_tensor, cute.slice_(tile_shape, coord), (None, None, (hidx, batch_idx))) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_B(gX) + tXsX, tXgX = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 3), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + @cute.jit + def _epilog_partition_varlen(self, atom, gC_2d, epi_tile, sC): + """Partition for varlen epilog TMA load (2D tensor with domain_offset). + + Uses local_tile instead of flat_divide to correctly preserve TMA basis + stride coordinates through domain_offset. Matches Flash Attention's + pattern: slice mode2 → domain_offset(2D) → local_tile → tma_partition. + + Uses (None, None) to keep all tile-count modes, producing the same + rank as _epilog_partition (flat_divide) so copy indexing is unchanged. + """ + gC_tiled = cute.local_tile(gC_2d, epi_tile, (None, None)) + sC_g = cute.group_modes(sC, 0, 2) + gC_g = cute.group_modes(gC_tiled, 0, 2) + bSG_sC, bSG_gC = cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + sC_g, + gC_g, + ) + return bSG_sC, bSG_gC + + +# ===================================================================== +# Compilation & Cache +# ===================================================================== + +_bwd_wy_kernel_cache: dict = {} + + +def _compile_bwd_wy_variant(H, HV, K, V, scale, chunk_size, beta_dtype, use_fast_math): + """Compile one ChunkKdaBwdWyDqkgFused kernel variant. + + Uses make_fake_compact_tensor and make_fake_stream for compilation with + TVM-FFI. At runtime, torch tensors are passed directly (zero-copy). + Uses sym_int() for dynamic B, T, NT dimensions. + """ + kernel_obj = ChunkKdaBwdWyDqkgFused( + chunk_size=chunk_size, + head_dim_k=K, + head_dim_v=V, + scale=scale, + beta_dtype=beta_dtype, + use_fast_math=use_fast_math, + ) + + sym_b = cute.sym_int() # T (non-varlen) or T_total (varlen) + sym_nt = cute.sym_int() # NT_total + sym_cu = cute.sym_int() # cu_seqlens size + sym_ci = cute.sym_int() # chunk_indices rows + + BT = chunk_size + + # only support varlen for real-world use cases + # varlen: data tensors are [1, T_total, H, ...] + q_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + k_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + v_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128) + vnew_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128) + g_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128) + beta_fake = make_fake_compact_tensor(beta_dtype, (1, sym_b, HV), stride_order=(2, 1, 0), assumed_align=128) + A_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, HV, BT), stride_order=(3, 2, 1, 0), assumed_align=128) + do_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128) + dv_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128) + + dq_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128) + dk_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128) + dv2_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, HV, V), stride_order=(3, 2, 1, 0), assumed_align=128) + dg_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, HV, K), stride_order=(3, 2, 1, 0), assumed_align=128) + db_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, HV), stride_order=(2, 1, 0), assumed_align=128) + dA_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, HV, BT), stride_order=(3, 2, 1, 0), assumed_align=128) + + h_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_nt, HV, K, V), stride_order=(4, 3, 2, 1, 0), assumed_align=128) + dh_fake = make_fake_compact_tensor( + cutlass.BFloat16, (1, sym_nt, HV, K, V), stride_order=(4, 3, 2, 1, 0), assumed_align=128 + ) + + cu_fake = make_fake_compact_tensor(cutlass.Int32, (sym_cu,), assumed_align=128) + ci_fake = make_fake_compact_tensor(cutlass.Int32, (sym_ci, 2), stride_order=(1, 0), assumed_align=128) + stream_fake = make_fake_stream(use_tvm_ffi_env_stream=True) + + compiled_fn = cute.compile( + kernel_obj, + # Inputs + q_fake, + k_fake, + v_fake, + vnew_fake, + g_fake, + beta_fake, + A_fake, + h_fake, + do_fake, + dh_fake, + dv_fake, + # Outputs + dq_fake, + dk_fake, + dv2_fake, + dg_fake, + db_fake, + dA_fake, + # Metadata + cu_fake, + ci_fake, + (Int32(1), Int32(1), Int32(H), Int32(HV), Int32(K), Int32(V)), + Int32(1), # total_nt dummy + stream_fake, + options=COMPILE_OPTIONS, + ) + return compiled_fn + + +def _get_compiled_bwd_wy(H, HV, K, V, scale, chunk_size, beta_dtype): + """Get a compiled ChunkKdaBwdWyDqkgFused kernel with on-demand (lazy) compilation. + + Cache key: (H, HV, K, V, scale, chunk_size, beta_dtype, USE_FAST_MATH) + """ + key = (H, HV, K, V, scale, chunk_size, beta_dtype, USE_FAST_MATH) + if key not in _bwd_wy_kernel_cache: + _bwd_wy_kernel_cache[key] = _compile_bwd_wy_variant( + H, + HV, + K, + V, + scale, + chunk_size, + _torch_to_cutlass_dtype[beta_dtype], + USE_FAST_MATH, + ) + return _bwd_wy_kernel_cache[key] + + +# ===================================================================== +# Python API (FLA-compatible) +# ===================================================================== + +_bwd_wy_dummy_cu_seqlens = None +_bwd_wy_dummy_chunk_indices = None + + +def chunk_kda_bwd_wy_dqkg_fused( + 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: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + ChunkKdaBwdWyDqkgFused — FLA-compatible Python API. + + Computes backward gradients dq, dk, dv2, db, dg, dA for the KDA + chunkwise delta-rule backward pass using the CuTe DSL Blackwell kernel. + + Returns: + (dq, dk, dv2, db, dg, dA) matching FLA's chunk_kda_bwd_wy_dqkg_fused output order. + """ + B, T, H, K = q.shape + V = v.shape[3] + HV = v.shape[2] + BT = chunk_size + beta_dtype = beta.dtype + device = q.device + + if cu_seqlens is None: + cu_seqlens = prepare_uniform_cu_seqlens(B, T, device, torch.int32) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + + if scale is None: + scale = K**-0.5 + + assert cu_seqlens is not None and chunk_indices is not None + # Ensure cu_seqlens is int32 + assert cu_seqlens.dtype == torch.int32, "cu_seqlens must be int32" + T_total = B * T + num_seqs = cu_seqlens.shape[0] - 1 + total_nt_val = chunk_indices.shape[0] + ps = (Int32(num_seqs), Int32(T_total), Int32(H), Int32(HV), Int32(K), Int32(V)) + + # Allocate output tensors + dq = torch.empty(1, T_total, HV, K, dtype=torch.float32, device=device) + dk = torch.empty(1, T_total, HV, K, dtype=torch.float32, device=device) + dv2 = torch.empty(1, T_total, HV, V, dtype=torch.bfloat16, device=device) + dg = torch.empty(1, T_total, HV, K, dtype=torch.float32, device=device) + db = torch.empty(1, T_total, HV, dtype=torch.float32, device=device) + dA = torch.empty(1, T_total, HV, BT, dtype=torch.float32, device=device) + + compiled_fn = _get_compiled_bwd_wy( + H, + HV, + K, + V, + scale, + chunk_size, + beta_dtype, + ) + + if B != 1: + q = q.reshape(1, T_total, H, K) + k = k.reshape(1, T_total, H, K) + v = v.reshape(1, T_total, HV, V) + v_new = v_new.reshape(1, T_total, HV, V) + g = g.reshape(1, T_total, HV, K) + beta = beta.reshape(1, T_total, HV) + A = A.reshape(1, T_total, HV, BT) + h = h.reshape(1, total_nt_val, HV, K, V) + do = do.reshape(1, T_total, HV, V) + dh = dh.reshape(1, total_nt_val, HV, K, V) + dv = dv.reshape(1, T_total, HV, V) + + # TVM-FFI call + compiled_fn( + # Inputs + q, + k, + v, + v_new, + g, + beta, + A, + h, + do, + dh, + dv, + # Outputs + dq, + dk, + dv2, + dg, + db, + dA, + # Metadata + cu_seqlens, + chunk_indices, + ps, + Int32(total_nt_val), + ) + + # rearrange back + if B != 1: + dq = dq.reshape(B, T, HV, K) + dk = dk.reshape(B, T, HV, K) + dv2 = dv2.reshape(B, T, HV, V) + dg = dg.reshape(B, T, HV, K) + db = db.reshape(B, T, HV) + dA = dA.reshape(B, T, HV, BT) + + return dq, dk, dv2, db, dg, dA + + +# ===================================================================== +# Main (test entry point) +# ===================================================================== + + +def main(): + parser = argparse.ArgumentParser(description="Chunk KDA BWD WY DqKG Fused kernel test") + parser.add_argument("--B", type=int, default=1) + parser.add_argument("--T", type=int, default=64) + parser.add_argument("--H", type=int, default=1) + parser.add_argument("--HV", type=int, default=None, help="Number of value heads (default: H, i.e. no GVA)") + parser.add_argument("--K", type=int, default=128) + parser.add_argument("--V", type=int, default=128) + parser.add_argument("--scale", type=float, default=None) + parser.add_argument("--chunk_size", type=int, default=64) + args = parser.parse_args() + + if args.scale is None: + args.scale = args.K**-0.5 + B, T, H, K, V = args.B, args.T, args.H, args.K, args.V + HV = args.HV if args.HV is not None else H + BT = args.chunk_size + seq_lens = [63, 63, 63] + seq_lens = [64] + total_len = sum(seq_lens) + T = total_len + scale = args.scale + NT = (T + BT - 1) // BT + dtype, device = torch.bfloat16, "cuda" + cu_seqlens = torch.tensor(_exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + + print(f"Config: B={B}, T={T}, H={H}, HV={HV}, K={K}, V={V}, BT={BT}, scale={scale:.4f}") + print(f" Chunks per seq: {NT}, Total chunks: {B * NT}") + print(f" BK={64}, BV={64}, NK={K // 64}, NV={V // 64}") + + # Generate test data (q/k use H heads; all others use HV heads) + torch.manual_seed(42) + q = torch.randn(B, T, H, K, dtype=dtype, device=device) + k = torch.randn(B, T, H, K, dtype=dtype, device=device) + v = torch.randn(B, T, HV, V, dtype=dtype, device=device) + v_new = torch.randn(B, T, HV, V, dtype=dtype, device=device) + g = torch.randn(B, T, HV, K, dtype=torch.float32, device=device) * 0.1 + beta = torch.randn(B, T, HV, dtype=torch.bfloat16, device=device) + A = torch.randn(B, T, HV, BT, dtype=dtype, device=device) * 0.1 + h = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + do_t = torch.randn(B, T, HV, V, dtype=dtype, device=device) + dh = torch.randn(B, NT, HV, K, V, dtype=dtype, device=device) * 0.01 + dv = torch.randn(B, T, HV, V, dtype=dtype, device=device) + + print("\n=== Compilation Test ===") + try: + dq, dk, dv2, db, dg, dA = chunk_kda_bwd_wy_dqkg_fused( + q=q, + k=k, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=A, + h=h, + do=do_t, + dh=dh, + dv=dv, + cu_seqlens=cu_seqlens, + scale=scale, + chunk_size=BT, + ) + torch.cuda.synchronize() + print(f" dq shape: {dq.shape}, dtype: {dq.dtype}") + print(f" dk shape: {dk.shape}, dtype: {dk.dtype}") + print(f" dv2 shape: {dv2.shape}, dtype: {dv2.dtype}") + print(f" dg shape: {dg.shape}, dtype: {dg.dtype}") + print(f" db shape: {db.shape}, dtype: {db.dtype}") + print(f" dA shape: {dA.shape}, dtype: {dA.dtype}") + except Exception as e: + import traceback + + print(f" ERROR: {e}") + traceback.print_exc() + + +if __name__ == "__main__": + main() diff --git a/cula/ops/intrinsics_sm100.py b/cula/ops/intrinsics_sm100.py new file mode 100644 index 00000000..e7c1ce5b --- /dev/null +++ b/cula/ops/intrinsics_sm100.py @@ -0,0 +1,418 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NVVM wrappers for SM100 (Blackwell) Tensor Memory intrinsics. + +Provides low-level, CuteDSL-compatible helpers that move data between +Tensor Memory (TMEM) and registers / shared memory via the native +``nvvm.tcgen05.*`` MLIR ops. + +**T2R / R2T** – ``tcgen05.ld`` / ``tcgen05.st`` with ``.32x32b`` shape. +**S2T** – ``tcgen05.cp`` with ``.128x256b`` shape (SMEM → TMEM) +PTX reference +------------- + tcgen05.ld.sync.aligned.32x32b.xN.b32 {r0, ..., rN-1}, [taddr]; + tcgen05.st.sync.aligned.32x32b.xN.b32 [taddr], {r0, ..., rN-1}; + +where ``N ∈ {2, 4, 8, 16, 32, 64, 128}`` and each ``r`` is a 32-bit +register. ``taddr`` encodes both the TMEM column index (bits [15:0]) +and the lane index (bits [31:16]). + +See https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld + +Usage inside a ``@cute.kernel`` or ``@cute.jit`` function:: + + from cula.ops.intrinsics_sm100 import ( + tcgen05_ld_32x32b, tcgen05_st_32x32b, + reinterpret_cast, subvec, store_256b, + ) + from cutlass.cute.typing import Float32, Int32 + + # Load 32 × 32-bit values from TMEM → opaque vector<32 x i32> + vec_i32 = tcgen05_ld_32x32b(32, taddr) + + # Zero-cost reinterpret as f32 (single vector.bitcast, no instructions) + vec_f32 = reinterpret_cast(vec_i32, Int32, 32, Float32) + + # Store to global via store_256b (4 × 256-bit stores) + # store_256b takes vector<8 x i32>, so reinterpret back and slice + vec_i32_back = reinterpret_cast(vec_f32, Float32, 32, Int32) + for chunk in range(4): # 32 / 8 = 4 chunks + store_256b(gmem_addr + chunk * 32, subvec(vec_i32_back, chunk * 8, 8)) + + # Store back to TMEM + tcgen05_st_32x32b(32, taddr, vec_i32_back) +""" + +__all__ = [ + "tcgen05_ld_32x32b", + "tcgen05_st_32x32b", + "tcgen05_cp_128x256b", + "reinterpret_cast", + "subvec", + "store_256b", + "umma_arrive", + "umma_arrive_noelect", +] + +import cutlass.cute as cute +from cutlass._mlir import ir as _ir_mod +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import nvvm as _nvvm +from cutlass._mlir.dialects import vector as _vector +from cutlass.cute.arch import elect_one +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.typing import Int32 +from cutlass.cutlass_dsl import dsl_user_op + +from cula.ops.ptx_umma_ext import Tcgen05SmemDescriptor + + +def _to_ir(val, loc=None, ip=None): + """Extract raw MLIR IR value from a CuteDSL wrapper.""" + return val.ir_value(loc=loc, ip=ip) if hasattr(val, "ir_value") else val + + +# --------------------------------------------------------------------------- +# tcgen05.ld.sync.aligned.32x32b.xN.b32 (via nvvm.tcgen05.ld) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05_ld_32x32b(num: int, taddr: int): + """Load *num* × 32-bit values from TMEM → an opaque ``vector``. + + ``num`` must be a **compile-time constant** in {2, 4, 8, 16, 32, 64, 128}. + Returns a single opaque MLIR vector value (``vector``). + + Use :func:`reinterpret_cast` to reinterpret the element type (zero-cost), + and :func:`subvec` to slice a contiguous sub-vector. + + Parameters + ---------- + num : int + Number of 32-bit registers to load. Must be a compile-time constant. + taddr : int + TMEM address (bits [31:16] = lane, bits [15:0] = column). + """ + + @dsl_user_op + def _do(addr_val, *, loc=None, ip=None): + i32_ty = _ir_mod.IntegerType.get_signless(32) + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + vec_i32_ty = _ir_mod.VectorType.get([num], i32_ty) + return _nvvm.tcgen05_ld( + res=vec_i32_ty, + shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, + num=num, + tmem_addr=tmem_ptr, + loc=loc, + ip=ip, + ) + + return _do(Int32(taddr)) + + +# --------------------------------------------------------------------------- +# tcgen05.st.sync.aligned.32x32b.xN.b32 (via nvvm.tcgen05.st) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05_st_32x32b(num: int, taddr: int, vec): + """Store *num* × 32-bit values from an opaque vector → TMEM. + + ``num`` must be a **compile-time constant** in {2, 4, 8, 16, 32, 64, 128}. + + Parameters + ---------- + num : int + Number of 32-bit registers to store. Must be a compile-time constant. + taddr : int + TMEM address (bits [31:16] = lane, bits [15:0] = column). + vec : opaque vector + An opaque ``vector`` value (from :func:`tcgen05_ld_32x32b` + or :func:`reinterpret_cast`). + """ + + @dsl_user_op + def _do(addr_val, vec_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + _nvvm.tcgen05_st( + shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, + num=num, + tmem_addr=tmem_ptr, + r=_to_ir(vec_val, loc, ip), + loc=loc, + ip=ip, + ) + + _do(Int32(taddr), vec) + + +# --------------------------------------------------------------------------- +# reinterpret_cast (zero-cost vector.bitcast) +# --------------------------------------------------------------------------- + + +@cute.jit +def reinterpret_cast(vec, src_type, src_num, tgt_type): + """Zero-cost reinterpret of a vector's element type (single ``vector.bitcast``). + + Analogous to C++ ``reinterpret_cast``: no instructions emitted, just + re-labels the bits. The total bit-width is preserved: + ``src_num * src_type.width == tgt_num * tgt_type.width``. + + Parameters + ---------- + vec : opaque vector + Source vector (e.g. ``vector`` from :func:`tcgen05_ld_32x32b`). + src_type : CuTeDSL type + Element type of *vec* (e.g. ``Int32``). + src_num : int + Number of elements in *vec* (compile-time constant). + tgt_type : CuTeDSL type + Desired element type (e.g. ``Float32``, ``BFloat16``, ``Float16``). + + Returns + ------- + opaque vector + ``vector`` where ``M = src_num * src_type.width // tgt_type.width``. + + Examples + -------- + :: + + vec_i32 = tcgen05_ld_32x32b(8, taddr) # vector<8 x i32> + vec_f32 = reinterpret_cast(vec_i32, Int32, 8, Float32) # vector<8 x f32> + vec_bf16 = reinterpret_cast(vec_i32, Int32, 8, BFloat16) # vector<16 x bf16> + vec_back = reinterpret_cast(vec_bf16, BFloat16, 16, Int32) # vector<8 x i32> + """ + tgt_num = src_num * src_type.width // tgt_type.width + + @dsl_user_op + def _do(v, *, loc=None, ip=None): + tgt_vec_ty = _ir_mod.VectorType.get([tgt_num], tgt_type.mlir_type) + return _vector.bitcast(tgt_vec_ty, _to_ir(v, loc, ip), loc=loc, ip=ip) + + return _do(vec) + + +# --------------------------------------------------------------------------- +# subvec (extract a contiguous sub-vector) +# --------------------------------------------------------------------------- + + +@cute.jit +def subvec(vec, offset, size): + """Extract a contiguous sub-vector (``vector.extract_strided_slice``). + + Parameters + ---------- + vec : opaque vector + Source vector. + offset : int + Starting element index (compile-time constant). + size : int + Number of elements to extract (compile-time constant). + + Returns + ------- + opaque vector + ``vector``. + """ + + @dsl_user_op + def _do(v, *, loc=None, ip=None): + ir_v = _to_ir(v, loc, ip) + elem_ty = _ir_mod.VectorType(ir_v.type).element_type + res_ty = _ir_mod.VectorType.get([size], elem_ty) + return _vector.extract_strided_slice( + res_ty, + ir_v, + offsets=[offset], + sizes=[size], + strides=[1], + loc=loc, + ip=ip, + ) + + return _do(vec) + + +# --------------------------------------------------------------------------- +# st.global.L1::no_allocate.v8.f32 (256-bit direct R2G store) +# --------------------------------------------------------------------------- + +_STORE_256B_ASM = "st.global.L1::no_allocate.v8.f32 [$0], {$1, $2, $3, $4, $5, $6, $7, $8};" +_STORE_256B_CONSTRAINTS = "l,r,r,r,r,r,r,r,r" + + +@cute.jit +def store_256b(gmem_ptr, vec): + """Store 256 bits (8 × 32-bit) to global memory, bypassing L1 allocation. + + Issues ``st.global.L1::no_allocate.v8.f32`` with ``"r"`` (integer register) + constraints — type-agnostic, just like C++ ``reinterpret_cast``. + + Parameters + ---------- + gmem_ptr : pointer + Global-memory destination address (must be 32-byte aligned). + vec : opaque vector + A ``vector<8 x i32>`` (use :func:`subvec` to slice from a larger vector). + """ + + @dsl_user_op + def _do(addr, v, *, loc=None, ip=None): + i32_ty = _ir_mod.IntegerType.get_signless(32) + ir_v = _to_ir(v, loc, ip) + elems = [ + _vector.extractelement( + ir_v, + position=_arith.constant(i32_ty, i, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + for i in range(8) + ] + operands = [_to_ir(addr, loc, ip)] + elems + llvm.inline_asm( + _ir_mod.Type.parse("!llvm.void"), + operands, + _STORE_256B_ASM, + _STORE_256B_CONSTRAINTS, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + _do(gmem_ptr, vec) + + +# --------------------------------------------------------------------------- +# tcgen05.cp.cta_group::1.128x256b (via nvvm.tcgen05.cp) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05_cp_128x256b(taddr: int, smem_desc: Tcgen05SmemDescriptor): + """Async copy SMEM → TMEM with shape ``128x256b`` (``cta_group::1``). + + Issues ``tcgen05.cp.cta_group::1.128x256b [taddr], s-desc;`` + via the native ``nvvm.tcgen05.cp`` MLIR op. + + The instruction copies a 128-row × 256-bit tile from shared memory + (described by *smem_desc*) into Tensor Memory at *taddr*. The copy + is **asynchronous** — use ``tcgen05.commit`` + ``mbarrier.wait`` to + synchronize. + + PTX reference + ------------- + tcgen05.cp.cta_group::1.128x256b [taddr], s-desc; + + Parameters + ---------- + taddr : int + TMEM destination address (uint32, passed as ``!llvm.ptr<6>``). + smem_desc : Tcgen05SmemDescriptor + 64-bit SMEM matrix descriptor (same format as ``tcgen05.mma`` + descriptors — see ``Tcgen05SmemDescriptor``). + """ + + @dsl_user_op + def _do(addr_val, desc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + _nvvm.tcgen05_cp( + shape=_nvvm.Tcgen05CpShape.SHAPE_128x256b, + taddr=tmem_ptr, + smem_desc=_to_ir(desc_val, loc, ip), + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + loc=loc, + ip=ip, + ) + + _do(Int32(taddr), smem_desc.desc_i64[0]) + + +@cute.jit +def tcgen05_cp_128x128b(taddr: int, smem_desc: Tcgen05SmemDescriptor): + """Async copy SMEM → TMEM with shape ``128x128b`` (``cta_group::1``). + + Issues ``tcgen05.cp.cta_group::1.128x128b [taddr], s-desc;`` + via the native ``nvvm.tcgen05.cp`` MLIR op. + + The instruction copies a 128-row × 128-bit tile from shared memory + (described by *smem_desc*) into Tensor Memory at *taddr*. The copy + is **asynchronous** — use ``tcgen05.commit`` + ``mbarrier.wait`` to + synchronize. + + PTX reference + ------------- + tcgen05.cp.cta_group::1.128x128b [taddr], s-desc; + + Parameters + ---------- + taddr : int + TMEM destination address (uint32, passed as ``!llvm.ptr<6>``). + smem_desc : Tcgen05SmemDescriptor + 64-bit SMEM matrix descriptor (same format as ``tcgen05.mma`` + descriptors — see ``Tcgen05SmemDescriptor``). + """ + + @dsl_user_op + def _do(addr_val, desc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + _nvvm.tcgen05_cp( + shape=_nvvm.Tcgen05CpShape.SHAPE_128x128b, + taddr=tmem_ptr, + smem_desc=_to_ir(desc_val, loc, ip), + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + loc=loc, + ip=ip, + ) + + _do(Int32(taddr), smem_desc.desc_i64[0]) + + +@cute.jit +def tcgen05_fence_before(): + """tcgen05.fence::before_thread_sync — non-blocking ordering fence.""" + _nvvm.tcgen05_fence(kind=_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) + + +@cute.jit +def tcgen05_fence_after(): + """tcgen05.fence::after_thread_sync — non-blocking ordering fence.""" + _nvvm.tcgen05_fence(kind=_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC) + + +@cute.jit +def umma_arrive(mbar_ptr: cute.Pointer): + """tcgen05.commit.cta_group::1.mbarrier::arrive::one — signal MMA done.""" + with elect_one(): + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + + +@cute.jit +def umma_arrive_noelect(mbar_ptr: cute.Pointer): + """tcgen05.commit.cta_group::1.mbarrier::arrive::one — signal MMA done.""" + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) diff --git a/cula/ops/ptx_umma_ext.py b/cula/ops/ptx_umma_ext.py new file mode 100644 index 00000000..2e8caeea --- /dev/null +++ b/cula/ops/ptx_umma_ext.py @@ -0,0 +1,961 @@ +# Copyright (c) 2025 ANTGROUP. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CuteDSL UMMA extension wrappers for SM100 (Blackwell) ``tcgen05.mma``. + +CuteDSL's high-level ``cute.gemm()`` / ``make_tiled_mma()`` API does not +expose all ``tcgen05.mma`` instruction variants. This module provides +low-level wrappers for the two categories currently needed: + +1. **Masked MMA** – SS and TS forms with the 128-bit ``disable-output-lane`` + mask operand (``{m0, m1, m2, m3}``). Implemented via the native + ``nvvm.tcgen05_mma`` MLIR op with its ``write_disable_mask`` parameter + (``vector<4xi32>``). + +2. **Weight-stationary (WS) MMA** – ``tcgen05.mma.ws`` SS / TS forms for + both ``kind::tf32`` and ``kind::f16``. Implemented via + ``llvm.inline_asm``. + +---------------------------------------------------------------------- +PTX instruction forms +---------------------------------------------------------------------- +SS (SMEM A, SMEM B): + tcgen05.mma.cta_group::1.kind::tf32 [tmem_c], desc_a, desc_b, + desc_val, {m0,m1,m2,m3}, p; + +TS (TMEM A, SMEM B): + tcgen05.mma.cta_group::1.kind::tf32 [tmem_c], [tmem_a], desc_b, + desc_val, {m0,m1,m2,m3}, p; + +WS_SS (weight-stationary, SMEM A, SMEM B): + tcgen05.mma.ws.cta_group::1.kind::tf32 [tmem_c], desc_a, desc_b, + desc_val, p; + tcgen05.mma.ws.cta_group::1.kind::f16 [tmem_c], desc_a, desc_b, + desc_val, p; + +WS_TS (weight-stationary, TMEM A, SMEM B): + tcgen05.mma.ws.cta_group::1.kind::tf32 [tmem_c], [tmem_a], desc_b, + desc_val, p; + tcgen05.mma.ws.cta_group::1.kind::f16 [tmem_c], [tmem_a], desc_b, + desc_val, p; + +---------------------------------------------------------------------- +Disable-output-lane mask layout (4 × uint32 = 128 bits) +---------------------------------------------------------------------- +Each uint32 covers 32 M-dimension rows (8 rows × 4 elements per group). + 0x00000000 → group is ACTIVE (output written) + 0xFFFFFFFF → group is DISABLED (output suppressed) + +Predefined SS mask constants (SMEM A variants): + SS_NO_MASK = (0, 0, 0, 0) all rows active + SS_MASK0 = (0, 0xFF…, 0, 0xFF…) odd groups disabled + SS_MASK1 = (0xFF…, 0, 0xFF…, 0) even groups disabled + SS_MASK2 = (0xFF…, 0xFF…, 0, 0xFF…) group 2 only active + SS_MASK3 = (0xFF…, 0xFF…, 0xFF…, 0) group 3 only active + +Predefined TS mask constants (TMEM A variants): + TS_NO_MASK = (0, 0, 0, 0) all rows active + TS_MASK0 = (0, 0xFF…, 0xFF…, 0xFF…) group 0 only active + TS_MASK1 = (0xFF…, 0, 0xFF…, 0xFF…) group 1 only active + TS_MASK2 = (0xFF…, 0xFF…, 0, 0xFF…) group 2 only active + TS_MASK3 = (0xFF…, 0xFF…, 0xFF…, 0) group 3 only active + TS_MASK02 = (0, 0xFF…, 0, 0xFF…) groups 0,2 only active + TS_MASK13 = (0xFF…, 0, 0xFF…, 0) groups 1,3 only active + +Public API (all decorated with @cute.jit) +---------------------------------------------------------------------- +Descriptor helpers (call inside @cute.jit): + Tcgen05SmemDescriptor — 64-bit SMEM descriptor object + initialize_tcgen05_descriptor — fill descriptor bitfields + +Low-level primitives (pass mask words explicitly): + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, + mask0, mask1, mask2, mask3) + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, + mask0, mask1, mask2, mask3) + tcgen05mma_ws_ss_tf32(desc_a, desc_b, tmem_c, desc_val, scale_out) + tcgen05mma_ws_ts_tf32(tmem_a, desc_b, tmem_c, desc_val, scale_out) + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, desc_val, scale_out) + tcgen05mma_ws_ts_f16(tmem_a, desc_b, tmem_c, desc_val, scale_out) + +Named convenience wrappers (pre-set masks, pass only MMA operands): + tcgen05mma_ss_no_mask / tcgen05mma_ss_mask0 / …mask1 / …mask2 / …mask3 + tcgen05mma_ts_no_mask / tcgen05mma_ts_mask0 / …mask1 / …mask2 / …mask3 + tcgen05mma_ts_mask02 / tcgen05mma_ts_mask13 +""" + +__all__ = [ + # descriptor helpers + "Tcgen05SmemDescriptor", + "initialize_tcgen05_descriptor", + # low-level primitives + "tcgen05mma_ss", + "tcgen05mma_ts", + "tcgen05mma_ws_ss_tf32", + "tcgen05mma_ws_ts_tf32", + "tcgen05mma_ws_ss_f16", + "tcgen05mma_ws_ts_f16", + # SS named wrappers + "tcgen05mma_ss_no_mask", + "tcgen05mma_ss_mask0", + "tcgen05mma_ss_mask1", + "tcgen05mma_ss_mask2", + "tcgen05mma_ss_mask3", + # TS named wrappers + "tcgen05mma_ts_no_mask", + "tcgen05mma_ts_mask0", + "tcgen05mma_ts_mask1", + "tcgen05mma_ts_mask2", + "tcgen05mma_ts_mask3", + "tcgen05mma_ts_mask02", + "tcgen05mma_ts_mask13", + # collector enums (re-exported for convenience) + "CollectorBBuffer", + "CollectorOp", +] + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import nvvm as _nvvm +from cutlass.cutlass_dsl import dsl_user_op + +# Re-export collector enums for caller convenience. +CollectorBBuffer = _nvvm.Tcgen05MMACollectorBBuffer +CollectorOp = _nvvm.Tcgen05MMACollectorOp + +# --------------------------------------------------------------------------- +# Mask constants (4 × uint32). 0 = ACTIVE, 0xFFFFFFFF = DISABLED. +# --------------------------------------------------------------------------- +_ALL_ACTIVE = 0x00000000 +_ALL_OFF = 0xFFFFFFFF + +# SS masks (SMEM A, SMEM B) +SS_NO_MASK = (_ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE) +SS_MASK0 = (_ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {0,F,0,F} +SS_MASK1 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE) # {F,0,F,0} +SS_MASK2 = (_ALL_OFF, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {F,F,0,F} +SS_MASK3 = (_ALL_OFF, _ALL_OFF, _ALL_OFF, _ALL_ACTIVE) # {F,F,F,0} + +# TS masks (TMEM A, SMEM B) +TS_NO_MASK = (_ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE) +TS_MASK0 = (_ALL_ACTIVE, _ALL_OFF, _ALL_OFF, _ALL_OFF) # {0,F,F,F} +TS_MASK1 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_OFF) # {F,0,F,F} +TS_MASK2 = (_ALL_OFF, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {F,F,0,F} +TS_MASK3 = (_ALL_OFF, _ALL_OFF, _ALL_OFF, _ALL_ACTIVE) # {F,F,F,0} +TS_MASK02 = (_ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {0,F,0,F} +TS_MASK13 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE) # {F,0,F,0} + + +# --------------------------------------------------------------------------- +# Tcgen05SmemDescriptor — 64-bit SMEM descriptor stored as 2×Int32 +# --------------------------------------------------------------------------- + + +class Tcgen05SmemDescriptor: + """64-bit shared-memory descriptor for tcgen05 MMA (Blackwell / SM100). + + The descriptor encodes SMEM base address, leading/stride byte offsets, + swizzle mode, and other fields required by the ``tcgen05.mma`` PTX + instruction to locate a matrix tile in shared memory. + + 64-bit layout (PTX ISA Table 40):: + + Bit 63 Bit 0 + ┌──────────┬────────┬─────┬──────────┬────┬──────────┬──────┬──────────────┐ + │ 63 61 │ 60 53 │ 52 │ 51 49 │ 48 │ 45 32 │31 30 │ 29 16│15 14│ 13 0│ + │layout_typ│ reservd│l_abs│base_offst│ 46 │ SBO │ rsvd │ LBO │rsvd │start_adr│ + │ (3 bit) │ (8 bit)│(1b) │ (3 bit) │=0b001│(14 bit)│(2 b) │(14 bit)│(2b) │(14 bit) │ + └──────────┴────────┴─────┴──────────┴────┴──────────┴──────┴────────┴─────┴─────────┘ + + Field descriptions: + + - **start_address** [bits 0-13]: SMEM base pointer, encoded as + ``smem_ptr >> 4`` (16-byte aligned). The hardware reconstructs the + full address as ``encoded_value << 4``. + + - **LBO** (Leading Byte Offset) [bits 16-29]: distance in bytes between + consecutive elements along the leading dimension, encoded as + ``lbo_bytes >> 4``. When ``lbo_mode=1`` this is an absolute byte + address rather than a relative offset. + + - **SBO** (Stride Byte Offset) [bits 32-45]: distance in bytes between + consecutive elements along the stride dimension, encoded as + ``sbo_bytes >> 4``. + + - **version** [bits 46-48]: fixed constant ``0b001`` (= 1). + + - **base_offset** [bits 49-51]: 3-bit alignment correction when the + SMEM tile does not start at a natural swizzle-pattern boundary + (1024B for 128B swizzle, 512B for 64B, 256B for 32B). + Computed as ``(start_addr >> 7) & 0x7``. Usually 0. + + - **lbo_mode** (leading_abs) [bit 52]: 0 → LBO is a relative byte + offset; 1 → LBO is an absolute byte address. + + - **layout_type** (swizzle_mode) [bits 61-63]: + - 0 = SWIZZLE_NONE + - 1 = SWIZZLE_128B_BASE32B (128-byte pattern, 32-byte atom) + - 2 = SWIZZLE_128B (128-byte pattern) + - 4 = SWIZZLE_64B (64-byte pattern) + - 6 = SWIZZLE_32B (32-byte pattern) + + Storage: two Int32 registers (desc[0] = low 32 bits, desc[1] = high 32 + bits), recast to a single Int64 for the PTX ``l``-constraint operand. + + Usage inside a @cute.jit kernel:: + + desc = Tcgen05SmemDescriptor() + initialize_tcgen05_descriptor(desc, smem_ptr, lbo, sbo, 0, True, swizzle) + """ + + def __init__(self, desc_64: cute.Int64 = None): + # desc[0]: low 32 bits → start_address[0:14] | LBO[16:30] + # desc[1]: high 32 bits → SBO[0:14] | version[14:16] | base_offset[17:20] + # | lbo_mode[20] | layout_type[29:32] + self.desc = cute.make_rmem_tensor((2,), dtype=cutlass.Int32) + # Alias the 2×i32 as 1×i64 for PTX "l" constraint (64-bit operand) + self.desc_i64 = cute.make_tensor(cute.recast_ptr(self.desc.iterator, dtype=cute.Int64), (1,)) + if desc_64 is not None: + self.desc_i64[0] = desc_64 + + def __add__(self, byte_offset): + """Return a new descriptor offset by ``byte_offset`` bytes. + + Only the start_address field (bits 0-13 of desc[0]) is modified. + Since it is stored in 16-byte units, we add ``byte_offset >> 4``. + All other fields (LBO, SBO, swizzle, etc.) are copied unchanged. + """ + res = cute.make_rmem_tensor((2,), dtype=cutlass.Int32) + res_i64 = cute.make_tensor(cute.recast_ptr(res.iterator, dtype=cute.Int64), (1,)) + res[0] = self.desc[0] + (byte_offset >> 4) # adjust start_address + res[1] = self.desc[1] # high word unchanged + return Tcgen05SmemDescriptor(res_i64[0]) + + +# --------------------------------------------------------------------------- +# initialize_tcgen05_descriptor +# --------------------------------------------------------------------------- + + +def initialize_tcgen05_descriptor( + desc, + start_address, + leading_byte_offset, + stride_byte_offset, + base_offset, + leading_abs, + swizzle_mode, +): + """Pack SMEM descriptor bitfields into *desc* (a Tcgen05SmemDescriptor). + + Constructs the 64-bit descriptor in two 32-bit halves (desc[0] and desc[1]). + All address/offset fields must be pre-divided by 16 (``>> 4``) before + passing, because the hardware stores them in 16-byte granularity. + + Low 32 bits — desc[0]:: + + ┌────────────────┬──────┬──────────────────┐ + │ bits 29…16 │15…14 │ bits 13…0 │ + │ LBO (14 bits) │ rsvd │ start_addr >> 4 │ + └────────────────┴──────┴──────────────────┘ + + - [0:14) start_address >> 4 — SMEM tile base pointer in 16B units. + - [14:16) reserved (0). + - [16:30) leading_byte_offset — LBO in 16B units (caller passes >> 4). + + High 32 bits — desc[1]:: + + ┌────────┬────────┬─────┬──────────┬────────┬──────────────────┐ + │ 31…29 │ 28…21 │ 20 │ 19…17 │ 16…14 │ bits 13…0 │ + │ layout │ rsvd │l_abs│base_off │version │ SBO (14 bits) │ + │ (3 bit)│ (8 bit)│(1b) │ (3 bit) │=0b001 │ │ + └────────┴────────┴─────┴──────────┴────────┴──────────────────┘ + + - [0:14) stride_byte_offset — SBO in 16B units (caller passes >> 4). + - [14:16) version = 1 (fixed constant 0b001, only bit 14 set). + - [17:20) base_offset & 0x7 — swizzle alignment correction. + Typically 0. Non-zero when the tile doesn't start at + the natural swizzle boundary (1024B/512B/256B). + - [20:21) lbo_mode — 0 = LBO is relative offset, 1 = absolute address. + - [29:32) layout_type (swizzle_mode & 0x7): + 0 = SWIZZLE_NONE + 1 = SWIZZLE_128B_BASE32B (Swizzle<2,5,2>) + 2 = SWIZZLE_128B (Swizzle<3,4,3>) + 4 = SWIZZLE_64B (Swizzle<2,4,3>) + 6 = SWIZZLE_32B (Swizzle<1,4,3>) + + Args: + desc: Tcgen05SmemDescriptor to fill. + start_address: CuTeDSL Pointer to the SMEM tile start. + leading_byte_offset: Leading-dimension byte offset, already >> 4. + stride_byte_offset: Stride byte offset, already >> 4. + base_offset: Swizzle alignment correction (raw int, bits 17-19). + leading_abs: Bool — True → LBO is absolute address. + swizzle_mode: Swizzle layout_type integer (bits 29-31). + """ + # Encode start_address: take SMEM pointer, shift right by 4 to get 16B units + ptr_val = start_address.toint() >> 4 + + # --- Low 32 bits (desc[0]) --- + # bits [0:14) = start_address >> 4 + # bits [16:30) = leading_byte_offset (already in 16B units) + desc.desc[0] = cutlass.Int32(ptr_val) | cutlass.Int32(cutlass.Int32(leading_byte_offset) << 16) + + # --- High 32 bits (desc[1]) --- + # bits [0:14) = stride_byte_offset (already in 16B units) + # bit [14] = version = 1 (fixed) + # bits [17:20) = base_offset & 0x7 (swizzle alignment correction) + # bit [20] = lbo_mode (0=relative, 1=absolute) + # bits [29:32) = layout_type (swizzle mode) + desc.desc[1] = ( + cutlass.Int32(stride_byte_offset) + | cutlass.Int32(1 << 14) # version = 1 + | cutlass.Int32(cutlass.Int32(base_offset & 0x7) << 17) + | cutlass.Int32(cutlass.Int32(int(leading_abs)) << 20) + | cutlass.Int32(cutlass.Int32(swizzle_mode & 0x7) << 29) + ) + + +# --------------------------------------------------------------------------- +# Internal helper +# --------------------------------------------------------------------------- + + +def _ir(val, loc=None, ip=None): + """Extract raw MLIR IR value from a CuTeDSL wrapper.""" + return val.ir_value(loc=loc, ip=ip) if hasattr(val, "ir_value") else val + + +# =========================================================================== +# Low-level primitives +# =========================================================================== + +# --------------------------------------------------------------------------- +# tcgen05mma_ss — SMEM A, SMEM B (non-warp-specialised) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ss( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + mask0: int, + mask1: int, + mask2: int, + mask3: int, +): + """Issue ``tcgen05.mma.cta_group::1.kind::tf32`` with SMEM operands. + + ``mask{0-3}`` are the four uint32 words of the 128-bit + ``disable-output-lane`` mask (0=active, 0xFFFFFFFF=disabled). + + Caller must ensure single-thread execution (e.g. via ``elect_one``); + no internal ``elect.sync`` is performed. + + Args: + desc_a: 64-bit SMEM descriptor for matrix A. + desc_b: 64-bit SMEM descriptor for matrix B. + tmem_c: TMEM base address (uint32) for accumulators C/D. + desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). + scale_out: 1 → accumulate into C, 0 → overwrite C (clear accumulators). + mask0-3: Four uint32 words of the disable-output-lane mask. + """ + + @dsl_user_op + def _do(c_val, da_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i32_ty = ir.IntegerType.get_signless(32) + i1_ty = ir.IntegerType.get_signless(1) + vec4i32_ty = ir.VectorType.get([4], i32_ty) + + c_ir = _ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + da_ir = _ir(da_val, loc, ip) # i64 SMEM descriptor + db_ir = _ir(db_val, loc, ip) # i64 SMEM descriptor + dv_ir = _ir(dv_val, loc, ip) + sc_ir = _ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + m0_ir = _ir(m0_val, loc, ip) + m1_ir = _ir(m1_val, loc, ip) + m2_ir = _ir(m2_val, loc, ip) + m3_ir = _ir(m3_val, loc, ip) + + undef = llvm.mlir_undef(vec4i32_ty, loc=loc, ip=ip) + idx0 = _arith.constant(i32_ty, 0, loc=loc, ip=ip) + idx1 = _arith.constant(i32_ty, 1, loc=loc, ip=ip) + idx2 = _arith.constant(i32_ty, 2, loc=loc, ip=ip) + idx3 = _arith.constant(i32_ty, 3, loc=loc, ip=ip) + v = llvm.InsertElementOp(undef, m0_ir, idx0, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m1_ir, idx1, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) + mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) + + _nvvm.tcgen05_mma( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + d=d_ptr, + a=da_ir, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + write_disable_mask=mask, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + desc_a.desc_i64[0], + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + cutlass.Int32(mask0), + cutlass.Int32(mask1), + cutlass.Int32(mask2), + cutlass.Int32(mask3), + ) + + +# --------------------------------------------------------------------------- +# tcgen05mma_ts — TMEM A, SMEM B (non-warp-specialised) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ts( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + mask0: int, + mask1: int, + mask2: int, + mask3: int, +): + """Issue ``tcgen05.mma.cta_group::1.kind::tf32`` with TMEM A operand. + + Matrix A is read from TMEM via indirect addressing ``[tmem_a]``. + Matrix B is read from SMEM via descriptor. + Caller must ensure single-thread execution (e.g. via ``elect_one``). + + Args: + tmem_a: TMEM base address (uint32) for matrix A. + desc_b: 64-bit SMEM descriptor for matrix B. + tmem_c: TMEM base address (uint32) for accumulators C/D. + desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). + scale_out: 1 → accumulate into C, 0 → overwrite C. + mask0-3: Four uint32 words of the disable-output-lane mask. + """ + + @dsl_user_op + def _do(c_val, a_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i32_ty = ir.IntegerType.get_signless(32) + i1_ty = ir.IntegerType.get_signless(1) + vec4i32_ty = ir.VectorType.get([4], i32_ty) + + c_ir = _ir(c_val, loc, ip) + a_ir = _ir(a_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) + b_ir = _ir(db_val, loc, ip) + dv_ir = _ir(dv_val, loc, ip) + sc_ir = _ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + m0_ir = _ir(m0_val, loc, ip) + m1_ir = _ir(m1_val, loc, ip) + m2_ir = _ir(m2_val, loc, ip) + m3_ir = _ir(m3_val, loc, ip) + + undef = llvm.mlir_undef(vec4i32_ty, loc=loc, ip=ip) + idx0 = _arith.constant(i32_ty, 0, loc=loc, ip=ip) + idx1 = _arith.constant(i32_ty, 1, loc=loc, ip=ip) + idx2 = _arith.constant(i32_ty, 2, loc=loc, ip=ip) + idx3 = _arith.constant(i32_ty, 3, loc=loc, ip=ip) + v = llvm.InsertElementOp(undef, m0_ir, idx0, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m1_ir, idx1, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) + mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) + + _nvvm.tcgen05_mma( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + d=d_ptr, + a=a_ptr, + b=b_ir, + idesc=dv_ir, + enable_input_d=enable_d, + write_disable_mask=mask, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + cutlass.Int32(tmem_a), + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + cutlass.Int32(mask0), + cutlass.Int32(mask1), + cutlass.Int32(mask2), + cutlass.Int32(mask3), + ) + + +# --------------------------------------------------------------------------- +# tcgen05mma_ws_ss_tf32 — weight-stationary, SMEM A, SMEM B, kind::tf32 +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ws_ss_tf32( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::tf32`` (weight-stationary form). + + This variant does NOT take a ``disable-output-lane`` mask; the + optional ``zero-column-mask-desc`` operand is omitted. + + Args: + desc_a: 64-bit SMEM descriptor for matrix A. + desc_b: 64-bit SMEM descriptor for matrix B. + tmem_c: TMEM base address (uint32) for accumulators C/D. + desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). + scale_out: 1 → accumulate, 0 → overwrite. + collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). + Defaults to None (hardware default: ``b0::discard``). + collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). + Defaults to None (hardware default: discard). + """ + + @dsl_user_op + def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + da_ir = _ir(da_val, loc, ip) + db_ir = _ir(db_val, loc, ip) + dv_ir = _ir(dv_val, loc, ip) + sc_ir = _ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + d=d_ptr, + a=da_ir, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + desc_a.desc_i64[0], + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +# --------------------------------------------------------------------------- +# tcgen05mma_ws_ss_f16 — weight-stationary, SMEM A, SMEM B, kind::f16 +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ws_ss_f16( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::f16`` (weight-stationary form). + + Same as the tf32 variant but uses ``.kind::f16`` for half-precision + input types (f16 / bf16). K dimension is 16 instead of 8. + + This variant does NOT take a ``disable-output-lane`` mask; the + optional ``zero-column-mask-desc`` operand is omitted. + + Args: + desc_a: 64-bit SMEM descriptor for matrix A. + desc_b: 64-bit SMEM descriptor for matrix B. + tmem_c: TMEM base address (uint32) for accumulators C/D. + desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). + scale_out: 1 → accumulate, 0 → overwrite. + collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). + Defaults to None (hardware default: ``b0::discard``). + collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). + Defaults to None (hardware default: discard). + """ + + @dsl_user_op + def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + da_ir = _ir(da_val, loc, ip) + db_ir = _ir(db_val, loc, ip) + dv_ir = _ir(dv_val, loc, ip) + sc_ir = _ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.F16, + d=d_ptr, + a=da_ir, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + desc_a.desc_i64[0], + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +# --------------------------------------------------------------------------- +# tcgen05mma_ws_ts_tf32 — weight-stationary, TMEM A, SMEM B, kind::tf32 +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ws_ts_tf32( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::tf32`` with TMEM A (weight-stationary). + + Matrix A is read from TMEM via indirect addressing ``[tmem_a]``. + Matrix B is read from SMEM via descriptor. + This variant does NOT take a ``disable-output-lane`` mask; the + optional ``zero-column-mask-desc`` operand is omitted. + + Args: + tmem_a: TMEM base address (uint32) for matrix A. + desc_b: 64-bit SMEM descriptor for matrix B. + tmem_c: TMEM base address (uint32) for accumulators C/D. + desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). + scale_out: 1 → accumulate, 0 → overwrite. + collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). + Defaults to None (hardware default: ``b0::discard``). + collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). + Defaults to None (hardware default: discard). + """ + + @dsl_user_op + def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + a_ir = _ir(a_val, loc, ip) + a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) + db_ir = _ir(db_val, loc, ip) + dv_ir = _ir(dv_val, loc, ip) + sc_ir = _ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + d=d_ptr, + a=a_ptr, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + cutlass.Int32(tmem_a), + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +# --------------------------------------------------------------------------- +# tcgen05mma_ws_ts_f16 — weight-stationary, TMEM A, SMEM B, kind::f16 +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ws_ts_f16( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::f16`` with TMEM A (weight-stationary). + + Same as the tf32 variant but uses ``.kind::f16`` for half-precision + input types (f16 / bf16). K dimension is 16 instead of 8. + + Matrix A is read from TMEM via indirect addressing ``[tmem_a]``. + Matrix B is read from SMEM via descriptor. + This variant does NOT take a ``disable-output-lane`` mask; the + optional ``zero-column-mask-desc`` operand is omitted. + + Args: + tmem_a: TMEM base address (uint32) for matrix A. + desc_b: 64-bit SMEM descriptor for matrix B. + tmem_c: TMEM base address (uint32) for accumulators C/D. + desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). + scale_out: 1 → accumulate, 0 → overwrite. + collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). + Defaults to None (hardware default: ``b0::discard``). + collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). + Defaults to None (hardware default: discard). + """ + + @dsl_user_op + def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + a_ir = _ir(a_val, loc, ip) + a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) + db_ir = _ir(db_val, loc, ip) + dv_ir = _ir(dv_val, loc, ip) + sc_ir = _ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.F16, + d=d_ptr, + a=a_ptr, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + cutlass.Int32(tmem_a), + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +# =========================================================================== +# Named convenience wrappers +# =========================================================================== +# These call the low-level primitives with pre-set mask constants so callers +# do not need to repeat the literal values. Signature: same as the base +# function but without the mask0-3 args. + +# --------------------------------------------------------------------------- +# SS named wrappers (SMEM A) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ss_no_mask( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA with no output-lane disable (all rows active).""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_NO_MASK[0], SS_NO_MASK[1], SS_NO_MASK[2], SS_NO_MASK[3]) + + +@cute.jit +def tcgen05mma_ss_mask0( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0, 0xF…, 0, 0xF…} — groups 0,2 active (1,3 disabled).""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK0[0], SS_MASK0[1], SS_MASK0[2], SS_MASK0[3]) + + +@cute.jit +def tcgen05mma_ss_mask1( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0xF…, 0, 0xF…, 0} — groups 1,3 active (0,2 disabled).""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK1[0], SS_MASK1[1], SS_MASK1[2], SS_MASK1[3]) + + +@cute.jit +def tcgen05mma_ss_mask2( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0xF…, 0xF…, 0, 0xF…} — group 2 only active.""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK2[0], SS_MASK2[1], SS_MASK2[2], SS_MASK2[3]) + + +@cute.jit +def tcgen05mma_ss_mask3( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0xF…, 0xF…, 0xF…, 0} — group 3 only active.""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK3[0], SS_MASK3[1], SS_MASK3[2], SS_MASK3[3]) + + +# --------------------------------------------------------------------------- +# TS named wrappers (TMEM A) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ts_no_mask( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA with no output-lane disable (all rows active).""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_NO_MASK[0], TS_NO_MASK[1], TS_NO_MASK[2], TS_NO_MASK[3]) + + +@cute.jit +def tcgen05mma_ts_mask0( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0, 0xF…, 0xF…, 0xF…} — group 0 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK0[0], TS_MASK0[1], TS_MASK0[2], TS_MASK0[3]) + + +@cute.jit +def tcgen05mma_ts_mask1( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0, 0xF…, 0xF…} — group 1 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK1[0], TS_MASK1[1], TS_MASK1[2], TS_MASK1[3]) + + +@cute.jit +def tcgen05mma_ts_mask2( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0xF…, 0, 0xF…} — group 2 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK2[0], TS_MASK2[1], TS_MASK2[2], TS_MASK2[3]) + + +@cute.jit +def tcgen05mma_ts_mask3( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0xF…, 0xF…, 0} — group 3 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK3[0], TS_MASK3[1], TS_MASK3[2], TS_MASK3[3]) + + +@cute.jit +def tcgen05mma_ts_mask02( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0, 0xF…, 0, 0xF…} — groups 0,2 active (1,3 disabled). + + Used in the KDA intra-chunk backward kernel for the QK/KG phase where + only even row-groups of the M tile contribute to the triangular region. + """ + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK02[0], TS_MASK02[1], TS_MASK02[2], TS_MASK02[3]) + + +@cute.jit +def tcgen05mma_ts_mask13( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0, 0xF…, 0} — groups 1,3 active (0,2 disabled). + + Used in the KDA intra-chunk backward kernel for the QK/KG phase where + only odd row-groups of the M tile contribute to the triangular region. + """ + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK13[0], TS_MASK13[1], TS_MASK13[2], TS_MASK13[3]) diff --git a/tests/test_ptx_umma_masked.py b/tests/test_ptx_umma_masked.py new file mode 100644 index 00000000..e7bc3196 --- /dev/null +++ b/tests/test_ptx_umma_masked.py @@ -0,0 +1,326 @@ +# Copyright (c) 2025 ANTGROUP. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Standalone CuteDSL test for ptx_umma_masked.py inline PTX MMA wrappers. + +Tests: + 1. tcgen05mma_ss_no_mask -- M=64, N=64, K=8, TF32, all rows active → matches torch.mm + 2. tcgen05mma_ss_mask0 -- groups 0,2 active (rows 0-15, 32-47), groups 1,3 disabled + 3. tcgen05mma_ss_mask1 -- groups 1,3 active (rows 16-31, 48-63), groups 0,2 disabled + +SMEM layout: + A: swizzled K-major (Swizzle<1,4,3>, SWIZZLE_32B), descriptor LBO=1 SBO=16 layout=6 + B: swizzled MN-major (Swizzle<2,5,2>, SWIZZLE_128B_BASE32B), LBO=64 SBO=32 layout=1 + Data is loaded with M-major mapping for A, direct row-major for B. + +All descriptor values are computed via make_umma_smem_desc / smem_descriptor_to_int +(proven correct in test_umma_ptx_jit.py). Wrapped in Tcgen05SmemDescriptor for API +compatibility with ptx_umma_masked.py convenience wrappers. +""" + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass.cute.arch import ( + elect_one, + mbarrier_init, + mbarrier_init_fence, + mbarrier_wait, + sync_threads, +) +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.nvgpu.tcgen05 import ( + Pack, + Repetition, + make_umma_smem_desc, + smem_descriptor_to_int, +) +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Float32, Int32, Int64, TFloat32 + +from cula.ops.ptx_umma_ext import ( + Tcgen05SmemDescriptor, + tcgen05mma_ss_mask0, + tcgen05mma_ss_mask1, + tcgen05mma_ss_no_mask, +) + +M_DIM, N_DIM, K_DIM = 64, 64, 8 +TMEM_COLS = 64 + +IDESC_M64_N64 = (4 << 24) | (8 << 17) | (1 << 16) | (2 << 10) | (2 << 7) | (1 << 4) +assert IDESC_M64_N64 == 0x4110910 + + +class _Kernel: + def __init__(self, mask_mode: str = "none"): + self.mask_mode = mask_mode + + @cute.kernel + def kernel(self, A_in: cute.Tensor, B_in: cute.Tensor, C_out: cute.Tensor): + """ + For mask_mode == "none": single-phase MMA, all rows written. + For mask_mode == "mask0"/"mask1": two-phase MMA: + Phase 1 - full no-mask MMA with A_in used as A_zero (passed by caller as zeros), + scale_out=0 → zeroes TMEM for all rows. + Phase 2 - masked MMA using B column from B_in (same B, new A from A_in second half). + + To keep the interface simple, for mask tests A_in is [A_zero (64×8) || A_real (64×8)] + concatenated to shape (128, 8). Phase 1 loads rows [0:64], phase 2 loads rows [64:128]. + """ + M, N, K = M_DIM, N_DIM, K_DIM + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + smem = utils.SmemAllocator() + tmem_hold_ptr = smem.allocate(Int32) + mbar_ptr = smem.allocate(Int64, byte_alignment=8) + + # Build tiled_mma to get correct SMEM layout + tiled_mma = sm100_utils.make_trivial_tiled_mma( + TFloat32, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (M, N), + ) + mma_tiler = (M, N, K) + + # Allocate swizzled SMEM + a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler, TFloat32, 1) + b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler, TFloat32, 1) + bufferA = smem.allocate_tensor( + element_type=TFloat32, + layout=a_smem_layout.outer, + byte_alignment=128, + swizzle=a_smem_layout.inner, + ) + bufferB = smem.allocate_tensor( + element_type=TFloat32, + layout=b_smem_layout.outer, + byte_alignment=128, + swizzle=b_smem_layout.inner, + ) + bufA_s0 = bufferA[(None, None, None, 0)] + bufB_s0 = bufferB[(None, None, None, 0)] + + if tidx == cutlass.Int32(0): + mbarrier_init(mbar_ptr, 1) + mbarrier_init_fence() + + # gA_all: flat view of input A tensor (either M*K or 2M*K elements) + # For no_mask: caller passes (M,K) → total = M*K + # For mask tests: caller passes (2M,K) → total = 2M*K; row_offset selects which half + if cutlass.const_expr(self.mask_mode != "none"): + gA_all = cute.make_tensor(A_in.iterator, cute.make_layout(2 * M_DIM * K_DIM)) + else: + gA_all = cute.make_tensor(A_in.iterator, cute.make_layout(M_DIM * K_DIM)) + + # Load B once (shared by both phases) + gB_flat = cute.make_tensor(B_in.iterator, cute.make_layout(K * N)) + for step in cutlass.range(K * N // 128, unroll_full=False): + idx = tidx + step * 128 + bufB_s0[idx] = gB_flat[idx] + sync_threads() + + # TMEM allocation + alloc_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128) + tmem = utils.TmemAllocator( + tmem_hold_ptr, + barrier_for_retrieve=alloc_bar, + allocator_warp_id=0, + ) + tmem.allocate(TMEM_COLS) + tmem.wait_for_alloc() + tmem_ptr_f32 = tmem.retrieve_ptr(Float32) + + acc_shape = tiled_mma.partition_shape_C((M, N)) + acc_shape_staged = cute.append(acc_shape, 1) + tCtAcc = cute.make_tensor(tmem_ptr_f32, tiled_mma.make_fragment_C(acc_shape_staged).layout) + tmem_col_buf = cute.make_tensor(tmem_hold_ptr, cute.make_layout(1)) + tmem_col = tmem_col_buf[0] + + # Build descriptors + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufA_s0.iterator, bufA_s0.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufB_s0.iterator, bufB_s0.layout, "mn")) + desc_a = Tcgen05SmemDescriptor(desc_a_i64) + desc_b = Tcgen05SmemDescriptor(desc_b_i64) + + if cutlass.const_expr(self.mask_mode != "none"): + # Phase 1: Load A_zero (first M rows of gA_all = all zeros), no_mask MMA → zero TMEM + for step in cutlass.range(M_DIM * K_DIM // 128, unroll_full=False): + smem_idx = tidx + step * 128 + m = smem_idx % M_DIM + k = smem_idx // M_DIM + bufA_s0[smem_idx] = gA_all[m * K_DIM + k] # row_offset=0 + sync_threads() + if warp_idx == cutlass.Int32(0): + tcgen05mma_ss_no_mask(desc_a, desc_b, tmem_col, IDESC_M64_N64, 0) + with elect_one(): + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + mbarrier_wait(mbar_ptr, 0) + sync_threads() + # Re-arm mbar for second MMA + if tidx == cutlass.Int32(0): + mbarrier_init(mbar_ptr, 1) + mbarrier_init_fence() + + # Phase 2: Load A_real (rows M..2M of gA_all = real data), masked MMA + for step in cutlass.range(M_DIM * K_DIM // 128, unroll_full=False): + smem_idx = tidx + step * 128 + m = smem_idx % M_DIM + k = smem_idx // M_DIM + bufA_s0[smem_idx] = gA_all[(M_DIM + m) * K_DIM + k] # row_offset=M + sync_threads() + if warp_idx == cutlass.Int32(0): + if cutlass.const_expr(self.mask_mode == "mask0"): + tcgen05mma_ss_mask0(desc_a, desc_b, tmem_col, IDESC_M64_N64, 0) + elif cutlass.const_expr(self.mask_mode == "mask1"): + tcgen05mma_ss_mask1(desc_a, desc_b, tmem_col, IDESC_M64_N64, 0) + with elect_one(): + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + mbarrier_wait(mbar_ptr, 0) + sync_threads() + else: + # Simple single-phase no_mask + for step in cutlass.range(M_DIM * K_DIM // 128, unroll_full=False): + smem_idx = tidx + step * 128 + m = smem_idx % M_DIM + k = smem_idx // M_DIM + bufA_s0[smem_idx] = gA_all[m * K_DIM + k] + sync_threads() + if warp_idx == cutlass.Int32(0): + tcgen05mma_ss_no_mask(desc_a, desc_b, tmem_col, IDESC_M64_N64, 0) + with elect_one(): + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + mbarrier_wait(mbar_ptr, 0) + sync_threads() + + # T2R: TMEM → RMEM + t2r_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(Repetition(8), Pack.NONE), Float32) + fake_smem = cute.make_tensor(cute.make_ptr(Float32, 0, cute.AddressSpace.smem), cute.make_layout((M, N))) + tCtAcc_flat = tCtAcc[((None, None), 0, 0, None)] + tiled_t2r = tcgen05.make_tmem_copy(t2r_atom, tCtAcc_flat[(None, None, 0)]) + thr_t2r = tiled_t2r.get_slice(tidx) + tTR_tAcc = thr_t2r.partition_S(tCtAcc_flat) + tTR_sDummy = thr_t2r.partition_D(fake_smem) + tTR_rAcc = cute.make_rmem_tensor(tTR_sDummy.shape, Float32) + + cute.copy(tiled_t2r, tTR_tAcc[(None, None, None, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # R2G: RMEM → GMEM (row-major) + gC = cute.make_tensor(C_out.iterator, cute.make_layout((M, N), stride=(N, 1))) + tTR_gC = thr_t2r.partition_D(gC) + cute.copy(cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32), tTR_rAcc, tTR_gC) + + sync_threads() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr_f32, TMEM_COLS) + + @cute.jit + def _launch(self, A: cute.Tensor, B: cute.Tensor, C: cute.Tensor, stream): + self.kernel(A, B, C).launch(grid=(1, 1, 1), block=(128, 1, 1), stream=stream) + + def run(self, A_cpu, B_cpu): + """ + For no_mask: A_cpu is (M, K). + For mask tests: A_cpu is (2M, K) where [0:M] = zeros, [M:2M] = real A. + """ + A_gpu = A_cpu.contiguous().float().cuda() + B_gpu = B_cpu.contiguous().float().cuda() + C_gpu = torch.zeros(M_DIM, N_DIM, dtype=torch.float32, device="cuda") + stream = cutlass_torch.default_stream() + self._launch(from_dlpack(A_gpu), from_dlpack(B_gpu), from_dlpack(C_gpu), stream) + torch.cuda.synchronize() + return C_gpu.cpu() + + +def test_ss_no_mask(): + print("\n=== Test 1: tcgen05mma_ss_no_mask (all rows active) ===") + torch.manual_seed(42) + A = torch.randn(M_DIM, K_DIM) + B = torch.randn(K_DIM, N_DIM) + ref = torch.mm(A, B) + got = _Kernel("none").run(A, B) + rel = (got - ref).abs().max().item() / (ref.abs().max().item() + 1e-8) + print(f" got[0,:4]={got[0, :4].tolist()}") + print(f" ref[0,:4]={ref[0, :4].tolist()}") + print(f" max_rel_err={rel:.4f}") + assert rel < 0.02, f"FAIL: rel={rel:.4f}" + print(" PASSED") + + +def _run_masked(mask_mode, A_real, B): + """ + Two-phase: phase1 = no_mask with A_zero (rows 0..M-1 of combined A = zeros), + phase2 = masked with A_real (rows M..2M-1 of combined A). + Combined A is shape (2M, K): [zeros || A_real]. + """ + A_combined = torch.cat([torch.zeros_like(A_real), A_real], dim=0) + return _Kernel(mask_mode).run(A_combined, B) + + +def test_ss_mask0(): + print("\n=== Test 2: tcgen05mma_ss_mask0 ===") + torch.manual_seed(0) + A = torch.randn(M_DIM, K_DIM) + B = torch.randn(K_DIM, N_DIM) + ref = torch.mm(A, B) # expected for active rows + got = _run_masked("mask0", A, B) + + # SS_MASK0 = (0, 0xFF..., 0, 0xFF...) → mask words 1,3 disable rows 16-31 and 48-63 + # Active: rows 0-15 and 32-47, Disabled: rows 16-31 and 48-63 + active_rows = list(range(0, 16)) + list(range(32, 48)) + masked_rows = list(range(16, 32)) + list(range(48, 64)) + + rel_active = (got[active_rows] - ref[active_rows]).abs().max().item() / (ref[active_rows].abs().max().item() + 1e-8) + zero_max = got[masked_rows].abs().max().item() + + print(f" active rows 0-15,32-47: max_rel_err={rel_active:.4f} (expect <0.02)") + print(f" masked rows 16-31,48-63: max_abs={zero_max:.4f} (expect 0.0)") + assert rel_active < 0.02, f"FAIL active rows: {rel_active:.4f}" + assert zero_max == 0.0, f"FAIL masked rows not zero: {zero_max}" + print(" PASSED") + + +def test_ss_mask1(): + print("\n=== Test 3: tcgen05mma_ss_mask1 ===") + torch.manual_seed(7) + A = torch.randn(M_DIM, K_DIM) + B = torch.randn(K_DIM, N_DIM) + ref = torch.mm(A, B) + got = _run_masked("mask1", A, B) + + # SS_MASK1 = (0xFF..., 0, 0xFF..., 0) → mask words 0,2 disable rows 0-15 and 32-47 + # Active: rows 16-31 and 48-63, Disabled: rows 0-15 and 32-47 + active_rows = list(range(16, 32)) + list(range(48, 64)) + masked_rows = list(range(0, 16)) + list(range(32, 48)) + + rel_active = (got[active_rows] - ref[active_rows]).abs().max().item() / (ref[active_rows].abs().max().item() + 1e-8) + zero_max = got[masked_rows].abs().max().item() + + print(f" active rows 16-31,48-63: max_rel_err={rel_active:.4f} (expect <0.02)") + print(f" masked rows 0-15,32-47: max_abs={zero_max:.4f} (expect 0.0)") + assert rel_active < 0.02, f"FAIL active rows: {rel_active:.4f}" + assert zero_max == 0.0, f"FAIL masked rows not zero: {zero_max}" + print(" PASSED") + + +if __name__ == "__main__": + test_ss_no_mask() + test_ss_mask0() + test_ss_mask1() + print("\n=== All tests passed! ===") diff --git a/tests/test_ptx_umma_ws.py b/tests/test_ptx_umma_ws.py new file mode 100644 index 00000000..da211a9a --- /dev/null +++ b/tests/test_ptx_umma_ws.py @@ -0,0 +1,591 @@ +# Copyright (c) 2025 ANTGROUP. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Standalone CuteDSL test for tcgen05.mma.ws (weight-stationary) inline PTX wrappers. + +Tests: + 1. tcgen05mma_ws_ss_tf32 -- WS mode, SMEM A × SMEM B → TMEM C, kind::tf32 + 2. tcgen05mma_ws_ts_tf32 -- WS mode, TMEM A × SMEM B → TMEM C, kind::tf32 + 3. tcgen05mma_ws_ss_f16 -- WS mode, SMEM A × SMEM B → TMEM C, kind::f16 + 4. tcgen05mma_ws_ts_f16 -- WS mode, TMEM A × SMEM B → TMEM C, kind::f16 + +For the WS_TS test, matrix A is first loaded into TMEM via an SS MMA (identity- +like multiplication), then used as the A operand for the WS TS MMA. To keep +things simple we use a two-TMEM-column approach: + - tmem region 0: accumulator for both phases + - tmem region 1: holds A data for TS phase (populated via R2T store) + +SMEM layout follows the same conventions as test_ptx_umma_masked.py. +""" + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass.cute.arch import ( + elect_one, + mbarrier_init, + mbarrier_init_fence, + mbarrier_wait, + sync_threads, +) +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.nvgpu.tcgen05 import ( + make_umma_smem_desc, + smem_descriptor_to_int, +) +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import BFloat16, Float32, Int32, Int64, TFloat32 + +from cula.ops.intrinsics_sm100 import ( + store_256b, + subvec, + tcgen05_ld_32x32b, +) +from cula.ops.ptx_umma_ext import ( + CollectorBBuffer, + CollectorOp, + Tcgen05SmemDescriptor, + tcgen05mma_ws_ss_f16, + tcgen05mma_ws_ss_tf32, +) + +M_DIM, N_DIM = 64, 64 +# TODO: support arbitrary K +K_DIM_TF32 = 8 # kind::tf32 → K>=8, tile size +A_K_STEP_BYTES_TF32 = M_DIM * 8 * 4 # smem offset for each K-atom in operand A +B_K_STEP_BYTES_TF32 = N_DIM * 8 * 4 # smem offset for each K-atom in operand B +K_DIM_F16 = 128 # default after sweep +# NOTE: per-K-atom byte offsets are derived from the SMEM layout at runtime +# (see _WsSsF16Kernel) so K_DIM_F16 can be any multiple of 16. The layout's +# k_iter mode becomes hierarchical at K≥128 (e.g. (4, K/64):(16, 4096) for A), +# which the layout-based offset computation handles transparently. + +# Instruction descriptor for M=64, N=64, TF32, dense, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=8 at [17:22], TransposeB at [16], +# btype=tf32(2) at [10:12], atype=tf32(2) at [7:9], dtype=f32(1) at [4:5] +IDESC_TF32_M64_N64 = (4 << 24) | (8 << 17) | (1 << 16) | (2 << 10) | (2 << 7) | (1 << 4) +assert IDESC_TF32_M64_N64 == 0x4110910 + +# Instruction descriptor for M=64, N=64, BF16, dense, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=8 at [17:22], TransposeB at [16], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N64 = (4 << 24) | (8 << 17) | (1 << 16) | (1 << 10) | (1 << 7) | (1 << 4) +assert IDESC_F16_M64_N64 == 0x4110490 + +# Instruction descriptor for M=64, N=128, BF16, dense, TransposeB=1 +# Bits: M>>4=4 at [24:28], N>>3=16 at [17:22], TransposeB at [16], +# btype=bf16(1) at [10:12], atype=bf16(1) at [7:9], dtype=f32(1) at [4:5] +IDESC_F16_M64_N128 = (4 << 24) | (16 << 17) | (1 << 16) | (1 << 10) | (1 << 7) | (1 << 4) +assert IDESC_F16_M64_N128 == 0x4210490 + + +# ===================================================================== +# Test 1: tcgen05mma_ws_ss_tf32 (weight-stationary, SMEM A, SMEM B, tf32) +# ===================================================================== + + +class _WsSsTf32Kernel: + @cute.kernel + def kernel(self, A_in: cute.Tensor, B_in: cute.Tensor, C_out: cute.Tensor): + M, N, K = M_DIM, N_DIM, K_DIM_TF32 + ACC_NUM_COLS = N // 2 + NUM_COLS = ACC_NUM_COLS + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + smem = utils.SmemAllocator() + tmem_hold_ptr = smem.allocate(Int32) + mbar_ptr = smem.allocate(Int64, byte_alignment=8) + + # --- SMEM layouts via sm100_utils (handles swizzle correctly for TF32) --- + # NOTE: we use non-ws mode TiledMMA for creating smem layout in a easy way, + # because smem layouts of ws mode and non-ws mode are the same + non_ws_tiled_mma = sm100_utils.make_trivial_tiled_mma( + TFloat32, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (M, N), + ) + mma_tiler = (M, N, K) + a_smem_layout = sm100_utils.make_smem_layout_a(non_ws_tiled_mma, mma_tiler, TFloat32, 1) + b_smem_layout = sm100_utils.make_smem_layout_b(non_ws_tiled_mma, mma_tiler, TFloat32, 1) + bufferA = smem.allocate_tensor( + element_type=TFloat32, + layout=a_smem_layout.outer, + byte_alignment=128, + swizzle=a_smem_layout.inner, + ) + bufferB = smem.allocate_tensor( + element_type=TFloat32, + layout=b_smem_layout.outer, + byte_alignment=128, + swizzle=b_smem_layout.inner, + ) + bufA_s0 = bufferA[(None, None, None, 0)] + bufB_s0 = bufferB[(None, None, None, 0)] + + if tidx == cutlass.Int32(0): + mbarrier_init(mbar_ptr, 1) + mbarrier_init_fence() + + # Load A (row-major input → K-major swizzled SMEM) and B + gA_flat = cute.make_tensor(A_in.iterator, cute.make_layout(M * K)) + gB_flat = cute.make_tensor(B_in.iterator, cute.make_layout(K * N)) + + for step in cutlass.range(M * K // 128, unroll_full=False): + smem_idx = tidx + step * 128 + m = smem_idx % M + k = smem_idx // M + bufA_s0[smem_idx] = gA_flat[m * K + k] + for step in cutlass.range(K * N // 128, unroll_full=False): + idx = tidx + step * 128 + bufB_s0[idx] = gB_flat[idx] + sync_threads() + + # --- TMEM allocation --- + alloc_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128) + tmem = utils.TmemAllocator( + tmem_hold_ptr, + barrier_for_retrieve=alloc_bar, + allocator_warp_id=0, + ) + tmem.allocate(NUM_COLS) + tmem.wait_for_alloc() + tmem_ptr_f32 = tmem.retrieve_ptr(Float32) + + tmem_col_buf = cute.make_tensor(tmem_hold_ptr, cute.make_layout(1)) + tmem_col = tmem_col_buf[0] + + # Build SMEM descriptors (rank-2 vec_mode layout required) + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufA_s0.iterator, bufA_s0.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufB_s0.iterator, bufB_s0.layout, "mn")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + + # Issue WS SS MMA (scale_out=0 → D = A*B, not accumulate) + if warp_idx == cutlass.Int32(0): + with elect_one(): + for ks in cutlass.range_constexpr(K // 8): + scale = 0 if ks == 0 else 1 + desc_a = desc_a_base + (ks * A_K_STEP_BYTES_TF32) + desc_b = desc_b_base + (ks * B_K_STEP_BYTES_TF32) + tcgen05mma_ws_ss_tf32(desc_a, desc_b, tmem_col, IDESC_TF32_M64_N64, scale) + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + mbarrier_wait(mbar_ptr, 0) + sync_threads() + + # T2R → R2G: tcgen05_ld directly into store_256b (type-agnostic, like C++ reinterpret_cast) + vec_i32 = tcgen05_ld_32x32b(ACC_NUM_COLS, tmem_col) + cute.arch.fence_view_async_tmem_load() + + # 1. reinterpret_cast to f32 (zero-cost bitcast) + # vec_f32 = reinterpret_cast(vec_i32, Int32, ACC_NUM_COLS, Float32) + + # 2. TensorSSA wrap → .to(BFloat16) (real CUDA core CVT) + # regs = TensorSSA(vec_f32, (ACC_NUM_COLS,), Float32) + + # Debug print: thread 0, first 4 register values + # if tidx == cutlass.Int32(0): + # cute.printf("[T2R] tid=0, regs[0..3] = %f, %f, %f, %f", + # regs[0], regs[1], regs[2], regs[3]) + + # R2G via store_256b (4 × 256-bit stores per thread) + # Layout E (column-major warp order): + # warp0->(M0,N0), warp1->(M1,N0), warp2->(M0,N1), warp3->(M1,N1) + lane_idx = tidx % 32 + row = (warp_idx % 2) * 32 + lane_idx + col_base = (warp_idx // 2) * 32 + base_addr = (C_out.iterator + row * N + col_base).toint() + for chunk in cutlass.range_constexpr(ACC_NUM_COLS // 8): + store_256b(base_addr + chunk * 32, subvec(vec_i32, chunk * 8, 8)) + + sync_threads() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr_f32, NUM_COLS) + + @cute.jit + def _launch(self, A: cute.Tensor, B: cute.Tensor, C: cute.Tensor, stream): + self.kernel(A, B, C).launch(grid=(1, 1, 1), block=(128, 1, 1), stream=stream) + + def run(self, A_cpu, B_cpu): + assert K_DIM_TF32 == 8, "TODO: support larger K-dimension" + A_gpu = A_cpu.contiguous().float().cuda() + B_gpu = B_cpu.contiguous().float().cuda() + C_gpu = torch.zeros(M_DIM, N_DIM, dtype=torch.float32, device="cuda") + stream = cutlass_torch.default_stream() + self._launch(from_dlpack(A_gpu), from_dlpack(B_gpu), from_dlpack(C_gpu), stream) + torch.cuda.synchronize() + return C_gpu.cpu() + + +# ===================================================================== +# Test 2: tcgen05mma_ws_ss_f16 (weight-stationary, SMEM A, SMEM B, f16) +# ===================================================================== + + +class _WsSsF16Kernel: + def __init__(self, M: int, N: int, K: int): + self.M = M + self.N = N + self.K = K + if N == 64: + self.idesc = IDESC_F16_M64_N64 + elif N == 128: + self.idesc = IDESC_F16_M64_N128 + else: + raise ValueError(f"Unsupported N={N} for F16 IDESC (expected 64 or 128)") + + @cute.kernel + def kernel(self, A_in: cute.Tensor, B_in: cute.Tensor, C_out: cute.Tensor): + M, N, K = self.M, self.N, self.K + idesc = self.idesc + ACC_NUM_COLS = N // 2 + NUM_COLS = ACC_NUM_COLS + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + smem = utils.SmemAllocator() + tmem_hold_ptr = smem.allocate(Int32) + mbar_ptr = smem.allocate(Int64, byte_alignment=8) + + # Create MMA SMEM Layouts + # NOTE: we use non-ws mode TiledMMA for creating smem layout in a easy way, + # because smem layouts of ws mode and non-ws mode are the same + mma_tiler = (M, N, K) + non_ws_tiled_mma = sm100_utils.make_trivial_tiled_mma( + BFloat16, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (M, N), + ) + a_smem_layout = sm100_utils.make_smem_layout_a(non_ws_tiled_mma, mma_tiler, BFloat16, 1) + b_smem_layout = sm100_utils.make_smem_layout_b(non_ws_tiled_mma, mma_tiler, BFloat16, 1) + bufferA = smem.allocate_tensor( + element_type=BFloat16, + layout=a_smem_layout.outer, + byte_alignment=128, + swizzle=a_smem_layout.inner, + ) + + bufferB = smem.allocate_tensor( + element_type=BFloat16, + layout=b_smem_layout.outer, + byte_alignment=128, + swizzle=b_smem_layout.inner, + ) + + bufA_s0 = bufferA[(None, None, None, 0)] + bufB_s0 = bufferB[(None, None, None, 0)] + + if tidx == cutlass.Int32(0): + mbarrier_init(mbar_ptr, 1) + mbarrier_init_fence() + + # Load A (row-major input → K-major swizzled SMEM) + gA_flat = cute.make_tensor(A_in.iterator, cute.make_layout(M * K)) + gB_flat = cute.make_tensor(B_in.iterator, cute.make_layout(K * N)) + + for step in cutlass.range(M * K // 128, unroll_full=False): + smem_idx = tidx + step * 128 + m = smem_idx % M + k = smem_idx // M + bufA_s0[smem_idx] = gA_flat[m * K + k] + for step in cutlass.range(K * N // 128, unroll_full=False): + idx = tidx + step * 128 + bufB_s0[idx] = gB_flat[idx] + sync_threads() + + # --- TMEM allocation --- + alloc_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128) + tmem = utils.TmemAllocator( + tmem_hold_ptr, + barrier_for_retrieve=alloc_bar, + allocator_warp_id=0, + ) + tmem.allocate(NUM_COLS) + tmem.wait_for_alloc() + tmem_ptr_f32 = tmem.retrieve_ptr(Float32) + + tmem_col_buf = cute.make_tensor(tmem_hold_ptr, cute.make_layout(1)) + tmem_col = tmem_col_buf[0] + + # Build SMEM descriptors (rank-2 vec_mode layout required) + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufA_s0.iterator, bufA_s0.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufB_s0.iterator, bufB_s0.layout, "mn")) + desc_a_base = Tcgen05SmemDescriptor(desc_a_i64) + desc_b_base = Tcgen05SmemDescriptor(desc_b_i64) + + # Per-K-atom byte offsets are derived from the (unswizzled) outer layout + # so we transparently handle every K size: + # K∈{16,32,64} → A k_iter is single-mode, uniform stride + # K≥128 → A k_iter is hierarchical e.g. (4,K/64):(16,4096) + # B is always uniform stride=1024 elem + # Coord ((0,0), 0, ks, 0) into outer layout gives the linear elem offset + # of the ks-th MMA-K atom; * sizeof(elem) → byte offset to add to desc. + ELEM_BYTES_F16 = BFloat16.width // 8 + a_outer = a_smem_layout.outer + b_outer = b_smem_layout.outer + + # Issue WS SS MMA (scale_out=0 → D = A*B, not accumulate) + if warp_idx == cutlass.Int32(0): + with elect_one(): + for ks in cutlass.range_constexpr(K // 16): + scale = 0 if ks == 0 else 1 + a_off = cute.crd2idx(((0, 0), 0, ks, 0), a_outer) * ELEM_BYTES_F16 + b_off = cute.crd2idx(((0, 0), 0, ks, 0), b_outer) * ELEM_BYTES_F16 + desc_a = desc_a_base + a_off + desc_b = desc_b_base + b_off + tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_col, idesc, scale) + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + mbarrier_wait(mbar_ptr, 0) + sync_threads() + + # T2R + # Layout E (M=64, ws mode): 128 lanes, 32 columns + # ref: https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-e + # .32x32b.x32 loads all 32 columns → 32 FP32 regs per thread + # Layout: warp0->(M0,N0), warp1->(M0,N1), warp2->(M1,N0), warp3->(M1,N1) + # for 64x64 Acc, each warp process 32x32, with 128 lanes in TMEM all used + + vec_i32 = tcgen05_ld_32x32b(ACC_NUM_COLS, tmem_col) + cute.arch.fence_view_async_tmem_load() + + # =======DEBUG======== + # # 1. reinterpret_cast to f32 (zero-cost bitcast) + # vec_f32 = reinterpret_cast(vec_i32, Int32, ACC_NUM_COLS, Float32) + + # # 2. TensorSSA wrap → .to(BFloat16) (real CUDA core CVT) + # regs = TensorSSA(vec_f32, (ACC_NUM_COLS,), Float32) + + # # Debug print: thread 0, first 4 register values + # if tidx == cutlass.Int32(0): + # cute.printf("[T2R] tid=0, regs[0..3] = %f, %f, %f, %f", + # regs[0], regs[1], regs[2], regs[3]) + + # R2G via store_256b (4 × 256-bit stores per thread) + # Layout E (column-major warp order): + # warp0->(M0,N0), warp1->(M1,N0), warp2->(M0,N1), warp3->(M1,N1) + # in each warp, each thread process one row, T0->[0, 0:31], T1->[1, 0:31], ..., T31->[31, 0:31] + lane_idx = tidx % 32 + row = (warp_idx % 2) * M // 2 + lane_idx # M0 or M1 + col_base = (warp_idx // 2) * ACC_NUM_COLS # N0 or N1 + # 32 regs = 4 chunks of 8 × 32-bit each (256 bits) + base_addr = (C_out.iterator + row * N + col_base).toint() + for chunk in cutlass.range_constexpr(ACC_NUM_COLS // 8): + store_256b(base_addr + chunk * 32, subvec(vec_i32, chunk * 8, 8)) + + sync_threads() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr_f32, NUM_COLS) + + @cute.jit + def _launch(self, A: cute.Tensor, B: cute.Tensor, C: cute.Tensor, stream): + self.kernel(A, B, C).launch(grid=(1, 1, 1), block=(128, 1, 1), stream=stream) + + def run(self, A_cpu, B_cpu): + M, N = self.M, self.N + A_gpu = A_cpu.cuda().to(torch.bfloat16).contiguous() + B_gpu = B_cpu.cuda().to(torch.bfloat16).contiguous() + C_gpu = torch.zeros(M, N, dtype=torch.float32, device="cuda") + stream = cutlass_torch.default_stream() + self._launch(from_dlpack(A_gpu), from_dlpack(B_gpu), from_dlpack(C_gpu), stream) + torch.cuda.synchronize() + return C_gpu.cpu() + + +# ===================================================================== +# Test 3: tcgen05mma_ws_ss_tf32 with explicit collector_b_buffer/collector_op +# ===================================================================== + + +class _WsSsTf32CollectorKernel: + """Same as _WsSsTf32Kernel but passes collector_b_buffer=B0, collector_op=DISCARD.""" + + @cute.kernel + def kernel(self, A_in: cute.Tensor, B_in: cute.Tensor, C_out: cute.Tensor): + M, N, K = M_DIM, N_DIM, 8 # default K with 8 + ACC_NUM_COLS = N // 2 + NUM_COLS = ACC_NUM_COLS + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + smem = utils.SmemAllocator() + tmem_hold_ptr = smem.allocate(Int32) + mbar_ptr = smem.allocate(Int64, byte_alignment=8) + + non_ws_tiled_mma = sm100_utils.make_trivial_tiled_mma( + TFloat32, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (M, N), + ) + mma_tiler = (M, N, K) + a_smem_layout = sm100_utils.make_smem_layout_a(non_ws_tiled_mma, mma_tiler, TFloat32, 1) + b_smem_layout = sm100_utils.make_smem_layout_b(non_ws_tiled_mma, mma_tiler, TFloat32, 1) + bufferA = smem.allocate_tensor( + element_type=TFloat32, layout=a_smem_layout.outer, byte_alignment=128, swizzle=a_smem_layout.inner + ) + bufferB = smem.allocate_tensor( + element_type=TFloat32, layout=b_smem_layout.outer, byte_alignment=128, swizzle=b_smem_layout.inner + ) + bufA_s0 = bufferA[(None, None, None, 0)] + bufB_s0 = bufferB[(None, None, None, 0)] + + if tidx == cutlass.Int32(0): + mbarrier_init(mbar_ptr, 1) + mbarrier_init_fence() + + gA_flat = cute.make_tensor(A_in.iterator, cute.make_layout(M * K)) + gB_flat = cute.make_tensor(B_in.iterator, cute.make_layout(K * N)) + for step in cutlass.range(M * K // 128, unroll_full=False): + smem_idx = tidx + step * 128 + m = smem_idx % M + k = smem_idx // M + bufA_s0[smem_idx] = gA_flat[m * K + k] + for step in cutlass.range(K * N // 128, unroll_full=False): + idx = tidx + step * 128 + bufB_s0[idx] = gB_flat[idx] + sync_threads() + + alloc_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128) + tmem = utils.TmemAllocator(tmem_hold_ptr, barrier_for_retrieve=alloc_bar, allocator_warp_id=0) + tmem.allocate(NUM_COLS) + tmem.wait_for_alloc() + tmem_ptr_f32 = tmem.retrieve_ptr(Float32) + tmem_col_buf = cute.make_tensor(tmem_hold_ptr, cute.make_layout(1)) + tmem_col = tmem_col_buf[0] + + desc_a_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufA_s0.iterator, bufA_s0.layout, "k")) + desc_b_i64 = smem_descriptor_to_int(make_umma_smem_desc(bufB_s0.iterator, bufB_s0.layout, "mn")) + desc_a = Tcgen05SmemDescriptor(desc_a_i64) + desc_b = Tcgen05SmemDescriptor(desc_b_i64) + + if warp_idx == cutlass.Int32(0): + with elect_one(): + tcgen05mma_ws_ss_tf32( + desc_a, + desc_b, + tmem_col, + IDESC_TF32_M64_N64, + 0, + collector_b_buffer=CollectorBBuffer.B0, + collector_op=CollectorOp.DISCARD, + ) + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + mbarrier_wait(mbar_ptr, 0) + sync_threads() + + vec_i32 = tcgen05_ld_32x32b(NUM_COLS, tmem_col) + cute.arch.fence_view_async_tmem_load() + lane_idx = tidx % 32 + row = (warp_idx % 2) * M // 2 + lane_idx + col_base = (warp_idx // 2) * ACC_NUM_COLS + base_addr = (C_out.iterator + row * N + col_base).toint() + for chunk in cutlass.range_constexpr(ACC_NUM_COLS // 8): + store_256b(base_addr + chunk * 32, subvec(vec_i32, chunk * 8, 8)) + + sync_threads() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr_f32, NUM_COLS) + + @cute.jit + def _launch(self, A: cute.Tensor, B: cute.Tensor, C: cute.Tensor, stream): + self.kernel(A, B, C).launch(grid=(1, 1, 1), block=(128, 1, 1), stream=stream) + + def run(self, A_cpu, B_cpu): + A_gpu = A_cpu.contiguous().float().cuda() + B_gpu = B_cpu.contiguous().float().cuda() + C_gpu = torch.zeros(M_DIM, N_DIM, dtype=torch.float32, device="cuda") + stream = cutlass_torch.default_stream() + self._launch(from_dlpack(A_gpu), from_dlpack(B_gpu), from_dlpack(C_gpu), stream) + torch.cuda.synchronize() + return C_gpu.cpu() + + +# ===================================================================== +# Test functions +# ===================================================================== + + +def test_ws_ss_tf32(): + print("\n=== Test 1: tcgen05mma_ws_ss_tf32 (weight-stationary, SMEM A × SMEM B, tf32) ===") + torch.manual_seed(42) + A = torch.randn(M_DIM, K_DIM_TF32) + B = torch.randn(K_DIM_TF32, N_DIM) + ref = torch.mm(A, B) + got = _WsSsTf32Kernel().run(A, B) + err = (got - ref).abs() + rel = err.max().item() / (ref.abs().max().item() + 1e-8) + max_idx = err.argmax().item() + mi, mj = max_idx // N_DIM, max_idx % N_DIM + print(f" got[0,:4]={got[0, :4].tolist()}") + print(f" ref[0,:4]={ref[0, :4].tolist()}") + print(f" max_rel_err={rel:.4f} at ({mi},{mj}): got={got[mi, mj]:.6f} ref={ref[mi, mj]:.6f}") + assert rel < 0.02, f"FAIL: rel={rel:.4f}" + print(" PASSED") + + +def test_ws_ss_f16(): + print("\n=== Test 3: tcgen05mma_ws_ss_f16 (weight-stationary, SMEM A × SMEM B, f16) ===") + torch.manual_seed(42) + for N in [64, 128]: + for K in [64, 128]: + print(f" --- N={N}, K={K} ---") + A = torch.randn(M_DIM, K) + B = torch.randn(K, N) + ref = torch.mm(A, B) + got = _WsSsF16Kernel(M_DIM, N, K).run(A, B) + err = (got - ref).abs() + rel = err.max().item() / (ref.abs().max().item() + 1e-8) + max_idx = err.argmax().item() + mi, mj = max_idx // N, max_idx % N + print(f" got[0,:4]={got[0, :4].tolist()}") + print(f" ref[0,:4]={ref[0, :4].tolist()}") + print(f" max_rel_err={rel:.4f} at ({mi},{mj}): got={got[mi, mj]:.6f} ref={ref[mi, mj]:.6f}") + assert rel < 0.02, f"FAIL N={N}, K={K}: rel={rel:.4f}" + print(f" PASSED (N={N}, K={K})") + + +def test_ws_ss_tf32_collector(): + """Explicit collector_b_buffer=B0, collector_op=DISCARD should match default.""" + print("\n=== Test 2: tcgen05mma_ws_ss_tf32 + collector (B0::DISCARD) ===") + torch.manual_seed(42) + A = torch.randn(M_DIM, K_DIM_TF32) + B = torch.randn(K_DIM_TF32, N_DIM) + ref = torch.mm(A, B) + got = _WsSsTf32CollectorKernel().run(A, B) + err = (got - ref).abs() + rel = err.max().item() / (ref.abs().max().item() + 1e-8) + max_idx = err.argmax().item() + mi, mj = max_idx // N_DIM, max_idx % N_DIM + print(f" got[0,:4]={got[0, :4].tolist()}") + print(f" ref[0,:4]={ref[0, :4].tolist()}") + print(f" max_rel_err={rel:.4f} at ({mi},{mj}): got={got[mi, mj]:.6f} ref={ref[mi, mj]:.6f}") + assert rel < 0.02, f"FAIL: rel={rel:.4f}" + print(" PASSED") + + +if __name__ == "__main__": + test_ws_ss_tf32() + test_ws_ss_tf32_collector() + test_ws_ss_f16() + print("\n=== All tests passed! ===") From 85dbcb7f4e38cdc0c28bf9a4b9ceb70a3ed60bc0 Mon Sep 17 00:00:00 2001 From: Longxmas <92327126+Longxmas@users.noreply.github.com> Date: Fri, 29 May 2026 09:47:42 +0800 Subject: [PATCH 21/34] [KDA] split unit tests into fast/slow modes * [KDA] split unit tests into fast/slow modes * add kda_fast / kda_slow / kda_backward markers; default mode deselects kda_slow * fast default runs 18 cases, slow 164, full sweep all 192 * full sweep reproduces the original test suite * mark beta fp32 and disable_recompute=True as kda_slow * fast mode always runs full backward + gradient checks; drop unneeded retain_graph * document fast / slow / full test commands in README * default fast mode vs full sweep, ~92% latency reduction --- README.md | 10 +- tests/conftest.py | 41 +++++++ tests/test_kda.py | 206 ++++++++++++++++++++------------- tests/test_kda_compare_fla.py | 208 +++++++++++++++++++++------------- 4 files changed, 303 insertions(+), 162 deletions(-) diff --git a/README.md b/README.md index b894cd32..7bed61e2 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-l **Blackwell (SM10X)** See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA 13.0 for detailed results. - + **Hopper (SM90)** See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 12.9 for detailed results. @@ -142,6 +142,14 @@ python -m pytest tests/test_kda_fused_fwd.py -v python tests/test_lightning_attn.py # Tests for Lightning Attention decode python -m pytest tests/test_la_decode.py -v + +# test_kda.py and test_kda_compare_fla.py support a fast/slow split. +# Fast (default) — representative correctness paths for default CI and local iteration +python -m pytest tests/test_kda.py tests/test_kda_compare_fla.py -v +# Slow — broader stress coverage for nightly or manual runs +python -m pytest -m kda_slow tests/test_kda.py tests/test_kda_compare_fla.py -v +# Full sweep (fast + slow) — run before submitting a PR +python -m pytest -m kda_full tests/test_kda.py tests/test_kda_compare_fla.py -v ```
diff --git a/tests/conftest.py b/tests/conftest.py index 38e19314..f144c10b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,4 @@ +import re import pytest import torch @@ -9,6 +10,24 @@ def _is_sm100() -> bool: def pytest_configure(config): config.addinivalue_line("markers", "sm100_only: only run on SM100 devices") config.addinivalue_line("markers", "sm90_only: skip on SM100 devices") + config.addinivalue_line( + "markers", + "kda_fast: KDA test case included in fast (default) mode", + ) + config.addinivalue_line( + "markers", + "kda_slow: KDA test case excluded from fast (default) mode; " + "include via 'pytest -m kda_slow' or run the full sweep with '-m kda_full'", + ) + config.addinivalue_line( + "markers", + "kda_fast_norecomp: fast-mode KDA config that also runs the disable_recompute=True " + "variant in fast mode (other fast configs run disable_recompute=False only)", + ) + + markexpr = config.option.markexpr + if markexpr and "kda_full" in markexpr: + config.option.markexpr = re.sub(r"\bkda_full\b", "(kda_fast or kda_slow)", markexpr) def pytest_collection_modifyitems(config, items): @@ -16,8 +35,30 @@ def pytest_collection_modifyitems(config, items): skip_non_sm100 = pytest.mark.skip(reason="SM100-only test: skip on non-SM100 devices") skip_on_sm100 = pytest.mark.skip(reason="SM90-only test: skip on SM100") + marker_expr = config.option.markexpr or "" + include_slow = "kda_slow" in marker_expr + skip_slow = pytest.mark.skip( + reason="kda_slow case: run 'pytest -m kda_slow' or the full sweep with '-m kda_full' to include" + ) + skip_fast_norecomp = pytest.mark.skip( + reason="disable_recompute=True runs in fast mode only for kda_fast_norecomp configs; " + "include the rest via '-m kda_slow' or '-m kda_full'" + ) + for item in items: if "sm100_only" in item.keywords and not is_sm100: item.add_marker(skip_non_sm100) if "sm90_only" in item.keywords and is_sm100: item.add_marker(skip_on_sm100) + if include_slow: + continue + if "kda_slow" in item.keywords: + item.add_marker(skip_slow) + continue + callspec = getattr(item, "callspec", None) + if ( + callspec is not None + and callspec.params.get("disable_recompute") + and "kda_fast_norecomp" not in item.keywords + ): + item.add_marker(skip_fast_norecomp) diff --git a/tests/test_kda.py b/tests/test_kda.py index d2b17d32..b503f922 100644 --- a/tests/test_kda.py +++ b/tests/test_kda.py @@ -15,6 +15,10 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # Tests for the chunk-decomposed KDA implementation (python/kda/chunk.py) +# Test modes (see issue #78): +# default: fast subset (kda_fast cases only) +# pytest -m kda_slow: slow stress traces and wider parameter grids +# pytest -m kda_full: full sweep (fast + slow) import pytest import torch @@ -27,8 +31,108 @@ pytestmark = pytest.mark.sm100_only +_FAST = [pytest.mark.kda_fast] +_FAST_NR = [pytest.mark.kda_fast, pytest.mark.kda_fast_norecomp] +_SLOW = [pytest.mark.kda_slow] -@pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) +# (B, T, H, HV, D, gln, mask_p, l2norm, gate, safe_gate, dtype), marks +_FIXED_CONFIGS = [ + ((1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), _FAST_NR), # small fixed (+ no_recomp) + ((2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), _SLOW), + ((2, 1000, 3, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), _SLOW), + ((3, 1024, 4, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), _SLOW), + ((4, 1024, 4, 4, 128, 1, 0, False, False, True, torch.bfloat16), _SLOW), + ((4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), _FAST), # l2norm medium + ((2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), _FAST_NR), # gated (+ no_recomp) + ((4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW), + # GVA cases: HV > H + ((2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), _FAST), # GVA medium + ((2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), _SLOW), + ((2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW), +] + +# (H, HV, D, mask_p, cu_seqlens, dtype, safe_gate), marks +_VARLEN_CONFIGS = [ + ((4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST), # short varlen smoke + ((4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW), + ((4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW), + ((4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), _FAST), # multi-batch varlen + ((4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW), + # ======Varlen test with simulated trace======= + ( + ( + 32, + 32, + 128, + 0, + [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096, 4824, 5101, 5919, 6426, 7137, 7392, 7800, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + ( + ( + 32, + 32, + 128, + 0, + [0, 652, 1255, 1600, 2083, 2345, 2756, 3172, 3767, 4096, 4891, 5236, 5543, 6255, 6480, 6947, 7616, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + ( + ( + 32, + 32, + 128, + 0, + [0, 315, 973, 1283, 2162, 2459, 2678, 2998, 3781, 4096, 4503, 5459, 6318, 6669, 6979, 7583, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + ( + ( + 32, + 32, + 128, + 0, + [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + # ======GVA varlen cases: HV > H======= + ((2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST), # GVA varlen smoke + ((4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW), + ((4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW), + ( + ( + 8, + 32, + 128, + 0, + [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), +] + + +@pytest.mark.parametrize( + "beta_dtype", + [ + pytest.param(torch.float32, id="beta_fp32", marks=pytest.mark.kda_slow), + pytest.param(torch.bfloat16, id="beta_bf16"), + ], +) @pytest.mark.parametrize("disable_recompute", [True, False], ids=["no_recomp", "recomp"]) @pytest.mark.parametrize( ( @@ -46,23 +150,11 @@ ), [ pytest.param( - *test, - id="B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), + *params, + id="B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*params), + marks=marks, ) - for test in [ - (1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 1000, 3, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), - (3, 1024, 4, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 4, 128, 1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), - # GVA cases: HV > H - (2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), - ] + for params, marks in _FIXED_CONFIGS ], ) def test_safe_gate_chunk( @@ -118,7 +210,7 @@ def test_safe_gate_chunk( initial_state=h0.clone(), output_final_state=True, ) - ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() if use_gate_in_kernel: ref_dA, A_log.grad = A_log.grad, None ref_dbias, dt_bias.grad = dt_bias.grad, None @@ -141,7 +233,7 @@ def test_safe_gate_chunk( lower_bound=lower_bound, disable_recompute=disable_recompute, ) - ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() if use_gate_in_kernel: tri_dA, A_log.grad = A_log.grad, None tri_dbias, dt_bias.grad = dt_bias.grad, None @@ -161,69 +253,23 @@ def test_safe_gate_chunk( assert_close("dh0", ref_dh0, tri_dh0, 0.008) -@pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) +@pytest.mark.parametrize( + "beta_dtype", + [ + pytest.param(torch.float32, id="beta_fp32", marks=pytest.mark.kda_slow), + pytest.param(torch.bfloat16, id="beta_bf16"), + ], +) @pytest.mark.parametrize("disable_recompute", [True, False], ids=["no_recomp", "recomp"]) @pytest.mark.parametrize( ("H", "HV", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), [ - pytest.param(*test, id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) - for test in [ - (4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), - (4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), - # ======Varlen test with simulated trace======= - ( - 32, - 32, - 128, - 0, - [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096, 4824, 5101, 5919, 6426, 7137, 7392, 7800, 8192], - torch.bfloat16, - True, - ), - ( - 32, - 32, - 128, - 0, - [0, 652, 1255, 1600, 2083, 2345, 2756, 3172, 3767, 4096, 4891, 5236, 5543, 6255, 6480, 6947, 7616, 8192], - torch.bfloat16, - True, - ), - ( - 32, - 32, - 128, - 0, - [0, 315, 973, 1283, 2162, 2459, 2678, 2998, 3781, 4096, 4503, 5459, 6318, 6669, 6979, 7583, 8192], - torch.bfloat16, - True, - ), - ( - 32, - 32, - 128, - 0, - [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], - torch.bfloat16, - True, - ), - # ======GVA varlen cases: HV > H======= - (2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), - ( - 8, - 32, - 128, - 0, - [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], - torch.bfloat16, - True, - ), - ] + pytest.param( + *params, + id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*params), + marks=marks, + ) + for params, marks in _VARLEN_CONFIGS ], ) def test_safe_gate_chunk_varlen( @@ -273,7 +319,7 @@ def test_safe_gate_chunk_varlen( lower_bound=-5.0 if safe_gate else None, disable_recompute=disable_recompute, ) - ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None @@ -294,7 +340,7 @@ def test_safe_gate_chunk_varlen( ref = torch.cat(ref, 1) ref_ht = torch.cat(ref_ht, 0) - ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad assert_close("o", ref, tri, 0.005) diff --git a/tests/test_kda_compare_fla.py b/tests/test_kda_compare_fla.py index c88a40af..70b2ca9a 100644 --- a/tests/test_kda_compare_fla.py +++ b/tests/test_kda_compare_fla.py @@ -14,7 +14,11 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -# Precision Tests for FlashKDA implementation compared with FLA Triton baselin +# Precision Tests for FlashKDA implementation compared with FLA Triton baseline +# Test modes (see issue #78): +# default: fast subset (kda_fast cases only) +# pytest -m kda_slow: slow stress traces and wider parameter grids +# pytest -m kda_full: full sweep (fast + slow) import pytest import torch @@ -26,8 +30,108 @@ pytestmark = pytest.mark.sm100_only +_FAST = [pytest.mark.kda_fast] +_FAST_NR = [pytest.mark.kda_fast, pytest.mark.kda_fast_norecomp] +_SLOW = [pytest.mark.kda_slow] -@pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) +# (B, T, H, HV, D, gln, mask_p, l2norm, gate, safe_gate, dtype), marks +_FIXED_CONFIGS = [ + ((1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), _FAST_NR), # small fixed (+ no_recomp) + ((2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), _SLOW), + ((2, 1000, 3, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), _SLOW), + ((3, 1024, 4, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), _SLOW), + ((4, 1024, 4, 4, 128, 1, 0, False, False, True, torch.bfloat16), _SLOW), + ((4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), _FAST), # l2norm medium + ((2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), _FAST_NR), # gated (+ no_recomp) + ((4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW), + # GVA cases: HV > H + ((2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), _FAST), # GVA medium + ((2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), _SLOW), + ((2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW), +] + +# (H, HV, D, mask_p, cu_seqlens, dtype, safe_gate), marks +_VARLEN_CONFIGS = [ + ((4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST), # short varlen smoke + ((4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW), + ((4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW), + ((4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), _FAST), # multi-batch varlen + ((4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW), + # ======Varlen test with simulated trace======= + ( + ( + 32, + 32, + 128, + 0, + [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096, 4824, 5101, 5919, 6426, 7137, 7392, 7800, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + ( + ( + 32, + 32, + 128, + 0, + [0, 652, 1255, 1600, 2083, 2345, 2756, 3172, 3767, 4096, 4891, 5236, 5543, 6255, 6480, 6947, 7616, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + ( + ( + 32, + 32, + 128, + 0, + [0, 315, 973, 1283, 2162, 2459, 2678, 2998, 3781, 4096, 4503, 5459, 6318, 6669, 6979, 7583, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + ( + ( + 32, + 32, + 128, + 0, + [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), + # ======GVA varlen cases: HV > H======= + ((2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST), # GVA varlen smoke + ((4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW), + ((4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW), + ( + ( + 8, + 32, + 128, + 0, + [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], + torch.bfloat16, + True, + ), + _SLOW, + ), +] + + +@pytest.mark.parametrize( + "beta_dtype", + [ + pytest.param(torch.float32, id="beta_fp32", marks=pytest.mark.kda_slow), + pytest.param(torch.bfloat16, id="beta_bf16"), + ], +) @pytest.mark.parametrize("disable_recompute", [True, False], ids=["no_recomp", "recomp"]) @pytest.mark.parametrize( ( @@ -45,23 +149,11 @@ ), [ pytest.param( - *test, - id="B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*test), + *params, + id="B{}-T{}-H{}-HV{}-D{}-gln{}-mask_p{}-l2norm{}-gate{}-safe_gate{}-{}".format(*params), + marks=marks, ) - for test in [ - (1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), - (2, 1000, 3, 3, 128, 1, 0.5, False, False, True, torch.bfloat16), - (3, 1024, 4, 4, 128, 0.1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 4, 128, 1, 0, False, False, True, torch.bfloat16), - (4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), - # GVA cases: HV > H - (2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), - (2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), - (2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), - ] + for params, marks in _FIXED_CONFIGS ], ) def test_safe_gate_chunk( @@ -122,7 +214,7 @@ def test_safe_gate_chunk( lower_bound=lower_bound, disable_recompute=disable_recompute, ) - ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() if use_gate_in_kernel: ref_dA, A_log.grad = A_log.grad, None ref_dbias, dt_bias.grad = dt_bias.grad, None @@ -145,7 +237,7 @@ def test_safe_gate_chunk( lower_bound=lower_bound, disable_recompute=disable_recompute, ) - ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() if use_gate_in_kernel: tri_dA, A_log.grad = A_log.grad, None tri_dbias, dt_bias.grad = dt_bias.grad, None @@ -165,69 +257,23 @@ def test_safe_gate_chunk( assert_close("dh0", ref_dh0, tri_dh0, 0.008) -@pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) +@pytest.mark.parametrize( + "beta_dtype", + [ + pytest.param(torch.float32, id="beta_fp32", marks=pytest.mark.kda_slow), + pytest.param(torch.bfloat16, id="beta_bf16"), + ], +) @pytest.mark.parametrize("disable_recompute", [True, False], ids=["no_recomp", "recomp"]) @pytest.mark.parametrize( ("H", "HV", "D", "mask_p", "cu_seqlens", "dtype", "safe_gate"), [ - pytest.param(*test, id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*test)) - for test in [ - (4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), - (4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), - # ======Varlen test with simulated trace======= - ( - 32, - 32, - 128, - 0, - [0, 247, 699, 982, 1688, 1985, 2383, 3081, 3526, 3973, 4096, 4824, 5101, 5919, 6426, 7137, 7392, 7800, 8192], - torch.bfloat16, - True, - ), - ( - 32, - 32, - 128, - 0, - [0, 652, 1255, 1600, 2083, 2345, 2756, 3172, 3767, 4096, 4891, 5236, 5543, 6255, 6480, 6947, 7616, 8192], - torch.bfloat16, - True, - ), - ( - 32, - 32, - 128, - 0, - [0, 315, 973, 1283, 2162, 2459, 2678, 2998, 3781, 4096, 4503, 5459, 6318, 6669, 6979, 7583, 8192], - torch.bfloat16, - True, - ), - ( - 32, - 32, - 128, - 0, - [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], - torch.bfloat16, - True, - ), - # ======GVA varlen cases: HV > H======= - (2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), - (4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), - (4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), - ( - 8, - 32, - 128, - 0, - [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192], - torch.bfloat16, - True, - ), - ] + pytest.param( + *params, + id="H{}-HV{}-D{}-mask_p{}-cu_seqlens{}-{}-safe_gate{}".format(*params), + marks=marks, + ) + for params, marks in _VARLEN_CONFIGS ], ) def test_safe_gate_chunk_varlen( @@ -277,7 +323,7 @@ def test_safe_gate_chunk_varlen( lower_bound=-5.0 if safe_gate else None, disable_recompute=disable_recompute, ) - ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None @@ -295,7 +341,7 @@ def test_safe_gate_chunk_varlen( lower_bound=-5.0 if safe_gate else None, disable_recompute=disable_recompute, ) - ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad assert_close("o", ref, tri, 0.005) From b6d8b2d891c6d973c7daf17c51be043bb90c67c2 Mon Sep 17 00:00:00 2001 From: cher <117337477+cherhh@users.noreply.github.com> Date: Wed, 3 Jun 2026 11:46:34 +0800 Subject: [PATCH 22/34] [KDA] Add intra-card CP for chunk_delta_h forward in SM100 (#70) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: intracard context parallel for chunk_delta_h * bench: add bench_intracard_cp.py * test: add intracard CP tests * fix: avoid repeated D2H sync in intracard CP auto-dispatch * fix Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix: remove unused T_total param from should_use_intracard_cp * fix Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * refactor(merge): remove secondary metadata cache Replace _get_meta_tensors LRU cache with direct torch.tensor calls. The secondary cache required converting lists to tuples for hashing on every merge_fwd call. Since the metadata lists are tiny (1-4 int32 values), the hashing overhead outweighs any benefit. Each torch.tensor call produces an independently aligned allocation required by from_dlpack(assumed_align=16). A single merged tensor with slice views would not guarantee per-slice 16-byte alignment. * docs: add intracard CP usage section to USAGE.md * fix: empty warp setmaxnreg * [Feat] upgrade FLA to v0.5.0 (#72) * upgrade fla and update b200 bench, update readme and fix lightning test param * update h200 bench result with fla bug fixed * update b200 bench * update b200 bench * fix readme * remove useless repeat_interleave for fla * fix readme --------- Co-authored-by: boyu.zbw * bench: add 128K+Nx1K configs for reviewer request * fix: simplify compute_subseq_len and add sub-seq length guard 1. compute_subseq_len: remove power-of-2 snap, use floor division for target_splits. This ensures ceil(seq_chunks / target_splits) <= target_splits so Guard 3 in intracard_fwd_h no longer fires spuriously for the common single-long-seq case. 2. should_use_intracard_cp: add Guard 3 (expected sub-seq length check). CP merge work scales with H; require expected_subseq_c >= 12*H chunks. Add MIN_SUBSEQ_CHUNKS_PER_HEAD=12 constant. Restructure Guard 2 from an inline return to an explicit if so Guard 3 can follow. This fixes worst-case degradation (28K+4K H=8: 0.87x ??? 1.00x). All previously-degraded configs now correctly bypass CP. 3. tests: update ACCURACY_CONFIGS, FINAL_STATE_CONFIGS, STRESS and h0_none_equiv to use 65536 (instead of 32768) as the long sequence so they still exercise the CP kernel path under the new Guard 3 threshold. 29/29 tests pass. Worst triggered speedup: 1.11x. Max bypass overhead: 1.016x. * fix: thread-safe cache, SM-aware pre_scan dispatch, refactor bench * fix: update imports after chunk_delta_h → chunk_delta_h_sm100 rename * fix: address review comments — .tolist() loop, remove unused param, NotImplementedError for SM90 pre_scan --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Kevinzz <2538015266@qq.com> Co-authored-by: boyu.zbw --- USAGE.md | 51 ++ benchmarks/bench_intracard_cp.py | 321 +++++++ cula/kda/chunk_fwd.py | 1 + cula/ops/chunk_delta_h_sm100.py | 40 + cula/ops/cp/__init__.py | 3 + cula/ops/cp/chunk_delta_h.py | 595 +++++++++++++ cula/ops/cp/merge.py | 533 ++++++++++++ cula/ops/cp/pre_scan.py | 1335 ++++++++++++++++++++++++++++++ cula/utils.py | 26 + tests/test_intracard_cp.py | 466 +++++++++++ 10 files changed, 3371 insertions(+) create mode 100644 benchmarks/bench_intracard_cp.py create mode 100644 cula/ops/cp/__init__.py create mode 100644 cula/ops/cp/chunk_delta_h.py create mode 100644 cula/ops/cp/merge.py create mode 100644 cula/ops/cp/pre_scan.py create mode 100644 tests/test_intracard_cp.py diff --git a/USAGE.md b/USAGE.md index 80274b29..86d2813c 100644 --- a/USAGE.md +++ b/USAGE.md @@ -112,3 +112,54 @@ print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] - Mainly **suitable for large-batch inference**; performance is limited when both batch size and head count are small, because we do not parallelize over the sequence-length dimension. - **Matrix inversion uses fp16 precision**, which is faster and occupies less shared memory but introduces minor numerical differences compared to tf32 inversion. - **Intra-subchunk attention uses g-first as anchor**, which causes some numerical differences compared with the FLA Triton implementation (FLA uses g-half as anchor in the diagonal). + +--- + +## Intra-Card Context Parallel (chunk_delta_h) + +cuLA includes an intra-card context parallel (CP) path for `chunk_gated_delta_rule_fwd_h`. Long sequences are split into sub-sequences, processed independently in parallel, then merged via a prefix-scan step — unlocking sequence-dimension parallelism on a single GPU. + +**Requirements** + +| Condition | Detail | +|---|---| +| Environment variable | `CULA_INTRACARD_CP=1` | +| Execution context | Inside `torch.inference_mode()` | +| Input mode | Varlen only (`cu_seqlens` must be provided) | +| Global gate | `g=None` (scalar gate `g` not supported; key-dim gate `gk` is supported) | + +If the heuristic decides CP would not help (e.g. batch already saturates SMs, or sequences are too short), it silently falls back to the standard single-pass kernel. + +**Example** + +```python +import os +os.environ["CULA_INTRACARD_CP"] = "1" + +import torch +from cula.ops.chunk_delta_h import chunk_gated_delta_rule_fwd_h + +B, T, H, K, V = 1, 65536, 8, 128, 128 +device = 'cuda' + +k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) +w = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) +u = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) +cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=device) + +with torch.inference_mode(): + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, w=w, u=u, + cu_seqlens=cu_seqlens, + output_final_state=True, + ) + +print(f'h shape: {h.shape}') # [1, NT, H, K, V] +print(f'final_state shape: {final_state.shape}') # [1, H, K, V] +``` + +**Notes** + +- CP is only beneficial when a small number of long sequences under-utilise the SM array. The built-in heuristic checks SM saturation, minimum sequence length (≥ 256 chunks), and effective batch size before enabling CP. +- Currently **inference-only**; the backward pass is not supported through the CP path. +- `cu_seqlens` must be **`int32`**. diff --git a/benchmarks/bench_intracard_cp.py b/benchmarks/bench_intracard_cp.py new file mode 100644 index 00000000..f24fb7bf --- /dev/null +++ b/benchmarks/bench_intracard_cp.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +bench_intracard_cp.py — Benchmark: Intracard Context-Parallel speedup + for chunk_kda (KDA forward) + +Measures the speedup of cuLA's intracard context-parallel path against the +non-CP baseline across a range of varlen configurations. Also verifies that +the heuristic does not regress throughput when CP is correctly bypassed. + +Usage: + python bench_intracard_cp.py [--ncu] [--sanitizer] + +With --ncu, warmup=1 and iters=1 for ncu profiling: + ncu --set full -o report python bench_intracard_cp.py --ncu +""" + +import argparse +import contextlib +import os +import pathlib +import sys + +os.environ.setdefault("CULA_INTRACARD_CP", "1") + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch + +from benchmarks.utils import ( + SEED, + exclusive_cumsum, + prepare_safe_gate_inputs, + set_seed, +) +from cula.kda.chunk_fwd import chunk_kda_fwd +from cula.ops.cp.chunk_delta_h import ( + compute_subseq_len, + prepare_subseq_cu_seqlens, + should_use_intracard_cp, +) +from cula.utils import get_device_sm_count + +# ============================================================ +# Constants +# ============================================================ +BT, D = 64, 128 +H_VALUES = [4, 8] +WARMUP = 10 +N_ITERS = 100 +NCU_MODE = False +SANITIZER_MODE = False + +# (tag, seq_lens) — each entry is tested at every H in H_VALUES +CONFIGS = [ + # --- single seq (ascending length) --- + ("T=4K", [4096]), + ("T=8K", [8192]), + ("T=32K", [32768]), + ("T=64K", [65536]), + ("T=128K", [131072]), + # --- equal-length batches (~32K total) --- + ("8x4K", [4096] * 8), + ("4x8K", [8192] * 4), + ("2x16K", [16384] * 2), + # --- asymmetric multi-seq --- + ("16K+16K", [16384, 16384]), + ("24K+8K", [24576, 8192]), + ("28K+4K", [28672, 4096]), + ("32K+256+256", [32768, 256, 256]), + ("40K+1K+8K", [40960, 1024, 8192]), + ("64K+512+256+128", [65536, 512, 256, 128]), + ("128K+1K", [131072, 1024]), + # --- 128K + several short seqs --- + ("128K+2x1K", [131072, 1024, 1024]), + ("128K+5x1K", [131072] + [1024] * 5), + ("128K+10x1K", [131072] + [1024] * 10), +] + + +# ============================================================ +# CP toggle +# ============================================================ +@contextlib.contextmanager +def cp_on(enable: bool): + old = os.environ.get("CULA_INTRACARD_CP") + os.environ["CULA_INTRACARD_CP"] = "1" if enable else "0" + try: + if enable: + with torch.inference_mode(): + yield + else: + yield + finally: + if old is None: + os.environ.pop("CULA_INTRACARD_CP", None) + else: + os.environ["CULA_INTRACARD_CP"] = old + + +# ============================================================ +# Helpers +# ============================================================ +def time_kernel(fn, warmup=None, n_iters=None): + if warmup is None: + warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + if n_iters is None: + n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() + for _ in range(n_iters): + fn() + end_evt.record() + torch.cuda.synchronize() + return start_evt.elapsed_time(end_evt) / n_iters + + +def run_cp(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, enable_cp): + with cp_on(enable_cp): + chunk_kda_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens.cpu(), + safe_gate=True, + lower_bound=lower_bound, + use_gate_in_kernel=True, + A_log=A_log, + dt_bias=dt_bias, + ) + + +def predict_cp(seq_lens, H, num_sms): + cu = torch.tensor( + exclusive_cumsum(seq_lens), + dtype=torch.int32, + ) + if not should_use_intracard_cp(cu, num_sms, H, BT): + return False, 0 + max_len = int(torch.diff(cu).max().item()) + subseq_len = compute_subseq_len(max_len, num_sms, H, BT, num_seqs=len(seq_lens)) + _, split_info, total_subseqs = prepare_subseq_cu_seqlens(cu, subseq_len, BT) + return bool(split_info), total_subseqs + + +# ============================================================ +# Benchmark +# ============================================================ +def bench_cp(h_values, configs): + print("\n" + "=" * 100) + print(" Intracard CP Benchmark: CP-on vs CP-off") + print("=" * 100) + + device = torch.device("cuda") + num_sms = get_device_sm_count(device) + results = [] + + for H in h_values: + for tag, seq_lens in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + total_T = sum(seq_lens) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens, seed=SEED) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, lower_bound = inputs["scale"], inputs["lower_bound"] + + pred, n_sub = predict_cp(seq_lens, H, num_sms) + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + + ms_off = time_kernel(lambda: run_cp(**common, enable_cp=False)) + ms_on = time_kernel(lambda: run_cp(**common, enable_cp=True)) + + r = { + "tag": tag, + "H": H, + "total_T": total_T, + "pred": pred, + "n_sub": n_sub, + "ms_off": ms_off, + "ms_on": ms_on, + "speedup": ms_off / ms_on if ms_on > 0 else float("inf"), + } + results.append(r) + + del q, k, v, g, beta, A_log, dt_bias, inputs + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Report +# ============================================================ +def print_report(results, h_values): + sep = "=" * 110 + print(f"\n\n{sep}") + print(" BENCHMARK REPORT: Intracard CP") + print(" CP-on vs CP-off (same kernel, different code paths)") + print(f" D={D} dtype=bf16 safe_gate=True") + wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") + print(f" Warmup={wu} Iters={ni}{mode_tag}") + print(sep) + + for H_val in h_values: + h_results = [r for r in results if r["H"] == H_val] + if not h_results: + continue + + print(f"\n [H={H_val}]") + print(f" {'─' * 95}") + print( + f" {'config':<24s} {'T':>7s} {'pred':>4s} {'sub':>4s}" + f" │ {'CP_off(ms)':>10s} {'CP_on(ms)':>10s} {'Speedup':>8s}" + ) + print(f" {'─' * 95}") + for r in h_results: + pred_s = "Y" if r["pred"] else "N" + print( + f" {r['tag']:<24s} {r['total_T']:>7d} {pred_s} {r['n_sub']:>4d}" + f" │ {r['ms_off']:>10.4f} {r['ms_on']:>10.4f} {r['speedup']:>7.2f}x" + ) + print(f" {'─' * 95}") + + # Summary + triggered = [r for r in results if r["pred"]] + bypassed = [r for r in results if not r["pred"]] + + if triggered: + speedups = [r["speedup"] for r in triggered] + geo = 1.0 + for s in speedups: + geo *= s + geo = geo ** (1 / len(speedups)) + print( + f"\n CP triggered ({len(triggered)} configs): " + f"geo-mean={geo:.2f}x best={max(speedups):.2f}x worst={min(speedups):.2f}x" + ) + + if bypassed: + ratios = [r["ms_on"] / r["ms_off"] for r in bypassed] + print( + f" CP bypassed ({len(bypassed)} configs): " + f"mean overhead={sum(ratios) / len(ratios):.3f}x max={max(ratios):.3f}x " + f"(1.00 = no regression)" + ) + + print(f"\n{sep}\n") + + +# ============================================================ +# Main +# ============================================================ +def main(): + parser = argparse.ArgumentParser(description="bench_intracard_cp: CP-on vs CP-off benchmark") + parser.add_argument( + "--ncu", + action="store_true", + help="NCU profiling mode: warmup=1, iters=1", + ) + parser.add_argument( + "--sanitizer", + action="store_true", + help="Sanitizer mode: warmup=1, iters=1", + ) + args = parser.parse_args() + + global NCU_MODE, SANITIZER_MODE + if args.ncu: + NCU_MODE = True + print("[NCU mode] warmup=1, iters=1") + if args.sanitizer: + SANITIZER_MODE = True + print("[Sanitizer mode] warmup=1, iters=1") + + results = bench_cp(H_VALUES, CONFIGS) + print_report(results, H_VALUES) + return results + + +if __name__ == "__main__": + main() diff --git a/cula/kda/chunk_fwd.py b/cula/kda/chunk_fwd.py index f49364bb..d0a93fac 100644 --- a/cula/kda/chunk_fwd.py +++ b/cula/kda/chunk_fwd.py @@ -117,6 +117,7 @@ def chunk_kda_fwd( output_final_state=output_final_state, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + cu_seqlens_cpu=cu_seqlens_cpu, ) if cp_context is not None: diff --git a/cula/ops/chunk_delta_h_sm100.py b/cula/ops/chunk_delta_h_sm100.py index a34f7d84..c4c84af1 100644 --- a/cula/ops/chunk_delta_h_sm100.py +++ b/cula/ops/chunk_delta_h_sm100.py @@ -18,6 +18,7 @@ """ import argparse +import os as _os import cutlass import cutlass.cute as cute @@ -40,6 +41,16 @@ COMPILE_OPTIONS = "--enable-tvm-ffi --generate-line-info --ptxas-options '--verbose'" +# Intracard CP auto-dispatch +def _intracard_cp_enabled() -> bool: + """Return whether intracard-CP is currently enabled (runtime check). + + Env var truthiness matches FLA: any value other than "0" enables it. + Default (unset) is "0" → disabled. + """ + return _os.environ.get("CULA_INTRACARD_CP", "0") != "0" + + # in FLA, cumsum returns int64 tensor by default @tensor_cache def prepare_chunk_offsets_i32( @@ -2015,6 +2026,8 @@ def chunk_gated_delta_rule_fwd_h( cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, persistent: bool = True, + _no_cp: bool = False, + cu_seqlens_cpu: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """ ChunkDeltaRuleFwdH forward pass — FLA-compatible API. @@ -2046,6 +2059,33 @@ def chunk_gated_delta_rule_fwd_h( v_new: [B, T, HV, V] bf16 (or None if save_new_value=False) final_state: [N, HV, K, V] fp32 (or None if output_final_state=False) """ + # --- Intracard CP auto-dispatch --- + if _intracard_cp_enabled() and not _no_cp and cu_seqlens is not None and g is None and torch.is_inference_mode_enabled(): + from cula.ops.cp.chunk_delta_h import intracard_fwd_h, should_use_intracard_cp + from cula.utils import get_device_sm_count + + # Materialize cu_seqlens_cpu once here to avoid repeated D2H sync inside intracard_fwd_h. + _cu_seqlens_cpu = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.cpu() + if should_use_intracard_cp( + _cu_seqlens_cpu, + get_device_sm_count(k.device), + k.shape[2], + chunk_size, + ): + return intracard_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + cu_seqlens_cpu=_cu_seqlens_cpu, + ) + B, T, H, K_dim = k.shape HV = u.shape[2] V_dim = u.shape[3] diff --git a/cula/ops/cp/__init__.py b/cula/ops/cp/__init__.py new file mode 100644 index 00000000..5727f9d4 --- /dev/null +++ b/cula/ops/cp/__init__.py @@ -0,0 +1,3 @@ +from cula.ops.cp.chunk_delta_h import intracard_fwd_h + +__all__ = ["intracard_fwd_h"] diff --git a/cula/ops/cp/chunk_delta_h.py b/cula/ops/cp/chunk_delta_h.py new file mode 100644 index 00000000..6762476e --- /dev/null +++ b/cula/ops/cp/chunk_delta_h.py @@ -0,0 +1,595 @@ +# Copyright (c) 2025 ANTGROUP. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Intra-Card Context Parallel (CP) for Chunk Delta H. + +Overview: + Long sequences on a single card are split into sub-sequences, each processed + independently via cuLA's CuTeDSL chunk_delta_h kernel. A prefix-scan merge + step propagates initial states across sub-sequences, eliminating the sequential + bottleneck of the original single-pass recurrence. + +Pipeline (3 stages): + 1. Pre-Scan: For each sub-sequence, compute packed (he, m) state: + he [K, V] = cumulative delta-rule update (the "h-exit" state) + m [K, K] = cumulative decay matrix + Packed as hm [S_split, H, K, K+V] where columns [0:V]=he, [V:V+K]=m + + 2. Merge: Prefix scan across sub-sequences of the same original sequence. + For sub-sequence j: h0_j = m_j @ h0_{j-1} + he_j + Produces per-sub-sequence initial states. + + 3. Forward H: Run cuLA's existing chunk_gated_delta_rule_fwd_h on the + split sub-sequences with the merged initial states. + +Reference: + - FLA intra-card CP: fla/ops/common/intracard_cp.py + - FLA CP kernels: fla/ops/cp/chunk_delta_h.py + - cuLA chunk_delta_h: cula/ops/chunk_delta_h.py +""" + +from __future__ import annotations + +import threading +import weakref +from collections import OrderedDict +from typing import NamedTuple + +import torch + +from cula.utils import get_device_sm_count, get_pre_scan + +# Lazy import to avoid circular dependency with cula.ops.chunk_delta_h +_chunk_gated_delta_rule_fwd_h = None + + +def _get_fwd_h(): + global _chunk_gated_delta_rule_fwd_h + if _chunk_gated_delta_rule_fwd_h is None: + from cula.ops.chunk_delta_h_sm100 import chunk_gated_delta_rule_fwd_h + + _chunk_gated_delta_rule_fwd_h = chunk_gated_delta_rule_fwd_h + return _chunk_gated_delta_rule_fwd_h + + +class SplitSeqInfo(NamedTuple): + """Metadata for sequences split into sub-sequences.""" + + split_seq_ids: list[int] # original sequence indices that were split + start_subseq_idx: list[int] # first sub-seq index in expanded cu_seqlens per split seq + num_subseqs: list[int] # number of sub-sequences per split seq + + +class _CacheEntry(NamedTuple): + """Cached precomputed indices and GPU tensors for a given cu_seqlens layout.""" + + cu_seqlens_ref: weakref.ref + cu_seqlens_subseq_values: list[int] + split_info: SplitSeqInfo + total_subseqs: int + non_first_indices: torch.Tensor # [num_non_first] int64 GPU + first_subseq_indices: torch.Tensor # [N_orig] int64 GPU + last_subseq_indices: torch.Tensor # [N_orig] int64 GPU + num_non_first: int + merge_seq_starts: list[int] + merge_seq_counts: list[int] + merge_init_offsets: list[int] + cu_seqlens_subseq_gpu: torch.Tensor + chunk_indices_subseq: torch.Tensor # [NT_subseq, 2] int32 + + +_intracard_cache: OrderedDict[tuple, _CacheEntry] = OrderedDict() +_intracard_cache_lock = threading.Lock() +_INTRACARD_CACHE_MAXSIZE = 8 + + +def _prepare_chunk_indices( + cu_seqlens_values: list[int], + chunk_size: int, + device: torch.device, +) -> torch.Tensor: + """Build chunk_indices [NT, 2] int32 from cu_seqlens CPU list.""" + num_seqs = len(cu_seqlens_values) - 1 + seq_ids: list[int] = [] + chunk_ids: list[int] = [] + for i in range(num_seqs): + nc = (cu_seqlens_values[i + 1] - cu_seqlens_values[i] + chunk_size - 1) // chunk_size + seq_ids.extend([i] * nc) + chunk_ids.extend(range(nc)) + return torch.stack( + [torch.tensor(seq_ids, dtype=torch.int32, device=device), torch.tensor(chunk_ids, dtype=torch.int32, device=device)], + dim=1, + ) + + +# Tunable thresholds — empirically calibrated on B200 SM100 (SM=152). +NUM_V_BLOCKS = 2 # fwd_h grid V-tile factor: grid = (NUM_V_BLOCKS, N*H) +MIN_SUBSEQ_CHUNKS = 16 # min chunks per sub-sequence +MIN_LONG_SEQ_CHUNKS = 256 # min chunks of the longest seq to consider CP +MAX_BE_H = 10 # max Be*H; above this CP gain < overhead (~3%) +MIN_SUBSEQ_CHUNKS_PER_HEAD = 12 # min expected sub-seq chunks per head (H-scaled Guard 3) + + +def should_use_intracard_cp( + cu_seqlens_cpu: torch.Tensor, + num_sms: int, + H: int, + chunk_size: int = 64, +) -> bool: + """Pure-Python predicate: should we dispatch to intracard CP? + + Four cheap CPU-only guards (a fifth post-split guard lives in intracard_fwd_h): + Guard 0: baseline already saturates SMs. + Guard 1: longest sequence too short to amortize CP overhead. + Guard 2: Be*H > MAX_BE_H — other seqs already provide enough parallelism. + Guard 3: expected sub-seq too short — CP merge overhead exceeds gain. + """ + cu_list = cu_seqlens_cpu.tolist() + num_seqs = len(cu_list) - 1 + if num_seqs == 0: + return False + + if NUM_V_BLOCKS * H * num_seqs >= num_sms: # Guard 0 + return False + + chunks = [(cu_list[i + 1] - cu_list[i] + chunk_size - 1) // chunk_size for i in range(num_seqs)] + max_c = max(chunks) + + if max_c < MIN_LONG_SEQ_CHUNKS: # Guard 1 + return False + + # Guard 2: Be = effective batch size (as if every seq were max_c chunks long) + Be = sum(chunks) / max_c + if Be * H > MAX_BE_H: + return False + + # Guard 3: expected sub-seq length must be long enough to amortise CP overhead. + # CP merge work scales with H, so the minimum is proportional to H. + per_seq_units = NUM_V_BLOCKS * H + sm_budget = max(num_sms - per_seq_units * max(num_seqs - 1, 0), per_seq_units) + target_splits = max(2, sm_budget // per_seq_units) + expected_subseq_c = max((max_c + target_splits - 1) // target_splits, MIN_SUBSEQ_CHUNKS) + return expected_subseq_c >= MIN_SUBSEQ_CHUNKS_PER_HEAD * H + + +def compute_subseq_len( + seq_len: int, + num_sms: int, + num_heads: int, + chunk_size: int = 64, + num_seqs: int = 1, +) -> int: + """Compute target sub-sequence length for intracard splitting. + + Targets enough splits to saturate remaining SMs after other sequences + in the batch occupy their share. Uses floor division so that + actual sub-seqs never exceed target_splits, guaranteeing Guard 3 + (total_subseqs * NUM_V_BLOCKS * H <= num_sms) is always satisfied + for a single split sequence in the batch. + Result is floored at MIN_SUBSEQ_CHUNKS * chunk_size. + """ + seq_chunks = (seq_len + chunk_size - 1) // chunk_size + + if seq_chunks < 8: + return seq_len + + per_seq_units = NUM_V_BLOCKS * num_heads + sm_budget = max(num_sms - per_seq_units * max(num_seqs - 1, 0), per_seq_units) + target_splits = max(2, sm_budget // per_seq_units) + + subseq_chunks = (seq_chunks + target_splits - 1) // target_splits + subseq_chunks = max(subseq_chunks, MIN_SUBSEQ_CHUNKS) + + return subseq_chunks * chunk_size + + +def prepare_subseq_cu_seqlens( + cu_seqlens_cpu: torch.Tensor, + subseq_len: int, + chunk_size: int = 64, + max_splits: int = 32, +) -> tuple[list[int], SplitSeqInfo | bool, int]: + """Insert sub-sequence split points into cu_seqlens. + + Sequences >= 3 * subseq_len are split into evenly-sized sub-sequences + (each a multiple of chunk_size); shorter sequences are kept intact. + Returns (expanded boundaries, SplitSeqInfo or False, total_subseqs). + """ + cu_list = cu_seqlens_cpu.tolist() + N = len(cu_list) - 1 + if N == 0: + return cu_list, False, 0 + + subseq_chunks = (subseq_len + chunk_size - 1) // chunk_size + threshold_subseq_len = 3 * subseq_len + + split_seq_ids: list[int] = [] + start_subseq_idxs: list[int] = [] + num_subseqs_list: list[int] = [] + + boundaries: list[int] = [0] + cumsum_offset = 0 + + for i in range(N): + seq_start = cu_list[i] + seq_end = cu_list[i + 1] + seq_len_i = seq_end - seq_start + seq_chunks_i = (seq_len_i + chunk_size - 1) // chunk_size + + if seq_len_i >= threshold_subseq_len: + num_ss = min(max_splits, (seq_chunks_i + subseq_chunks - 1) // subseq_chunks) + chunks_per = (seq_chunks_i + num_ss - 1) // num_ss + actual_ssl = chunks_per * chunk_size + split_seq_ids.append(i) + start_subseq_idxs.append(cumsum_offset) + num_subseqs_list.append(num_ss) + for j in range(num_ss): + boundary = min(seq_start + (j + 1) * actual_ssl, seq_end) + boundaries.append(boundary) + cumsum_offset += num_ss + else: + boundaries.append(seq_end) + cumsum_offset += 1 + + if not split_seq_ids: + return cu_list, False, 0 + + total_subseqs = cumsum_offset + split_info = SplitSeqInfo( + split_seq_ids=split_seq_ids, + start_subseq_idx=start_subseq_idxs, + num_subseqs=num_subseqs_list, + ) + return boundaries, split_info, total_subseqs + + +class _PrecomputedIndices(NamedTuple): + """Derived scatter/gather indices for the CP orchestrator.""" + + non_first_indices: list[int] # where to scatter merge results + first_subseq_indices: list[int] # first sub-seq index per original seq + last_subseq_indices: list[int] # last sub-seq index per original seq + num_non_first: int + merge_seq_starts: list[int] + merge_seq_counts: list[int] + merge_init_offsets: list[int] + + +def _precompute_intracard_indices( + split_info: SplitSeqInfo, + N_orig: int, +) -> _PrecomputedIndices: + """Precompute scatter/gather indices from split metadata.""" + starts = split_info.start_subseq_idx + num_ss = split_info.num_subseqs + split_ids = split_info.split_seq_ids + + num_subseqs_per_seq = [1] * N_orig + for sid, nss in zip(split_ids, num_ss): + num_subseqs_per_seq[sid] = nss + + non_first_indices: list[int] = [] + for s, n in zip(starts, num_ss): + for j in range(1, n): + non_first_indices.append(s + j) + + first_subseq_indices: list[int] = [0] + running = 0 + for i in range(N_orig - 1): + running += num_subseqs_per_seq[i] + first_subseq_indices.append(running) + + last_subseq_indices: list[int] = [] + running = 0 + for n in num_subseqs_per_seq: + running += n + last_subseq_indices.append(running - 1) + + # merge_seq_starts/counts use per-seq start indices (not CSR offsets) because + # split sub-seqs may be non-contiguous in hm when unsplit seqs exist in between. + merge_seq_starts: list[int] = list(starts) + merge_seq_counts: list[int] = list(num_ss) + merge_init_offsets: list[int] = [0] + for n in num_ss: + merge_init_offsets.append(merge_init_offsets[-1] + n - 1) + num_non_first = merge_init_offsets[-1] + + return _PrecomputedIndices( + non_first_indices=non_first_indices, + first_subseq_indices=first_subseq_indices, + last_subseq_indices=last_subseq_indices, + num_non_first=num_non_first, + merge_seq_starts=merge_seq_starts, + merge_seq_counts=merge_seq_counts, + merge_init_offsets=merge_init_offsets, + ) + + +def intracard_pre_scan( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + gk: torch.Tensor | None, + cu_seqlens_subseq_split: torch.Tensor, + S_split: int, + chunk_size: int = 64, +) -> torch.Tensor: + """Compute packed (he, m) exit state for each sub-sequence. + + Returns hm [S_split, H, K, V+K] fp32 where columns [0:V]=he, [V:V+K]=m. + Dispatches to the SM-appropriate pre_scan implementation via get_pre_scan(). + """ + chunk_delta_rule_pre_scan = get_pre_scan(k.device) + return chunk_delta_rule_pre_scan( + k=k, + w=w, + u=u, + gk=gk, + cu_seqlens_split=cu_seqlens_subseq_split, + S_split=S_split, + chunk_size=chunk_size, + ) + + +def intracard_merge( + hm: torch.Tensor, + split_info: SplitSeqInfo, + num_non_first: int, + merge_seq_starts: list[int], + merge_seq_counts: list[int], + merge_init_offsets: list[int], + device: torch.device, + initial_state: torch.Tensor | None = None, +) -> tuple[torch.Tensor | None, int]: + """Prefix scan across sub-sequences to produce per-sub-sequence initial states. + + For split seq [s0, s1, ..., s_{n-1}]: h0_sj = m_{j-1} @ h0_{j-1} + he_{j-1}. + Returns (initial_states_merge [num_non_first, H, K, V] fp32, num_non_first). + """ + from cula.ops.cp.merge import merge_fwd + + if num_non_first == 0: + return None, 0 + + initial_states_merge = merge_fwd( + hm=hm, + seq_starts=merge_seq_starts, + seq_counts=merge_seq_counts, + init_offsets=merge_init_offsets, + split_seq_ids=split_info.split_seq_ids, + h0=initial_state, + num_non_first=num_non_first, + ) + + return initial_states_merge, num_non_first + + +def _scatter_initial_states( + initial_state: torch.Tensor | None, + initial_states_merge: torch.Tensor | None, + num_non_first: int, + total_subseqs: int, + first_subseq_indices: torch.Tensor, + non_first_indices: torch.Tensor, + H: int, + K: int, + V: int, + device: torch.device, +) -> torch.Tensor: + """Build initial_state_expanded [total_subseqs, H, K, V] for all sub-sequences.""" + initial_state_expanded = torch.zeros(total_subseqs, H, K, V, device=device, dtype=torch.float32) + + if initial_state is not None: + initial_state_expanded[first_subseq_indices] = initial_state + + if initial_states_merge is not None and num_non_first > 0: + initial_state_expanded[non_first_indices] = initial_states_merge + + return initial_state_expanded + + +def _gather_final_states( + final_state_subseq: torch.Tensor | None, + last_subseq_indices: torch.Tensor, + output_final_state: bool, +) -> torch.Tensor | None: + """Gather final state from last sub-sequence of each original sequence.""" + if not output_final_state or final_state_subseq is None: + return None + return final_state_subseq[last_subseq_indices] + + +def intracard_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + max_splits: int = 32, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Intra-card CP chunk_delta_h forward; drop-in replacement for chunk_gated_delta_rule_fwd_h. + + Splits long sequences, runs pre_scan → merge → fwd_h on sub-sequences. + Falls back to the non-CP path when guards indicate no benefit. + """ + assert cu_seqlens is not None, "intracard_fwd_h requires cu_seqlens (varlen mode)" + + _, _, H, K = k.shape + V = u.shape[3] + device = k.device + num_sms = get_device_sm_count(device) + + if cu_seqlens_cpu is None: + cu_seqlens_cpu = cu_seqlens.cpu() + + if not should_use_intracard_cp(cu_seqlens_cpu, num_sms, H, chunk_size): + return _get_fwd_h()( + k=k, + w=w, + u=u, + gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + _no_cp=True, + ) + + cu_list = cu_seqlens_cpu.tolist() + num_seqs = len(cu_list) - 1 + max_seq_len = max(cu_list[i + 1] - cu_list[i] for i in range(num_seqs)) + subseq_len = compute_subseq_len(max_seq_len, num_sms, H, chunk_size, num_seqs=num_seqs) + + cached = None + cache_key = (id(cu_seqlens), subseq_len, chunk_size, max_splits, str(device)) + with _intracard_cache_lock: + cached = _intracard_cache.get(cache_key) + if cached is not None: + if cached.cu_seqlens_ref() is cu_seqlens: + _intracard_cache.move_to_end(cache_key) + else: + _intracard_cache.pop(cache_key, None) + cached = None + + if cached is None: + cu_seqlens_subseq_values, split_info, total_subseqs = prepare_subseq_cu_seqlens( + cu_seqlens_cpu, subseq_len, chunk_size, max_splits=max_splits + ) + else: + split_info = cached.split_info + total_subseqs = cached.total_subseqs + + # Post-split occupancy guard (total_subseqs only known after prepare_subseq_cu_seqlens) + if split_info and total_subseqs * NUM_V_BLOCKS * H > num_sms: + split_info = False + + if not split_info: + return _get_fwd_h()( + k=k, + w=w, + u=u, + gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + _no_cp=True, + ) + + N_orig = len(cu_seqlens_cpu) - 1 + + if cached is not None: + cu_seqlens_subseq_values = cached.cu_seqlens_subseq_values + total_subseqs = cached.total_subseqs + non_first_indices = cached.non_first_indices + first_subseq_indices = cached.first_subseq_indices + last_subseq_indices = cached.last_subseq_indices + num_non_first = cached.num_non_first + merge_seq_starts = cached.merge_seq_starts + merge_seq_counts = cached.merge_seq_counts + merge_init_offsets = cached.merge_init_offsets + cu_seqlens_subseq_gpu = cached.cu_seqlens_subseq_gpu + chunk_indices_subseq = cached.chunk_indices_subseq + else: + ( + non_first_indices, + first_subseq_indices, + last_subseq_indices, + num_non_first, + merge_seq_starts, + merge_seq_counts, + merge_init_offsets, + ) = _precompute_intracard_indices(split_info, N_orig) + + non_first_indices = torch.tensor(non_first_indices, dtype=torch.int64, device=device) + first_subseq_indices = torch.tensor(first_subseq_indices, dtype=torch.int64, device=device) + last_subseq_indices = torch.tensor(last_subseq_indices, dtype=torch.int64, device=device) + + cu_seqlens_subseq_gpu = torch.tensor(cu_seqlens_subseq_values, dtype=torch.int32, device=device) + chunk_indices_subseq = _prepare_chunk_indices(cu_seqlens_subseq_values, chunk_size, device) + + with _intracard_cache_lock: + _intracard_cache[cache_key] = _CacheEntry( + cu_seqlens_ref=weakref.ref(cu_seqlens), + cu_seqlens_subseq_values=cu_seqlens_subseq_values, + split_info=split_info, + total_subseqs=total_subseqs, + non_first_indices=non_first_indices, + first_subseq_indices=first_subseq_indices, + last_subseq_indices=last_subseq_indices, + num_non_first=num_non_first, + merge_seq_starts=merge_seq_starts, + merge_seq_counts=merge_seq_counts, + merge_init_offsets=merge_init_offsets, + cu_seqlens_subseq_gpu=cu_seqlens_subseq_gpu, + chunk_indices_subseq=chunk_indices_subseq, + ) + while len(_intracard_cache) > _INTRACARD_CACHE_MAXSIZE: + _intracard_cache.popitem(last=False) + + hm = intracard_pre_scan( + k=k, + w=w, + u=u, + gk=gk, + cu_seqlens_subseq_split=cu_seqlens_subseq_gpu, + S_split=total_subseqs, + chunk_size=chunk_size, + ) + + initial_states_merge, num_non_first = intracard_merge( + hm=hm, + split_info=split_info, + num_non_first=num_non_first, + merge_seq_starts=merge_seq_starts, + merge_seq_counts=merge_seq_counts, + merge_init_offsets=merge_init_offsets, + device=device, + initial_state=initial_state, + ) + + initial_state_expanded = _scatter_initial_states( + initial_state=initial_state, + initial_states_merge=initial_states_merge, + num_non_first=num_non_first, + total_subseqs=total_subseqs, + first_subseq_indices=first_subseq_indices, + non_first_indices=non_first_indices, + H=H, + K=K, + V=V, + device=device, + ) + + h, v_new, final_state_subseq = _get_fwd_h()( + k=k, + w=w, + u=u, + gk=gk, + initial_state=initial_state_expanded, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens_subseq_gpu, + chunk_indices=chunk_indices_subseq, + _no_cp=True, + ) + + final_state = _gather_final_states( + final_state_subseq=final_state_subseq, + last_subseq_indices=last_subseq_indices, + output_final_state=output_final_state, + ) + + return h, v_new, final_state diff --git a/cula/ops/cp/merge.py b/cula/ops/cp/merge.py new file mode 100644 index 00000000..1136b501 --- /dev/null +++ b/cula/ops/cp/merge.py @@ -0,0 +1,533 @@ +# Copyright (c) 2025 ANTGROUP. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Merge step for Intra-Card Context Parallel chunk_delta_h. + +Implements the prefix-scan merge: + For each original sequence split into sub-sequences [s0, s1, ..., s_{n-1}]: + h0_s0 = initial_state (or zero) + h0_s1 = m_s0 @ h0_s0 + he_s0 + ... + +Input: hm [S_split, H, K, V+K] fp32 — packed (he, m) from pre_scan +Output: h [num_non_first, H, K, V] fp32 +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm as _llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor, make_fake_stream +from cutlass.cutlass_dsl import T as _T + + +# --------------------------------------------------------------------------- +# Inline PTX helpers: SM80 warp-level TF32 MMA (mma.sync.m16n8k8.tf32.tf32.f32) +# --------------------------------------------------------------------------- +def _to_ir(v, loc=None, ip=None): + """Convert DSL Numeric to an MLIR Value; pass through if already a Value.""" + if hasattr(v, "ir_value"): + return v.ir_value(loc=loc, ip=ip) + return v + + +@cutlass.dsl_user_op +def _cvt_f32_to_tf32(f, *, loc=None, ip=None): + """Round-to-nearest convert fp32 -> tf32 (stored as i32 bit pattern).""" + f_ir = _to_ir(f, loc=loc, ip=ip) + result = _llvm.inline_asm( + _T.i32(), + [f_ir], + "cvt.rna.tf32.f32 $0, $1;", + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@cutlass.dsl_user_op +def _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """One mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 instruction. + + Inputs: + a0..a3: tf32 bits (Int32) — A fragment of 16x8 tile + b0..b1: tf32 bits (Int32) — B fragment of 8x8 tile + c0..c3: Float32 — accumulator in + Returns: + (d0, d1, d2, d3) Float32 — accumulator out + """ + ins = [ + _to_ir(a0, loc=loc, ip=ip), + _to_ir(a1, loc=loc, ip=ip), + _to_ir(a2, loc=loc, ip=ip), + _to_ir(a3, loc=loc, ip=ip), + _to_ir(b0, loc=loc, ip=ip), + _to_ir(b1, loc=loc, ip=ip), + _to_ir(c0, loc=loc, ip=ip), + _to_ir(c1, loc=loc, ip=ip), + _to_ir(c2, loc=loc, ip=ip), + _to_ir(c3, loc=loc, ip=ip), + ] + struct_ty = ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>") + ret = _llvm.inline_asm( + struct_ty, + ins, + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = _llvm.extractvalue(_T.f32(), ret, [0], loc=loc, ip=ip) + d1 = _llvm.extractvalue(_T.f32(), ret, [1], loc=loc, ip=ip) + d2 = _llvm.extractvalue(_T.f32(), ret, [2], loc=loc, ip=ip) + d3 = _llvm.extractvalue(_T.f32(), ret, [3], loc=loc, ip=ip) + return ( + cutlass.Float32(d0), + cutlass.Float32(d1), + cutlass.Float32(d2), + cutlass.Float32(d3), + ) + + +# --------------------------------------------------------------------------- +# Compile-time constants (thread/vector layout) +# --------------------------------------------------------------------------- +_BV_DEFAULT = 64 +_M_THR = 8 # threads along rows of the (K, BV) tile +_N_THR = 16 # threads along cols of the (K, BV) tile +_NUM_THREADS = _M_THR * _N_THR # 128 +_VEC = 4 # 128-bit vectorized fp32 cp.async + + +class ChunkDeltaRuleMerge: + """Prefix-scan merge kernel. + + H/K/V/BV kept as Python ints on ``self`` so layout construction is static. + """ + + def __init__(self, H: int, K: int, V: int, BV: int = _BV_DEFAULT, has_h0: int = 0): + assert V % BV == 0, f"V={V} not divisible by BV={BV}" + assert K % _M_THR == 0, f"K={K} not divisible by M_THR={_M_THR}" + assert BV % _N_THR == 0, f"BV={BV} not divisible by N_THR={_N_THR}" + assert (BV // _N_THR) == _VEC, f"BV/N_THR must equal VEC={_VEC}" + assert K % _N_THR == 0, f"K={K} not divisible by N_THR={_N_THR}" + assert (K // _N_THR) % _VEC == 0, "K/N_THR must be a multiple of VEC" + self.H = H + self.K = K + self.V = V + self.BV = BV + self.has_h0 = int(has_h0) + self.rows_per_thr = K // _M_THR + self.cols_per_thr = BV // _N_THR # == _VEC + self.num_v_tiles = V // BV + + # ------------------------------------------------------------------ + @cute.jit + def __call__( + self, + hm: cute.Tensor, + h_out: cute.Tensor, + h0: cute.Tensor, + seq_starts: cute.Tensor, + seq_counts: cute.Tensor, + init_offsets: cute.Tensor, + split_seq_ids: cute.Tensor, + num_split_seqs: cutlass.Int32, + stream: cuda.CUstream, + ): + # +8 fp32 pad on the leading dim to eliminate SMEM bank conflicts: + # without padding, row strides 128 / 64 are both multiples of 32 banks, + # causing 4-8-way conflicts in the mma fragment loads/stores. + _PAD: cutlass.Constexpr[int] = 8 + sM_layout = cute.make_layout((self.K, self.K), stride=(self.K + _PAD, 1)) + sHe_layout = cute.make_layout((self.K, self.BV), stride=(self.BV + _PAD, 1)) + sH_layout = cute.make_layout((self.K, self.BV), stride=(self.BV + _PAD, 1)) + + @cute.struct + class SharedStorage: + sM: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, cute.cosize(sM_layout)], + 128, + ] + sHe: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, cute.cosize(sHe_layout)], + 128, + ] + sH: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, cute.cosize(sH_layout)], + 128, + ] + + self.shared_storage_ty = SharedStorage + + # cp.async 128-bit vectorized copy atom (G->S loads). + copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=_VEC * 32, + ) + thr_layout = cute.make_layout((_M_THR, _N_THR), stride=(_N_THR, 1)) + val_layout = cute.make_layout((1, _VEC)) + tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + # Universal 128-bit copy atom (S->G stores) sharing the same T/V layout + # so gmem writes to h_out are coalesced (128B/warp). + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=_VEC * 32, + ) + tiled_store = cute.make_tiled_copy_tv(store_atom, thr_layout, val_layout) + + self.kernel( + hm, + h_out, + h0, + seq_starts, + seq_counts, + init_offsets, + split_seq_ids, + sM_layout, + sHe_layout, + sH_layout, + tiled_copy, + tiled_store, + ).launch( + grid=(self.num_v_tiles, num_split_seqs, self.H), + block=(_NUM_THREADS, 1, 1), + stream=stream, + ) + + # ------------------------------------------------------------------ + @cute.kernel + def kernel( + self, + hm: cute.Tensor, + h_out: cute.Tensor, + h0: cute.Tensor, + seq_starts: cute.Tensor, + seq_counts: cute.Tensor, + init_offsets: cute.Tensor, + split_seq_ids: cute.Tensor, + sM_layout: cute.Layout, + sHe_layout: cute.Layout, + sH_layout: cute.Layout, + tiled_copy: cute.TiledCopy, + tiled_store: cute.TiledCopy, + ): + tidx, _, _ = cute.arch.thread_idx() + i_v, i_seq, i_h = cute.arch.block_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage_ty) + sM = cute.make_tensor(storage.sM.data_ptr(), sM_layout) + sHe = cute.make_tensor(storage.sHe.data_ptr(), sHe_layout) + sH = cute.make_tensor(storage.sH.data_ptr(), sH_layout) + + thr_copy = tiled_copy.get_slice(tidx) + thr_store = tiled_store.get_slice(tidx) + + ss_start = seq_starts[i_seq] + n_ss = seq_counts[i_seq] + init_base = init_offsets[i_seq] + + t_m = tidx // _N_THR + t_n = tidx % _N_THR + + # --- Initialize sH from h0 or zero --- + if cutlass.const_expr(self.has_h0): + orig_id = split_seq_ids[i_seq] + g_full = h0[orig_id, i_h, None, None] # (K, V) + gH0_tile = cute.local_tile( + g_full, + tiler=(self.K, self.BV), + coord=(0, i_v), + ) + tAgH = thr_copy.partition_S(gH0_tile) + tAsH = thr_copy.partition_D(sH) + cute.copy(tiled_copy, tAgH, tAsH) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + else: + for i in cutlass.range_constexpr(self.rows_per_thr): + r = t_m + _M_THR * i + for c in cutlass.range_constexpr(self.cols_per_thr): + sH[r, t_n * _VEC + c] = cutlass.Float32(0.0) + cute.arch.barrier() + + # --- Main prefix-scan loop --- + # Pre-declare loop-scratch scalars so their dsl types are stable across + # the has_h0 / !has_h0 control-flow merge and into the dynamic loop. + r = t_m + out_idx = cutlass.Int32(0) + i_ss = cutlass.Int32(0) + # Number of BV-wide column tiles in b_m (K cols). + m_col_tiles: cutlass.Constexpr[int] = self.K // self.BV + for idx in cutlass.range(0, n_ss, unroll=0): + i_ss = ss_start + idx + + g_hm = hm[i_ss, i_h, None, None] # (K, V+K) + + # Load b_he [K, BV] from cols [i_v*BV, (i_v+1)*BV) of g_hm. + gHe_tile = cute.local_tile( + g_hm, + tiler=(self.K, self.BV), + coord=(0, i_v), + ) + tAgHe = thr_copy.partition_S(gHe_tile) + tAsHe = thr_copy.partition_D(sHe) + cute.copy(tiled_copy, tAgHe, tAsHe) + + # Load b_m [K, K] as m_col_tiles BV-wide tiles (cols V..V+K). + base_tile = self.num_v_tiles # col-tile index where m starts + for j in cutlass.range_constexpr(m_col_tiles): + gM_j = cute.local_tile( + g_hm, + tiler=(self.K, self.BV), + coord=(0, base_tile + j), + ) + sM_j = cute.local_tile( + sM, + tiler=(self.K, self.BV), + coord=(0, j), + ) + tAgM = thr_copy.partition_S(gM_j) + tAsM = thr_copy.partition_D(sM_j) + cute.copy(tiled_copy, tAgM, tAsM) + + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + # --- Compute new_b_h = b_m @ b_h + b_he via SM80 TF32 MMA --- + # Warp-level mma.sync.m16n8k8 tiling of the (K=128, BV=64) output. + # 4 warps per CTA: each warp owns rows [warp*32, warp*32 + 32) and + # all BV cols. Within a warp, 2 M-tiles × 8 N-tiles × 16 K-iters. + warp_id = tidx // 32 + lane = tidx % 32 + q = lane // 4 + rp = lane % 4 + + M_TILES: cutlass.Constexpr[int] = 2 # (warp rows = 32) / 16 + N_TILES: cutlass.Constexpr[int] = self.BV // 8 + K_TILES: cutlass.Constexpr[int] = self.K // 8 + + # Accumulator: [M_TILES, N_TILES, 4] fp32 per lane. + acc = cute.make_rmem_tensor( + cute.make_layout((M_TILES, N_TILES, 4)), + cutlass.Float32, + ) + # Initialize acc from sHe using the MMA D-fragment ownership. + for mi in cutlass.range_constexpr(M_TILES): + row_a = warp_id * 32 + mi * 16 + q + row_b = row_a + 8 + for nj in cutlass.range_constexpr(N_TILES): + col_a = nj * 8 + rp * 2 + acc[mi, nj, 0] = sHe[row_a, col_a] + acc[mi, nj, 1] = sHe[row_a, col_a + 1] + acc[mi, nj, 2] = sHe[row_b, col_a] + acc[mi, nj, 3] = sHe[row_b, col_a + 1] + + # K-reduction. For each k-tile: pre-cvt A (per M-tile) and B + # (per N-tile) once, then call 2*8 MMAs reusing them. + a_frag = cute.make_rmem_tensor( + cute.make_layout((M_TILES, 4)), + cutlass.Int32, # tf32 bits + ) + b_frag = cute.make_rmem_tensor( + cute.make_layout((N_TILES, 2)), + cutlass.Int32, # tf32 bits + ) + for ki in cutlass.range_constexpr(K_TILES): + k_base = ki * 8 + # Pre-load + cvt A. For m16n8k8 TF32, A[16x8] per-lane: + # a0: (q, rp), a1: (q+8, rp) + # a2: (q, rp+4), a3: (q+8, rp+4) + for mi in cutlass.range_constexpr(M_TILES): + row_a = warp_id * 32 + mi * 16 + q + row_b = row_a + 8 + a_frag[mi, 0] = _cvt_f32_to_tf32(sM[row_a, k_base + rp]) + a_frag[mi, 1] = _cvt_f32_to_tf32(sM[row_b, k_base + rp]) + a_frag[mi, 2] = _cvt_f32_to_tf32(sM[row_a, k_base + rp + 4]) + a_frag[mi, 3] = _cvt_f32_to_tf32(sM[row_b, k_base + rp + 4]) + # Pre-load + cvt B. For m16n8k8 TF32, B[8x8] per-lane (col-major): + # b0: (rp, q) + # b1: (rp+4, q) + for nj in cutlass.range_constexpr(N_TILES): + col_b = nj * 8 + q + b_frag[nj, 0] = _cvt_f32_to_tf32(sH[k_base + rp, col_b]) + b_frag[nj, 1] = _cvt_f32_to_tf32(sH[k_base + rp + 4, col_b]) + # MMAs + for mi in cutlass.range_constexpr(M_TILES): + for nj in cutlass.range_constexpr(N_TILES): + d0, d1, d2, d3 = _mma_m16n8k8_tf32( + a_frag[mi, 0], + a_frag[mi, 1], + a_frag[mi, 2], + a_frag[mi, 3], + b_frag[nj, 0], + b_frag[nj, 1], + acc[mi, nj, 0], + acc[mi, nj, 1], + acc[mi, nj, 2], + acc[mi, nj, 3], + ) + acc[mi, nj, 0] = d0 + acc[mi, nj, 1] = d1 + acc[mi, nj, 2] = d2 + acc[mi, nj, 3] = d3 + + # --- Write acc → sH (for next iter) and h_out (when not last) --- + cute.arch.barrier() + for mi in cutlass.range_constexpr(M_TILES): + row_a = warp_id * 32 + mi * 16 + q + row_b = row_a + 8 + for nj in cutlass.range_constexpr(N_TILES): + col_a = nj * 8 + rp * 2 + sH[row_a, col_a] = acc[mi, nj, 0] + sH[row_a, col_a + 1] = acc[mi, nj, 1] + sH[row_b, col_a] = acc[mi, nj, 2] + sH[row_b, col_a + 1] = acc[mi, nj, 3] + + if idx < n_ss - 1: + # Coalesced 128-bit stores from sH -> h_out via shared thread + # layout (matches loader). acc was already scattered to sH + # above, so read from sH (same barrier covers the hand-off). + cute.arch.barrier() + out_idx = init_base + idx + g_out = h_out[out_idx, i_h, None, None] # (K, V) + gOut_tile = cute.local_tile( + g_out, + tiler=(self.K, self.BV), + coord=(0, i_v), + ) + tSsH = thr_store.partition_S(sH) + tSgO = thr_store.partition_D(gOut_tile) + cute.copy(tiled_store, tSsH, tSgO) + + cute.arch.barrier() + + +# --------------------------------------------------------------------------- +# Compile cache +# --------------------------------------------------------------------------- +def _compile_merge_variant(H: int, K: int, V: int, has_h0: int): + kernel_obj = ChunkDeltaRuleMerge(H=H, K=K, V=V, BV=_BV_DEFAULT, has_h0=has_h0) + + sym_s = cute.sym_int() + sym_nnf = cute.sym_int() + sym_nss = cute.sym_int() + sym_nss1 = cute.sym_int() + sym_n = cute.sym_int() + + hm_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_s, H, K, V + K), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + h_out_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_nnf, H, K, V), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + h0_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_n, H, K, V), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + starts_fake = make_fake_compact_tensor(cutlass.Int32, (sym_nss,), assumed_align=16) + counts_fake = make_fake_compact_tensor(cutlass.Int32, (sym_nss,), assumed_align=16) + init_fake = make_fake_compact_tensor(cutlass.Int32, (sym_nss1,), assumed_align=16) + sid_fake = make_fake_compact_tensor(cutlass.Int32, (sym_nss,), assumed_align=16) + + stream_fake = make_fake_stream() + + return cute.compile( + kernel_obj, + hm_fake, + h_out_fake, + h0_fake, + starts_fake, + counts_fake, + init_fake, + sid_fake, + cutlass.Int32(1), + stream_fake, + ) + + +@functools.lru_cache(maxsize=32) +def _get_compiled_merge(H: int, K: int, V: int, has_h0: int): + return _compile_merge_variant(H, K, V, has_h0) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- +def merge_fwd( + hm: torch.Tensor, + seq_starts: list[int], + seq_counts: list[int], + init_offsets: list[int], + split_seq_ids: list[int], + h0: torch.Tensor | None, + num_non_first: int, +) -> torch.Tensor: + """Prefix-scan merge using a single CuTeDSL kernel launch.""" + assert hm.dtype == torch.float32, f"hm must be fp32, got {hm.dtype}" + _, H, K, VK = hm.shape + V = VK - K + device = hm.device + num_split_seqs = len(split_seq_ids) + + h_out = hm.new_empty(num_non_first, H, K, V) + + starts_gpu = torch.tensor(seq_starts, dtype=torch.int32, device=device) + counts_gpu = torch.tensor(seq_counts, dtype=torch.int32, device=device) + init_off_gpu = torch.tensor(init_offsets, dtype=torch.int32, device=device) + sid_gpu = torch.tensor(split_seq_ids, dtype=torch.int32, device=device) + + if h0 is not None: + h0_arg = h0 + has_h0 = 1 + else: + h0_arg = hm.new_zeros(1, H, K, V) + has_h0 = 0 + + compiled_fn = _get_compiled_merge(H, K, V, has_h0) + stream_ptr = torch.cuda.current_stream(device).cuda_stream + + compiled_fn( + from_dlpack(hm, assumed_align=128), + from_dlpack(h_out, assumed_align=128), + from_dlpack(h0_arg, assumed_align=128), + from_dlpack(starts_gpu, assumed_align=16), + from_dlpack(counts_gpu, assumed_align=16), + from_dlpack(init_off_gpu, assumed_align=16), + from_dlpack(sid_gpu, assumed_align=16), + cutlass.Int32(num_split_seqs), + cuda.CUstream(stream_ptr), + ) + + return h_out diff --git a/cula/ops/cp/pre_scan.py b/cula/ops/cp/pre_scan.py new file mode 100644 index 00000000..befef74d --- /dev/null +++ b/cula/ops/cp/pre_scan.py @@ -0,0 +1,1335 @@ +# Copyright (c) 2025 ANTGROUP. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Pre-Scan Kernel for Intra-Card Context Parallel chunk_delta_h. + +Single fused CuTeDSL kernel with grid-level dispatch: + blockIdx.x < num_v_tiles → he mode: computes he [K, V] = exit h-state + blockIdx.x >= num_v_tiles → m mode: computes m [K, K] = transition matrix + (8-warp SM100 MMA pipeline, identical MMA shapes for both modes) + +Output tensor: hm [S_split, H, K, V+K] fp32 + columns [0:V] = he (exit h-state) + columns [V:V+K] = m (transition matrix) +""" + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream +from cutlass.cute.typing import Float32, Int32, Int64 + +from cula.utils import USE_FAST_MATH, assert_blackwell + +PRINT_DEBUG = False + +LN2 = 0.6931471805599453 +INV_LN2 = 1.4426950408889634 + + +def make_thread_cooperative_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +# ===================================================================== +# Fused CuTeDSL Kernel: he + m with grid-level dispatch +# ===================================================================== + + +class ChunkDeltaRulePreScanFused: + """ + Fused pre-scan kernel: computes both he (exit h-state) and m (transition matrix). + + Grid-level dispatch: blockIdx.x < num_v_tiles → he mode, else → m mode. + Both modes share identical MMA structure (BS=BV=64, BT=64, BK=128). + MMA warp code is unchanged; only CUDA warps have mode-specific branches. + + Grid: (num_v_tiles + num_k_tiles, S_split * H, 1) — non-persistent. + Each CTA processes one (tile, sub-sequence, head) work unit. + """ + + def __init__( + self, + chunk_size: int = 64, + head_dim_k: int = 128, + head_dim_v: int = 128, + acc_dtype: type[cutlass.Numeric] = cutlass.Float32, + io_dtype: type[cutlass.Numeric] = cutlass.BFloat16, + use_fast_math: bool = True, + ): + assert head_dim_k == 128 and head_dim_v == 128 + assert_blackwell() + + self.use_fast_math = use_fast_math + self.chunk_size = chunk_size + self.head_dim_k = head_dim_k + self.head_dim_v = head_dim_v + self.acc_dtype = acc_dtype + self.io_dtype = io_dtype + + self.BT = chunk_size # 64 + self.BK = head_dim_k # 128 + self.BV = 64 # V tiling fixed at 64 + self.BS = 64 # K tiling for m mode (= BV) + + # Warp assignment (same as fwd_h) + self.threads_per_warp = 32 + self.cuda_warp_ids = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.load_warp_id = 5 + self.store_warp_id = 6 + self.empty_warp_id = 7 + self.min_occupancy = 1 + self.num_regs_cuda = 232 + self.num_regs_others = 40 + self.threads_per_cta = self.threads_per_warp * 8 + + # MMA tiling (same as fwd_h) + # WH MMA: state(BV,BK) @ W(BT,BK) → acc(BV,BT) + self.wh_mma_tiler = (self.BV, self.BT, self.BK) + # KV MMA: vnew(BV,BT) @ K^T(BK,BT) → update(BV,BK) + self.kv_mma_tiler = (self.BV, self.BK, self.BT) + + # Pipeline stages (simplified: no h_out, no vnew_store) + self.k_stage = 3 + self.w_stage = 3 + self.u_stage = 2 + self.gk_stage = 2 + self.acc_stage = 1 + self.cluster_shape_mnk = (1, 1, 1) + self.cta_group = tcgen05.CtaGroup.ONE + + self.buffer_align_bytes = 1024 + + # Barrier for TMEM dealloc sync + self.tmem_dealloc_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_cta, + ) + # Barrier for CUDA warp-group sync during gk_scale precomputation + self.gk_precompute_bar = pipeline.NamedBarrier( + barrier_id=3, + num_threads=self.threads_per_warp * len(self.cuda_warp_ids), # 128 + ) + + @staticmethod + def _plan_tmem_offsets(tiled_mma_wh, tile_wh, tiled_mma_kv, tile_kv, state_tmem_layout, vnew_tmem_layout, acc_stages): + """Plan TMEM column allocation. Same as fwd_h.""" + SM100_TMEM_CAPACITY_COLS = 512 + wh_shape = tiled_mma_wh.partition_shape_C(tile_wh[:2]) + wh_fake = tiled_mma_wh.make_fragment_C(cute.append(wh_shape, acc_stages)) + num_wh = tcgen05.find_tmem_tensor_col_offset(wh_fake) + + tCrState_fake = tiled_mma_wh.make_fragment_A(state_tmem_layout.outer.shape) + num_state = tcgen05.find_tmem_tensor_col_offset(tCrState_fake) + + tCrVnew_fake = tiled_mma_kv.make_fragment_A(vnew_tmem_layout.outer.shape) + num_vnew = tcgen05.find_tmem_tensor_col_offset(tCrVnew_fake) + + kv_shape = tiled_mma_kv.partition_shape_C(tile_kv[:2]) + kv_fake = tiled_mma_kv.make_fragment_C(cute.append(kv_shape, 1)) + num_kv = tcgen05.find_tmem_tensor_col_offset(kv_fake) + + wh_off = 0 + state_off = wh_off + num_wh + vnew_off = state_off + num_state + kv_off = vnew_off + num_vnew + total_tmp = kv_off + num_kv + total = 1 + while total < total_tmp: + total *= 2 + assert total <= SM100_TMEM_CAPACITY_COLS + return wh_off, state_off, vnew_off, kv_off, total + + def _compute_grid(self, S_split, H, K, V): + """Grid: (num_v_tiles + num_k_tiles, S_split * H, 1). Non-persistent.""" + num_v_tiles = (V + self.BV - 1) // self.BV + num_k_tiles = (K + self.BS - 1) // self.BS + return (num_v_tiles + num_k_tiles, S_split * H, 1) + + def _tma_partition_B(self, tma_atom, tma_tensor, smem, tile_shape, tiled_mma, batch_idx, hidx): + """Partition B operand tensors for TMA copy.""" + coord = (0, None, None) + gX = cute.local_tile(tma_tensor, cute.slice_(tile_shape, coord), (None, None, (hidx, batch_idx))) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_B(gX) + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 3), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + @cute.jit + def _epilog_partition(self, atom, gC_mnl, epi_tile, sC): + """Partition for epilogue-style TMA load.""" + gC_epi = cute.flat_divide(gC_mnl, epi_tile) + sC_g = cute.group_modes(sC, 0, 2) + gC_g = cute.group_modes(gC_epi, 0, 2) + bSG_sC, bSG_gC = cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + sC_g, + gC_g, + ) + return atom, bSG_sC, bSG_gC + + @cute.jit + def __call__( + self, + # ── Input tensors (varlen packed, B=1) ── + k_in: cute.Tensor, # [T_total, H, K] bf16 + w_in: cute.Tensor, # [T_total, H, K] bf16 + u_in: cute.Tensor, # [T_total, H, V] bf16 + gk_in: cute.Tensor, # [T_total, H, K] fp32 + # ── Output tensor ── + hm_in: cute.Tensor, # [S_split, H, K, V+K] fp32 (packed he+m) + # ── Sequence metadata ── + cu_seqlens_in: cute.Tensor, # [S_split+1] int32 + # ── Scalar parameters ── + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], # (S_split, T_total, H, K, V) + use_gk: Int32, # 1 if gk is provided, 0 otherwise + num_v_tiles: Int32, # cdiv(V, BV) — dispatch threshold + stream, + ): + """ + Launch the pre-scan kernel. + + Args: + k_in: key tensor, varlen packed [T_total, H, K] bf16 + w_in: decay weight tensor [T_total, H, K] bf16 + u_in: value tensor [T_total, H, V] bf16 + gk_in: key gate [T_total, H, K] fp32 (zeros if unused) + hm_in: output tensor [S_split, H, K, V+K] fp32 + he written to columns [0:V], m written to columns [V:V+K] + cu_seqlens_in: cumulative sequence lengths [S_split+1] int32 + problem_size: (S_split, T_total, H, K, V) + use_gk: flag for gk gating + num_v_tiles: number of V tiles (dispatch threshold for he vs m) + """ + k_ptr = k_in.iterator + w_ptr = w_in.iterator + u_ptr = u_in.iterator + gk_ptr = gk_in.iterator + hm_ptr = hm_in.iterator + cu_seqlens_ptr = cu_seqlens_in.iterator + + S_split, T_total, H, K, V = problem_size + + # ===================== GMEM layouts ===================== + # All data tensors are varlen packed [T_total, H, dim] + # K^T view: (K, T, (H, 1)) with K contiguous — for KV MMA B operand + kt_layout = cute.make_layout((K, T_total, (H, Int32(1))), stride=(1, H * K, (K, T_total * H * K))) + kt = cute.make_tensor(k_ptr, kt_layout) + + # W view: (T, K, (H, 1)) with K contiguous — for WH MMA B operand + w_layout = cute.make_layout((T_total, K, (H, Int32(1))), stride=(H * K, 1, (K, T_total * H * K))) + w = cute.make_tensor(w_ptr, w_layout) + + # U transposed view: (V, T, (H, 1)) with V contiguous — for TMA load + u_T_layout = cute.make_layout((V, T_total, (H, Int32(1))), stride=(1, H * V, (V, T_total * H * V))) + u_T = cute.make_tensor(u_ptr, u_T_layout) + + # U row-major view: (T, V, H) — for address computation in CUDA warps + u_layout = cute.make_layout((T_total, V, H), stride=(H * V, 1, V)) + u = cute.make_tensor(u_ptr, u_layout) + + # gk K-first view: (K, T_gk, (H, 1)) with K contiguous — for TMA load + # T_gk = 1 when gk is unused (dummy 1-row tensor), T_total otherwise + T_gk = gk_in.shape[0] + gk_K_layout = cute.make_layout((K, T_gk, (H, Int32(1))), stride=(1, H * K, (K, T_gk * H * K))) + gk_K = cute.make_tensor(gk_ptr, gk_K_layout) + + # he output: writes columns [0:V] of packed [S_split, H, K, V+K] + he_layout = cute.make_layout( + (K, V, (H, S_split)), + stride=(V + K, 1, (K * (V + K), H * K * (V + K))), + ) + he = cute.make_tensor(hm_ptr, he_layout) + + # m output: writes columns [V:V+K] of packed [S_split, H, K, V+K] + m_layout = cute.make_layout( + (K, K, (H, S_split)), + stride=(V + K, 1, (K * (V + K), H * K * (V + K))), + ) + m = cute.make_tensor(hm_ptr + V, m_layout) + + # cu_seqlens: [S_split+1] + cu_seqlens = cute.make_tensor(cu_seqlens_ptr, cute.make_layout((S_split + 1,))) + + self.k_dtype = kt.element_type + self.w_dtype = w.element_type + self.u_dtype = u.element_type + + # ===================== MMA setup (same as fwd_h) ===================== + wh_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + self.cta_group, + self.wh_mma_tiler[:2], + tcgen05.OperandSource.TMEM, + ) + kv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + self.cta_group, + self.kv_mma_tiler[:2], + tcgen05.OperandSource.TMEM, + ) + + vnew_tmem_layout = sm100_utils.make_smem_layout_a( + kv_tiled_mma, + self.kv_mma_tiler, + self.io_dtype, + 1, + ) + state_tmem_layout = sm100_utils.make_smem_layout_a( + wh_tiled_mma, + self.wh_mma_tiler, + self.io_dtype, + 1, + ) + + # ===================== TMEM offsets ===================== + (self.tmem_wh_off, self.tmem_state_off, self.tmem_vnew_off, self.tmem_kv_off, self.tmem_total) = ( + self._plan_tmem_offsets( + wh_tiled_mma, + self.wh_mma_tiler, + kv_tiled_mma, + self.kv_mma_tiler, + state_tmem_layout, + vnew_tmem_layout, + self.acc_stage, + ) + ) + + # ===================== SMEM layouts ===================== + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + + w_smem_staged = sm100_utils.make_smem_layout_b( + wh_tiled_mma, + self.wh_mma_tiler, + self.io_dtype, + self.w_stage, + ) + kt_smem_staged = sm100_utils.make_smem_layout_b( + kv_tiled_mma, + self.kv_mma_tiler, + self.io_dtype, + self.k_stage, + ) + u_epi_staged = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.COL_MAJOR, + (self.BV, self.BT), + self.u_stage, + ) + + # ===================== TMA descriptors ===================== + cluster_layout = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (wh_tiled_mma.thr_id.shape,), + ) + + w_smem = cute.select(w_smem_staged, mode=[0, 1, 2]) + tma_atom_w, tma_tensor_w = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + w, + w_smem, + self.wh_mma_tiler, + wh_tiled_mma, + cluster_layout.shape, + ) + kt_smem = cute.select(kt_smem_staged, mode=[0, 1, 2]) + tma_atom_kt, tma_tensor_kt = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + kt, + kt_smem, + self.kv_mma_tiler, + kv_tiled_mma, + cluster_layout.shape, + ) + u_smem = cute.select(u_epi_staged, mode=[0, 1]) + tma_atom_u, tma_tensor_u = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, + u_T, + u_smem, + (self.BV, self.BT), + ) + gk_smem_2d = cute.make_layout((self.BK, 1)) + tma_atom_gk, tma_tensor_gk = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, + gk_K, + gk_smem_2d, + (self.BK, 1), + ) + + self.tma_w_bytes = cute.size_in_bytes(self.io_dtype, w_smem) + self.tma_kt_bytes = cute.size_in_bytes(self.io_dtype, kt_smem) + self.tma_u_bytes = cute.size_in_bytes(self.io_dtype, u_smem) + self.tma_gk_bytes = self.BK * 4 + + # ===================== SharedStorage ===================== + @cute.struct + class SharedStorage: + # -- Pipelines: Load → MMA -- + load_w_mbar: cute.struct.MemRange[Int64, self.w_stage * 2] + load_kt_mbar: cute.struct.MemRange[Int64, self.k_stage * 2] + load_u_mbar: cute.struct.MemRange[Int64, self.u_stage * 2] + load_gk_mbar: cute.struct.MemRange[Int64, self.gk_stage * 2] + # -- Pipelines: CUDA ↔ MMA -- + state_tmem_mbar: cute.struct.MemRange[Int64, 1 * 2] + wh_done_mbar: cute.struct.MemRange[Int64, self.acc_stage * 2] + vnew_smem_mbar: cute.struct.MemRange[Int64, 1 * 2] + kv_done_mbar: cute.struct.MemRange[Int64, 1 * 2] + + # -- TMEM holding -- + tmem_holding_buf: Int32 + # -- Data buffers -- + sW: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(w_smem_staged)], + self.buffer_align_bytes, + ] + sKt: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(kt_smem_staged)], + self.buffer_align_bytes, + ] + sU: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(u_epi_staged)], + self.buffer_align_bytes, + ] + sGK: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BK * self.gk_stage], + 128, + ] + + self.shared_storage = SharedStorage + self.grid = self._compute_grid(S_split, H, K, V) + + self.kernel( + wh_tiled_mma, + kv_tiled_mma, + tma_atom_w, + tma_tensor_w, + tma_atom_kt, + tma_tensor_kt, + tma_atom_u, + tma_tensor_u, + tma_atom_gk, + tma_tensor_gk, + u, + u_T, + he, + m, + w_smem_staged, + kt_smem_staged, + state_tmem_layout, + vnew_tmem_layout, + u_epi_staged, + cu_seqlens, + problem_size, + use_gk, + num_v_tiles, + ).launch( + grid=self.grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=self.min_occupancy, + ) + + @cute.kernel + def kernel( + self, + wh_tiled_mma: cute.TiledMma, + kv_tiled_mma: cute.TiledMma, + # TMA atoms + descriptors + tma_atom_w: cute.CopyAtom, + tma_tensor_w: cute.Tensor, + tma_atom_kt: cute.CopyAtom, + tma_tensor_kt: cute.Tensor, + tma_atom_u: cute.CopyAtom, + tma_tensor_u: cute.Tensor, + tma_atom_gk: cute.CopyAtom, + tma_tensor_gk: cute.Tensor, + # GMEM tensors for address computation + u_tensor: cute.Tensor, # (T, V, H) + u_T_tensor: cute.Tensor, # (V, T, H) + he_tensor: cute.Tensor, # (K, V, (H, S_split)) — he columns of packed hm + m_tensor: cute.Tensor, # (K, K, (H, S_split)) — m columns of packed hm + # SMEM layouts + w_smem_staged: cute.ComposedLayout, + kt_smem_staged: cute.ComposedLayout, + state_tmem_layout: cute.ComposedLayout, + vnew_tmem_layout: cute.ComposedLayout, + u_epi_staged: cute.ComposedLayout, + # Sequence metadata + cu_seqlens: cute.Tensor, # (S_split+1,) + # Scalars + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + use_gk: Int32, + num_v_tiles: Int32, # dispatch: tile_idx < num_v_tiles → he mode + ): + """ + Device kernel. Each CTA processes one (tile, sub-seq, head) triple. + + Grid-level dispatch: + tile_idx < num_v_tiles → he mode (exit h-state) + tile_idx >= num_v_tiles → m mode (transition matrix) + + Both modes share identical MMA structure. Only CUDA warps and + Load warp (U TMA) differ between modes. + + Warp roles: + Load warp (5): TMA G2S for W, K^T, gk, U(he mode only). + MMA warp (4): WH/WM + KV/KM MMA (code unchanged). + CUDA warps (0-3): + he mode: h recursion (same as fwd_h minus outputs) + m mode: M^T recursion via associativity reformulation + Store warp (6): idle. + Empty warp (7): idle. + """ + S_split, T_total, H, K, V = problem_size + BT = self.BT + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + + # Prefetch TMA descriptors (Load warp) + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_w) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_kt) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_u) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_gk) + + # ===================== SMEM allocation ===================== + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sGK_smem = storage.sGK.get_tensor(cute.make_layout((self.BK, self.gk_stage))) + sGK_3d = storage.sGK.get_tensor(cute.make_layout((self.BK, 1, self.gk_stage), stride=(1, self.BK, self.BK))) + + # ===================== Pipelines ===================== + # Load → MMA: W, K^T (TmaUmma) + load_w_P, load_w_C = pipeline.PipelineTmaUmma.create( + num_stages=self.w_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(1), + tx_count=self.tma_w_bytes, + barrier_storage=storage.load_w_mbar.data_ptr(), + ).make_participants() + + load_kt_P, load_kt_C = pipeline.PipelineTmaUmma.create( + num_stages=self.k_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(1), + tx_count=self.tma_kt_bytes, + barrier_storage=storage.load_kt_mbar.data_ptr(), + ).make_participants() + + # CUDA → MMA: state TMEM (AsyncUmma) + state_smem_P, state_smem_C = pipeline.PipelineAsyncUmma.create( + num_stages=1, + producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.cuda_warp_ids)), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.state_tmem_mbar.data_ptr(), + ).make_participants() + + # MMA → CUDA: WH done (UmmaAsync) + wh_done_P, wh_done_C = pipeline.PipelineUmmaAsync.create( + num_stages=self.acc_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.cuda_warp_ids)), + barrier_storage=storage.wh_done_mbar.data_ptr(), + ).make_participants() + + # CUDA → MMA: vnew TMEM (AsyncUmma) + vnew_smem_P, vnew_smem_C = pipeline.PipelineAsyncUmma.create( + num_stages=1, + producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.cuda_warp_ids)), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.vnew_smem_mbar.data_ptr(), + ).make_participants() + + # MMA → CUDA: KV done (UmmaAsync) + kv_done_P, kv_done_C = pipeline.PipelineUmmaAsync.create( + num_stages=1, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.cuda_warp_ids)), + barrier_storage=storage.kv_done_mbar.data_ptr(), + ).make_participants() + + # Load → CUDA: U (TmaAsync) + load_u_P, load_u_C = pipeline.PipelineTmaAsync.create( + num_stages=self.u_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len(self.cuda_warp_ids)), + tx_count=self.tma_u_bytes, + barrier_storage=storage.load_u_mbar.data_ptr(), + ).make_participants() + + # Load → CUDA: gk (TmaAsync) + load_gk_P, load_gk_C = pipeline.PipelineTmaAsync.create( + num_stages=self.gk_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len(self.cuda_warp_ids)), + tx_count=self.tma_gk_bytes, + barrier_storage=storage.load_gk_mbar.data_ptr(), + ).make_participants() + + # ===================== TMEM allocation ===================== + tmem_alloc_bar = pipeline.NamedBarrier(barrier_id=1, num_threads=self.threads_per_cta) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_bar, + allocator_warp_id=self.load_warp_id, + ) + tmem.allocate(self.tmem_total) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + # ===================== SMEM views ===================== + sW = storage.sW.get_tensor(w_smem_staged.outer, swizzle=w_smem_staged.inner) + sKt = storage.sKt.get_tensor(kt_smem_staged.outer, swizzle=kt_smem_staged.inner) + sU_epi = storage.sU.get_tensor(u_epi_staged.outer, swizzle=u_epi_staged.inner) + + # ===================== MMA fragments ===================== + # WH MMA: A=state(TMEM), B=sW, acc=WH TMEM + tCrState_fake = wh_tiled_mma.make_fragment_A(state_tmem_layout.outer.shape) + tCrState = cute.make_tensor( + cute.recast_ptr(tmem_ptr + self.tmem_state_off, dtype=tCrState_fake.element_type), + tCrState_fake.layout, + ) + tCrW = wh_tiled_mma.make_fragment_B(sW) + wh_shape = wh_tiled_mma.partition_shape_C(self.wh_mma_tiler[:2]) + tCtAccWH_fake = wh_tiled_mma.make_fragment_C(cute.append(wh_shape, self.acc_stage)) + tCtAccWH = cute.make_tensor(tmem_ptr + self.tmem_wh_off, tCtAccWH_fake.layout) + + # KV MMA: A=v_new(TMEM), B=sKt, acc=KV TMEM + tCrVnew_fake = kv_tiled_mma.make_fragment_A(vnew_tmem_layout.outer.shape) + tCrVnew = cute.make_tensor( + cute.recast_ptr(tmem_ptr + self.tmem_vnew_off, dtype=tCrVnew_fake.element_type), + tCrVnew_fake.layout, + ) + tCrKt = kv_tiled_mma.make_fragment_B(sKt) + kv_shape = kv_tiled_mma.partition_shape_C(self.kv_mma_tiler[:2]) + tCtAccKV_fake = kv_tiled_mma.make_fragment_C(cute.append(kv_shape, 1)) + tCtAccKV = cute.make_tensor(tmem_ptr + self.tmem_kv_off, tCtAccKV_fake.layout) + + # ===================== Work unit decode (non-persistent) ===================== + # Release references to non-serializable Python objects before runtime if-blocks + del storage, smem + tile_idx = cute.arch.block_idx()[0] + combined = cute.arch.block_idx()[1] + i_subseq = combined // H + i_h = combined % H + bos = cu_seqlens[i_subseq] + eos = cu_seqlens[i_subseq + 1] + seq_len = eos - bos + NT = (seq_len + BT - 1) // BT + + # Grid-level dispatch: he mode vs m mode + is_he_mode = tile_idx < num_v_tiles + + # ========================================================================= + # LOAD WARP + # ========================================================================= + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_others) + + # TMA partition: shift by bos for varlen + tma_tensor_w_v = cute.domain_offset((bos, 0, (0, 0)), tma_tensor_w) + tma_tensor_kt_v = cute.domain_offset((0, bos, (0, 0)), tma_tensor_kt) + tma_tensor_u_v = cute.domain_offset((0, bos, (0, 0)), tma_tensor_u) + tma_tensor_gk_v = cute.domain_offset((0, bos, (0, 0)), tma_tensor_gk) + + tWsW, tWgW = self._tma_partition_B( + tma_atom_w, + tma_tensor_w_v, + sW, + self.wh_mma_tiler, + wh_tiled_mma, + Int32(0), + i_h, + ) + tKsK, tKgK = self._tma_partition_B( + tma_atom_kt, + tma_tensor_kt_v, + sKt, + self.kv_mma_tiler, + kv_tiled_mma, + Int32(0), + i_h, + ) + + # U TMA partition + gU_ld = tma_tensor_u_v[None, None, (i_h, Int32(0))] + _, bSG_sU, bSG_gU = self._epilog_partition( + tma_atom_u, + gU_ld, + (self.BV, self.BT), + sU_epi, + ) + + # gk TMA partition + gGK_ld = tma_tensor_gk_v[None, None, (i_h, Int32(0))] + _, bSG_sGK, bSG_gGK = self._epilog_partition( + tma_atom_gk, + gGK_ld, + (self.BK, 1), + sGK_3d, + ) + + # Chunk loop: issue TMA loads + for chunk_idx in cutlass.range(0, NT, unroll=0): + w_h = load_w_P.acquire_and_advance() + cute.copy( + atom=tma_atom_w, + src=tWgW[None, chunk_idx, 0], + dst=tWsW[None, w_h.index], + tma_bar_ptr=w_h.barrier, + ) + + kt_h = load_kt_P.acquire_and_advance() + cute.copy( + atom=tma_atom_kt, + src=tKgK[None, 0, chunk_idx], + dst=tKsK[None, kt_h.index], + tma_bar_ptr=kt_h.barrier, + ) + + # U TMA: he mode only (m mode skips U entirely) + if is_he_mode: + u_h = load_u_P.acquire_and_advance() + cute.copy( + atom=tma_atom_u, + src=bSG_gU[(None, tile_idx, chunk_idx)], + dst=bSG_sU[None, u_h.index], + tma_bar_ptr=u_h.barrier, + ) + + # Load gk only when gk gating is active + if use_gk != 0: + gk_t_idx = chunk_idx * self.BT + self.BT - 1 + remaining = seq_len - chunk_idx * self.BT + if remaining < self.BT: + gk_t_idx = seq_len - 1 + gk_h = load_gk_P.acquire_and_advance() + cute.copy( + atom=tma_atom_gk, + src=bSG_gGK[(None, 0, gk_t_idx)], + dst=bSG_sGK[None, gk_h.index], + tma_bar_ptr=gk_h.barrier, + ) + + # ========================================================================= + # MMA WARP + # ========================================================================= + elif warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_others) + + for chunk_idx in cutlass.range(0, NT, unroll=0): + # WH MMA: acc = state @ W + state_h = state_smem_C.wait_and_advance() + w_h = load_w_C.wait_and_advance() + wh_h = wh_done_P.acquire_and_advance() + for kp in cutlass.range(cute.size(tCrW, mode=[2]), unroll_full=True): + wh_tiled_mma.set(tcgen05.Field.ACCUMULATE, cutlass.Boolean(kp != 0)) + cute.gemm( + wh_tiled_mma, + tCtAccWH[None, None, None, wh_h.index], + tCrState[None, None, kp, state_h.index], + tCrW[None, None, kp, w_h.index], + tCtAccWH[None, None, None, wh_h.index], + ) + wh_h.commit() + w_h.release() + state_h.release() + + # KV MMA: update = vnew @ K^T + vnew_h = vnew_smem_C.wait_and_advance() + kt_h = load_kt_C.wait_and_advance() + kv_h = kv_done_P.acquire_and_advance() + for kp in cutlass.range(cute.size(tCrKt, mode=[2]), unroll_full=True): + kv_tiled_mma.set(tcgen05.Field.ACCUMULATE, cutlass.Boolean(kp != 0)) + cute.gemm( + kv_tiled_mma, + tCtAccKV[None, None, None, 0], + tCrVnew[None, None, kp, vnew_h.index], + tCrKt[None, None, kp, kt_h.index], + tCtAccKV[None, None, None, 0], + ) + kv_h.commit() + kt_h.release() + vnew_h.release() + + # ========================================================================= + # CUDA CORE WARPS (0-3) + # ========================================================================= + elif warp_idx in self.cuda_warp_ids: + cute.arch.setmaxregister_increase(self.num_regs_cuda) + local_tidx = tidx % (self.threads_per_warp * len(self.cuda_warp_ids)) + + # ----- T2R setup for KV acc (BV, BK fp32) → h update ----- + t2r_atom_kv = cute.make_copy_atom( + tcgen05.Ld16x256bOp(tcgen05.Repetition(16), tcgen05.Pack.NONE), + self.acc_dtype, + ) + tCtAccKV_flat = tCtAccKV[((None, None), 0, 0, None)] + fake_sKV = cute.make_tensor( + cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem), + cute.dice(self.kv_mma_tiler, (1, 1, None)), + ) + tiled_t2r_kv = tcgen05.make_tmem_copy(t2r_atom_kv, tCtAccKV_flat[(None, None, 0)]) + thr_t2r_kv = tiled_t2r_kv.get_slice(local_tidx) + tTR_tKV = thr_t2r_kv.partition_S(tCtAccKV_flat) + tTR_sKV = thr_t2r_kv.partition_D(fake_sKV) + # h state in registers (persistent across chunks) + tTR_rKV = cute.make_rmem_tensor(tTR_sKV.shape, self.acc_dtype) + + # ----- T2R setup for WH acc (BV, BT fp32) → v_new ----- + t2r_atom_wh = cute.make_copy_atom( + tcgen05.Ld16x256bOp(tcgen05.Repetition(8), tcgen05.Pack.NONE), + self.acc_dtype, + ) + tCtAccWH_flat = tCtAccWH[((None, None), 0, 0, None)] + fake_sWH = cute.make_tensor( + cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem), + cute.dice(self.wh_mma_tiler, (1, 1, None)), + ) + tiled_t2r_wh = tcgen05.make_tmem_copy(t2r_atom_wh, tCtAccWH_flat[(None, None, 0)]) + thr_t2r_wh = tiled_t2r_wh.get_slice(local_tidx) + tTR_tWH = thr_t2r_wh.partition_S(tCtAccWH_flat) + tTR_sWH = thr_t2r_wh.partition_D(fake_sWH) + + # ----- R2T: h regs → TMEM for WH MMA A operand ----- + copy_atom_r2t_state = cute.make_copy_atom( + tcgen05.St16x128bOp(tcgen05.Repetition(16), tcgen05.Unpack.NONE), + self.io_dtype, + ) + tiled_r2t_state = tcgen05.make_tmem_copy(copy_atom_r2t_state, tCrState) + thr_r2t_state = tiled_r2t_state.get_slice(local_tidx) + r2t_state_shape = cute.slice_(thr_r2t_state.partition_S(tCrState).shape, (None, None, None, None, 0)) + tRT_tState = thr_r2t_state.partition_D(tCrState) + + # ----- R2T: v_new regs → TMEM for KV MMA A operand ----- + copy_atom_r2t_vnew = cute.make_copy_atom( + tcgen05.St16x128bOp(tcgen05.Repetition(8), tcgen05.Unpack.NONE), + self.io_dtype, + ) + tiled_r2t_vnew = tcgen05.make_tmem_copy(copy_atom_r2t_vnew, tCrVnew) + thr_r2t_vnew = tiled_r2t_vnew.get_slice(local_tidx) + r2t_vnew_shape = cute.slice_(thr_r2t_vnew.partition_S(tCrVnew).shape, (None, None, None, None, 0)) + tRT_tVnew = thr_r2t_vnew.partition_D(tCrVnew) + + # ----- Identity tensors for coordinate mapping ----- + vnew_tile = cute.dice(self.wh_mma_tiler, (1, 1, None)) # (BV, BT) + cM_vnew = cute.make_identity_tensor(vnew_tile) + tTR_cM = thr_t2r_wh.partition_D(cM_vnew) + + h_tile = cute.dice(self.kv_mma_tiler, (1, 1, None)) # (BV, BK) + cM_h = cute.make_identity_tensor(h_tile) + tTR_cM_h = thr_t2r_kv.partition_D(cM_h) + + # ----- Initialize state: h=0 (he mode) or M^T=I (m mode) ----- + if is_he_mode: + for ei in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + tTR_rKV[ei] = Float32(0.0) + else: + k_col_tile = tile_idx - num_v_tiles + for ei in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + v_coord, k_coord = tTR_cM_h[ei] + col_global = v_coord + k_col_tile * self.BS + if k_coord == col_global: + tTR_rKV[ei] = Float32(1.0) + else: + tTR_rKV[ei] = Float32(0.0) + + # ===== Main chunk loop ===== + for chunk_idx in cutlass.range(0, NT, unroll=0): + # ======================================== + # Phase 1: Publish state for WH/WM MMA + # ======================================== + tRT_rState = cute.make_rmem_tensor(r2t_state_shape, self.io_dtype) + h_vec = tTR_rKV.load() + h_vec_bf16 = h_vec.to(self.io_dtype) + + # R2T state → TMEM (triggers WH/WM MMA) + tRT_rState.store(h_vec_bf16) + state_h = state_smem_P.acquire_and_advance() + cute.copy(tiled_r2t_state, tRT_rState, tRT_tState[(None, None, None, None, 0)]) + cute.arch.fence_view_async_tmem_store() + state_h.commit() + + # Preload U from SMEM → registers (he mode only, overlapping WH MMA) + tTR_rU = cute.make_rmem_tensor(tTR_sWH.shape, self.acc_dtype) + if is_he_mode: + u_handle = load_u_C.wait_and_advance() + for ei in cutlass.range_constexpr(cute.size(tTR_cM)): + v_coord, t_coord = tTR_cM[ei] + tTR_rU[ei] = sU_epi[(v_coord, t_coord, u_handle.index)].to(self.acc_dtype) + u_handle.release() + + # ======================================== + # Phase 2: Process WH/WM result → triggers KV/KM MMA + # ======================================== + wh_h = wh_done_C.wait_and_advance() + tTR_rWH = cute.make_rmem_tensor(tTR_sWH.shape, self.acc_dtype) + cute.copy(tiled_t2r_wh, tTR_tWH[(None, None, None, wh_h.index)], tTR_rWH) + cute.arch.fence_view_async_tmem_load() + wh_h.release() + + if is_he_mode: + # he mode: v_new = u - WH + for ei in cutlass.range_constexpr(cute.size(tTR_rWH)): + tTR_rWH[ei] = tTR_rU[ei] - tTR_rWH[ei] + # else: m mode — tTR_rWH = WM result, used as-is for KM MMA + + # Varlen tail chunk zero mask (both modes) + valid_len_chunk = seq_len - chunk_idx * self.BT + if valid_len_chunk < self.BT: + for ei in cutlass.range_constexpr(cute.size(tTR_cM)): + v_coord, t_coord = tTR_cM[ei] + if t_coord >= valid_len_chunk: + tTR_rWH[ei] = Float32(0.0) + + # R2T vnew/temp → TMEM (triggers KV/KM MMA) + vnew_vec_bf16 = tTR_rWH.load().to(self.io_dtype) + tRT_rVnew = cute.make_rmem_tensor(r2t_vnew_shape, self.io_dtype) + tRT_rVnew.store(vnew_vec_bf16) + vnew_h = vnew_smem_P.acquire_and_advance() + cute.copy(tiled_r2t_vnew, tRT_rVnew, tRT_tVnew[(None, None, None, None, 0)]) + cute.arch.fence_view_async_tmem_store() + vnew_h.commit() + + # ======================================== + # Phase 3: gk decay (overlapping with KV/KM MMA) + # ======================================== + if use_gk != 0: + gk_h = load_gk_C.wait_and_advance() + gk_raw = sGK_smem[(tidx, gk_h.index)] + sGK_smem[(tidx, gk_h.index)] = cute.exp2(gk_raw, fastmath=self.use_fast_math) + self.gk_precompute_bar.arrive_and_wait() + for ei in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + v_coord, k_coord = tTR_cM_h[ei] + tTR_rKV[ei] = tTR_rKV[ei] * sGK_smem[(k_coord, gk_h.index)] + gk_h.release() + + # ======================================== + # Phase 4: KV/KM update + # ======================================== + kv_h = kv_done_C.wait_and_advance() + tTR_rUpdate = cute.make_rmem_tensor(tTR_sKV.shape, self.acc_dtype) + cute.copy(tiled_t2r_kv, tTR_tKV[(None, None, None, 0)], tTR_rUpdate) + cute.arch.fence_view_async_tmem_load() + kv_h.release() + + h_vec = tTR_rKV.load() + update_vec = tTR_rUpdate.load() + if is_he_mode: + tTR_rKV.store(h_vec + update_vec) # h += K^T @ v_new + else: + tTR_rKV.store(h_vec - update_vec) # M -= K^T @ (W @ M) + + # ===== After loop: write output to GMEM ===== + if is_he_mode: + # Write he (exit h-state) → hm[:, :, :, :V] + for ei in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + v_coord, k_coord = tTR_cM_h[ei] + he_tensor[(k_coord, v_coord + tile_idx * self.BV, (i_h, i_subseq))] = tTR_rKV[ei] + else: + # Write M^T (transition matrix, transposed) → hm[:, :, :, V:] + k_col_tile = tile_idx - num_v_tiles + for ei in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + v_coord, k_coord = tTR_cM_h[ei] + col_global = v_coord + k_col_tile * self.BS + m_tensor[(k_coord, col_global, (i_h, i_subseq))] = tTR_rKV[ei] + + # ========================================================================= + # STORE WARP + # ========================================================================= + elif warp_idx == self.store_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_others) + # Store warp idle — CUDA warps write hm directly to GMEM + pass + + # ========================================================================= + # EMPTY WARP + # ========================================================================= + else: + cute.arch.setmaxregister_decrease(self.num_regs_others) + # Empty warp idle + pass + + # ===================== TMEM dealloc ===================== + self.tmem_dealloc_sync_barrier.sync() + tmem.free(tmem_ptr) + + +# ===================================================================== +# Compile cache + Python API +# ===================================================================== + +_pre_scan_kernel_cache: dict = {} + + +def _compile_pre_scan_variant(H, K, V, chunk_size, use_fast_math): + """Compile one ChunkDeltaRulePreScanFused kernel variant.""" + kernel_obj = ChunkDeltaRulePreScanFused( + chunk_size=chunk_size, + head_dim_k=K, + head_dim_v=V, + use_fast_math=use_fast_math, + ) + + sym_t = cute.sym_int() # T_total + sym_s = cute.sym_int() # S_split + sym_cu = cute.sym_int() # cu_seqlens length = S_split+1 + + # varlen packed: [T_total, H, dim] + sym_gk = cute.sym_int() # independent: 1 when gk unused, T_total when used + + k_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, K), stride_order=(2, 1, 0), assumed_align=128) + w_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, K), stride_order=(2, 1, 0), assumed_align=128) + u_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, V), stride_order=(2, 1, 0), assumed_align=128) + gk_fake = make_fake_compact_tensor(cutlass.Float32, (sym_gk, H, K), stride_order=(2, 1, 0), assumed_align=128) + + # output: [S_split, H, K, V+K] fp32 (packed hm) + hm_fake = make_fake_compact_tensor(cutlass.Float32, (sym_s, H, K, V + K), stride_order=(3, 2, 1, 0), assumed_align=128) + + # cu_seqlens: [S_split+1] + cu_fake = make_fake_compact_tensor(cutlass.Int32, (sym_cu,), assumed_align=128) + + stream_fake = make_fake_stream(use_tvm_ffi_env_stream=True) + + compiled_fn = cute.compile( + kernel_obj, + k_fake, + w_fake, + u_fake, + gk_fake, + hm_fake, + cu_fake, + (Int32(1), Int32(1), Int32(H), Int32(K), Int32(V)), # problem_size + Int32(0), # use_gk + Int32(0), # num_v_tiles (concrete value passed at runtime) + stream_fake, + options="--enable-tvm-ffi", + ) + return compiled_fn + + +def _get_compiled_pre_scan(H, K, V, chunk_size): + """Get compiled pre-scan kernel with lazy compilation + caching.""" + key = (H, K, V, chunk_size, USE_FAST_MATH) + if key not in _pre_scan_kernel_cache: + _pre_scan_kernel_cache[key] = _compile_pre_scan_variant(H, K, V, chunk_size, USE_FAST_MATH) + return _pre_scan_kernel_cache[key] + + +# ===================================================================== +# Python API +# ===================================================================== + + +def chunk_delta_rule_pre_scan( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + gk: torch.Tensor | None = None, + cu_seqlens_split: torch.Tensor = None, + S_split: int = 0, + chunk_size: int = 64, +) -> torch.Tensor: + """ + Compute packed (he, m) state for each split sub-sequence. + + Single fused CuTeDSL kernel with grid-level dispatch: + blockIdx.x < num_v_tiles → he (exit h-state) → hm[:, :, :, :V] + blockIdx.x >= num_v_tiles → m (transition matrix) → hm[:, :, :, V:] + + Args: + k: [1, T, H, K] bf16 (varlen packed, B=1) + w: [1, T, H, K] bf16 + u: [1, T, H, V] bf16 + gk: [1, T, H, K] fp32 or None (key gate) + cu_seqlens_split: [S_split+1] int32 (sub-sequence boundaries) + S_split: number of sub-sequences + chunk_size: chunk size (default 64) + + Returns: + hm: [S_split, H, K, V+K] fp32 + hm[:, :, :, :V] = he (K×V exit h-state) + hm[:, :, :, V:] = m (K×K transition matrix) + """ + assert cu_seqlens_split is not None, "cu_seqlens_split is required" + assert k.shape[0] == 1, "pre_scan requires varlen mode (B=1)" + + T = k.shape[1] + H = k.shape[2] + K = k.shape[3] + V = u.shape[3] + device = k.device + + # Squeeze batch dim for kernel (varlen: [T, H, dim]) + k_kern = k[0] + w_kern = w[0] + u_kern = u[0] + + use_gk_flag = 1 if gk is not None else 0 + gk_kern = gk[0] if gk is not None else torch.zeros(1, H, K, device=device, dtype=torch.float32) + + # Ensure cu_seqlens is int32 + cu_seqlens_i32 = cu_seqlens_split.int() if cu_seqlens_split.dtype != torch.int32 else cu_seqlens_split + + # Allocate packed output: [S_split, H, K, V+K] fp32 + hm = torch.empty(S_split, H, K, V + K, device=device, dtype=torch.float32) + + # Single fused kernel: he + m via grid-level dispatch + BV = 64 + num_v_tiles = (V + BV - 1) // BV + + compiled_fn = _get_compiled_pre_scan(H, K, V, chunk_size) + compiled_fn( + k_kern, + w_kern, + u_kern, + gk_kern, + hm, + cu_seqlens_i32, + (S_split, T, H, K, V), + use_gk_flag, + num_v_tiles, + ) + + return hm + + +# ===================================================================== +# Reference Implementation + Main +# ===================================================================== + + +def reference_pre_scan(k, w, u, gk, cu_seqlens, S_split, chunk_size): + """Pure PyTorch reference: compute he and M for each sub-sequence.""" + H = k.shape[2] + K = k.shape[3] + V = u.shape[3] + BT = chunk_size + device = k.device + + hm = torch.zeros(S_split, H, K, V + K, device=device, dtype=torch.float32) + + for s in range(S_split): + bos = cu_seqlens[s].item() + eos = cu_seqlens[s + 1].item() + seq_len = eos - bos + NT = (seq_len + BT - 1) // BT + + for h in range(H): + h_state = torch.zeros(V, K, device=device, dtype=torch.float32) + M = torch.eye(K, device=device, dtype=torch.float32) + + for c in range(NT): + t_start = bos + c * BT + t_end = min(t_start + BT, eos) + actual_len = t_end - t_start + + k_chunk = k[0, t_start:t_end, h, :].float() + w_chunk = w[0, t_start:t_end, h, :].float() + u_chunk = u[0, t_start:t_end, h, :].float() + + if actual_len < BT: + k_chunk = torch.nn.functional.pad(k_chunk, (0, 0, 0, BT - actual_len)) + w_chunk = torch.nn.functional.pad(w_chunk, (0, 0, 0, BT - actual_len)) + u_chunk = torch.nn.functional.pad(u_chunk, (0, 0, 0, BT - actual_len)) + + gk_last_t = t_end - 1 + if gk is not None: + alpha = gk[0, gk_last_t, h, :].float().exp2() + else: + alpha = torch.ones(K, device=device, dtype=torch.float32) + + WH = h_state @ w_chunk.T + v_new = u_chunk.T - WH + h_state = h_state * alpha.unsqueeze(0) + update = v_new @ k_chunk + h_state = h_state + update + + KtW = k_chunk.T @ w_chunk + A_t = torch.diag(alpha) - KtW + M = A_t @ M + + hm[s, h, :, :V] = h_state.T + hm[s, h, :, V:] = M + + return hm + + +def main(): + import argparse + + parser = argparse.ArgumentParser(description="Pre-scan kernel test & benchmark") + parser.add_argument("--test", type=str, default="both", choices=["correctness", "benchmark", "both"]) + parser.add_argument("--S_split", type=int, default=4) + parser.add_argument("--T", type=int, default=4096) + parser.add_argument("--H", type=int, default=64) + parser.add_argument("--K", type=int, default=128) + parser.add_argument("--V", type=int, default=128) + parser.add_argument("--chunk_size", type=int, default=64) + args = parser.parse_args() + + S_split, T, H, K, V, BT = args.S_split, args.T, args.H, args.K, args.V, args.chunk_size + device = "cuda" + + # ===== Correctness ===== + if args.test in ("correctness", "both"): + configs = [ + ("basic (1 seq, 2 chunks, gk)", 1, 128, 4, True), + ("no_gk (1 seq, 1 chunk)", 1, 64, 2, False), + ("tail_chunk (T=100)", 1, 100, 2, True), + ("multi_subseq (3 seqs)", 3, 384, 4, True), + ("large (S=8, T=8192, H=64)", 8, 8192, 64, True), + ] + + all_pass = True + for name, s, t, h, use_gk in configs: + print(f"\n{'=' * 60}") + print(f"Test: {name} (S={s}, T={t}, H={h}, gk={use_gk})") + torch.manual_seed(42) + + # Build cu_seqlens: split T evenly into s sub-sequences + base_len = t // s + seq_lens = [base_len] * s + seq_lens[-1] = t - base_len * (s - 1) # remainder to last + cu = [0] + for sl in seq_lens: + cu.append(cu[-1] + sl) + cu_seqlens = torch.tensor(cu, device=device, dtype=torch.int32) + + k_t = torch.randn(1, t, h, K, device=device, dtype=torch.bfloat16) * 0.02 + w_t = torch.randn(1, t, h, K, device=device, dtype=torch.bfloat16) * 0.02 + u_t = torch.randn(1, t, h, V, device=device, dtype=torch.bfloat16) * 0.02 + gk_t = torch.randn(1, t, h, K, device=device, dtype=torch.float32) * 0.01 if use_gk else None + + hm_kernel = chunk_delta_rule_pre_scan(k_t, w_t, u_t, gk_t, cu_seqlens, S_split=s, chunk_size=BT) + hm_ref = reference_pre_scan(k_t, w_t, u_t, gk_t, cu_seqlens, s, BT) + + he_rel = (hm_kernel[:, :, :, :V] - hm_ref[:, :, :, :V]).abs().max().item() / ( + hm_ref[:, :, :, :V].abs().max().item() + 1e-8 + ) + m_rel = (hm_kernel[:, :, :, V:] - hm_ref[:, :, :, V:]).abs().max().item() / ( + hm_ref[:, :, :, V:].abs().max().item() + 1e-8 + ) + # m accumulates bf16 truncation over NT chunks; use 2% for large configs + he_tol, m_tol = 0.01, 0.02 + passed = he_rel < he_tol and m_rel < m_tol + all_pass = all_pass and passed + print(f" he rel err: {he_rel:.6e} m rel err: {m_rel:.6e} {'PASS' if passed else 'FAIL'}") + + print(f"\n{'=' * 60}") + print(f"{'ALL PASS' if all_pass else 'SOME FAILED'}") + + # ===== Benchmark ===== + if args.test in ("benchmark", "both"): + print(f"\n{'=' * 60}") + print(f"Benchmark: S_split={S_split}, T={T}, H={H}, K={K}, V={V}") + torch.manual_seed(999) + + base_len = T // S_split + seq_lens = [base_len] * S_split + seq_lens[-1] = T - base_len * (S_split - 1) + cu = [0] + for sl in seq_lens: + cu.append(cu[-1] + sl) + cu_seqlens = torch.tensor(cu, device=device, dtype=torch.int32) + + k_b = torch.randn(1, T, H, K, device=device, dtype=torch.bfloat16) * 0.02 + w_b = torch.randn(1, T, H, K, device=device, dtype=torch.bfloat16) * 0.02 + u_b = torch.randn(1, T, H, V, device=device, dtype=torch.bfloat16) * 0.02 + gk_b = torch.randn(1, T, H, K, device=device, dtype=torch.float32) * 0.01 + + def run_bench(): + chunk_delta_rule_pre_scan(k_b, w_b, u_b, gk_b, cu_seqlens, S_split=S_split, chunk_size=BT) + + # Warmup + for _ in range(3): + run_bench() + torch.cuda.synchronize() + + n_iter = 20 + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(n_iter): + run_bench() + end_event.record() + torch.cuda.synchronize() + elapsed_ms = start_event.elapsed_time(end_event) / n_iter + print(f" cuLA pre_scan: {elapsed_ms:.3f} ms") + + # FLA Triton kernel reference (call raw kernel directly) + try: + import triton + from fla.ops.cp.chunk_delta_h import pre_process_fwd_kernel_merged as fla_kernel + + BLOCK_SIZE_FLA = 32 if K <= 64 else 64 + BK1_FLA = triton.next_power_of_2(K) + fla_grid = (triton.cdiv(V, BLOCK_SIZE_FLA) + triton.cdiv(K, BLOCK_SIZE_FLA), S_split * H) + + # FLA expects [T, H, K/V] layout (no batch dim), HV=H for this case + k_fla = k_b[0] # [T, H, K] + w_fla = w_b[0] # [T, H, K] + u_fla = u_b[0] # [T, H, V] + gk_fla = gk_b[0] # [T, H, K] + hm_fla = torch.empty(S_split, H, K, V + K, device=device, dtype=torch.float32) + + def run_fla(): + fla_kernel[fla_grid]( + k=k_fla, + v=u_fla, + w=w_fla, + g=None, + gk=gk_fla, + hm=hm_fla, + cu_seqlens=cu_seqlens, + T=T, + H=H, + HV=H, + K=K, + V=V, + BT=BT, + BK1=BK1_FLA, + BLOCK_SIZE=BLOCK_SIZE_FLA, + USE_EXP2=True, + MULTI_SEQS=True, + ) + + for _ in range(3): + run_fla() + torch.cuda.synchronize() + start_event.record() + for _ in range(n_iter): + run_fla() + end_event.record() + torch.cuda.synchronize() + fla_ms = start_event.elapsed_time(end_event) / n_iter + print(f" FLA pre_scan: {fla_ms:.3f} ms") + print(f" Speedup vs FLA: {fla_ms / elapsed_ms:.2f}x") + except Exception as e: + print(f" FLA not available: {e}") + + +if __name__ == "__main__": + main() diff --git a/cula/utils.py b/cula/utils.py index eaa094da..8b8e0ab1 100644 --- a/cula/utils.py +++ b/cula/utils.py @@ -108,6 +108,32 @@ def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callabl ) +def get_pre_scan(device: torch.device | str | int | None = None) -> Callable: + """Return the appropriate ``chunk_delta_rule_pre_scan`` implementation for *device*. + + - sm100/sm103 (Blackwell) → cula.ops.cp.pre_scan (CuTeDSL SM100 kernel) + - sm90 (Hopper) → cula.ops.cp.pre_scan_sm90 (to be implemented) + + Args: + device: CUDA device to query. Defaults to the currently active device. + + Raises: + RuntimeError: If the device architecture is not supported. + """ + major, minor = get_device_sm_version(device) + if major == 10 and minor in (0, 3): + from cula.ops.cp.pre_scan import chunk_delta_rule_pre_scan + + return chunk_delta_rule_pre_scan + elif major == 9 and minor == 0: + raise NotImplementedError("The Hopper (SM90) implementation of pre_scan is not yet available.") + else: + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + f"Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) + + @cute.jit def print_tensor_2d(tensor: cute.Tensor): """ diff --git a/tests/test_intracard_cp.py b/tests/test_intracard_cp.py new file mode 100644 index 00000000..d478190a --- /dev/null +++ b/tests/test_intracard_cp.py @@ -0,0 +1,466 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# Licensed under the Apache License, Version 2.0. +"""Tests for intracard CP: dispatch routing + numerical accuracy. + +Two reference levels are used: + - cuLA no-CP baseline (same kernel, no CP scheduling) — verifies dispatch + plumbing and that CP scheduling is value-preserving. + - Pure-PyTorch fp32 reference — source of truth for kernel correctness; + any deviation here is a real CP / kernel bug, not a cross-impl gap. + +The CP path is exercised via two entry points: + - ``chunk_gated_delta_rule_fwd_h`` with ``CULA_INTRACARD_CP=1`` + inference_mode + - ``intracard_fwd_h`` (direct, bypasses the heuristic) +""" + +from __future__ import annotations + +import math +import os +import pathlib +import sys + +import pytest +import torch + +# Make cuLA importable when tests run from a fresh checkout (no `pip install -e`). +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h as fla_fwd_h # noqa: E402 +from fla.utils import assert_close # noqa: E402 (RMSE-relative + atol short-circuit + NaN check) + +from cula.ops.chunk_delta_h_sm100 import chunk_gated_delta_rule_fwd_h # noqa: E402 +from cula.ops.cp.chunk_delta_h import ( # noqa: E402 + compute_subseq_len, + intracard_fwd_h, + prepare_subseq_cu_seqlens, + should_use_intracard_cp, +) +from cula.utils import get_device_sm_count # noqa: E402 + +# Constants & tolerances — aligned with existing cuLA tests (see below). +BT, K, V = 64, 128, 128 +DEVICE = "cuda" +# Tolerances aligned with existing cuLA tests: +# * Same-kernel (CP scheduling only): torch.testing.assert_close(atol=1e-2, rtol=1e-2) +# — matches tests/test_chunk_delta_h.py CP block +# * Cross-impl / vs ref: fla.utils.assert_close(ratio=...) +# — matches tests/test_kda_compare_fla.py +ATOL_SAME_KERNEL = 1e-2 +RTOL_SAME_KERNEL = 1e-2 +RATIO_VS_REF = 0.005 # RMSE / RMS(ref) — matches FLA test_gated_delta.py fwd +RATIO_VS_FLA = 0.015 # cross-impl gap measured ~1.27% (TF32 MMA vs Triton fp32) +RATIO_STRESS = 1e-6 # deterministic re-run: drift would indicate race + + +pytestmark = [ + pytest.mark.sm100_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +# ============================== Helpers ============================== + + +def make_varlen_inputs(seq_lens, H, *, use_gk=False, use_h0=False, seed=42): + """Build varlen-packed B=1 inputs for chunk_gated_delta_rule_fwd_h.""" + total = sum(seq_lens) + N = len(seq_lens) + cu = [0] + for s in seq_lens: + cu.append(cu[-1] + s) + + torch.manual_seed(seed) + k = torch.randn(1, total, H, K, dtype=torch.bfloat16, device=DEVICE) * 0.02 + w = torch.randn(1, total, H, K, dtype=torch.bfloat16, device=DEVICE) * 0.02 + u = torch.randn(1, total, H, V, dtype=torch.bfloat16, device=DEVICE) * 0.02 + + gk = None + if use_gk: + gk = torch.zeros(1, total, H, K, dtype=torch.float32, device=DEVICE) + for i in range(N): + bos, eos = cu[i], cu[i + 1] + seg = torch.randn(1, eos - bos, H, K, dtype=torch.float32, device=DEVICE) * 0.01 + gk[:, bos:eos] = -torch.abs(seg).cumsum(dim=1) + + h0 = torch.randn(N, H, K, V, dtype=torch.float32, device=DEVICE) * 0.01 if use_h0 else None + return k, w, u, gk, h0, torch.tensor(cu, dtype=torch.int32, device=DEVICE) + + +def run_cula_no_cp(k, w, u, gk, h0, cu, **kw): + return chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=h0, + chunk_size=BT, + cu_seqlens=cu, + _no_cp=True, + **kw, + ) + + +def run_cula_cp(k, w, u, gk, h0, cu, **kw): + """Auto-dispatch via env + inference_mode.""" + old = os.environ.get("CULA_INTRACARD_CP") + os.environ["CULA_INTRACARD_CP"] = "1" + try: + with torch.inference_mode(): + return chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=h0, + chunk_size=BT, + cu_seqlens=cu, + **kw, + ) + finally: + if old is None: + os.environ.pop("CULA_INTRACARD_CP", None) + else: + os.environ["CULA_INTRACARD_CP"] = old + + +def run_intracard_direct(k, w, u, gk, h0, cu, *, output_final_state=True, save_new_value=True): + """Direct CP call — skips the auto-dispatch heuristic.""" + return intracard_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=h0, + output_final_state=output_final_state, + chunk_size=BT, + save_new_value=save_new_value, + cu_seqlens=cu, + cu_seqlens_cpu=cu.cpu(), + ) + + +def run_fla(k, w, u, gk, h0, cu, **kw): + return fla_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=h0, + chunk_size=BT, + cu_seqlens=cu, + **kw, + ) + + +def pytorch_ref(k, w, u, *, gk=None, initial_state=None, cu_seqlens, save_new_value=True): + """Pure-PyTorch fp32 reference for varlen chunk_gated_delta_rule_fwd_h. + + Mirrors the per-chunk math FLA's Triton kernel implements: + v_new = u - w @ h + h *= exp2(gk_last) # (if gk) + h += k^T @ v_new + """ + assert k.shape[0] == 1, "varlen reference expects packed B=1" + _, total, H, head_k = k.shape + head_v = u.shape[-1] + cu = cu_seqlens.cpu().tolist() + N = len(cu) - 1 + total_c = sum(math.ceil((cu[i + 1] - cu[i]) / BT) for i in range(N)) + + h_out = torch.empty(1, total_c, H, head_k, head_v, dtype=torch.bfloat16, device=k.device) + v_out = torch.empty_like(u) if save_new_value else None + ht_out = torch.empty(N, H, head_k, head_v, dtype=torch.float32, device=k.device) + + ci = 0 + for s in range(N): + bos, eos = cu[s], cu[s + 1] + h = ( + initial_state[s].float().clone() + if initial_state is not None + else torch.zeros(H, head_k, head_v, dtype=torch.float32, device=k.device) + ) + for cs in range(bos, eos, BT): + ce = min(cs + BT, eos) + h_out[0, ci] = h.to(torch.bfloat16) + w_c = w[0, cs:ce].permute(1, 0, 2).float() + k_c = k[0, cs:ce].permute(1, 0, 2).float() + u_c = u[0, cs:ce].permute(1, 0, 2).float() + v_new = u_c - torch.matmul(w_c, h) + if v_out is not None: + v_out[0, cs:ce] = v_new.permute(1, 0, 2).to(torch.bfloat16) + if gk is not None: + gk_last = gk[0, cs:ce].permute(1, 0, 2).float()[:, -1, :] + h = h * torch.exp2(gk_last).unsqueeze(-1) + h = h + torch.matmul(k_c.transpose(-2, -1), v_new) + ci += 1 + ht_out[s] = h + return h_out, v_out, ht_out + + +def _assert_same_kernel(name, actual, ref): + """torch.testing.assert_close — matches tests/test_chunk_delta_h.py.""" + if actual is None or ref is None: + assert actual is ref, f"{name}: one is None and other isn't" + return + torch.testing.assert_close( + actual.float(), + ref.float(), + atol=ATOL_SAME_KERNEL, + rtol=RTOL_SAME_KERNEL, + msg=lambda m: f"{name}: {m}", + ) + + +def assert_cp_splits(cu, H, total_T): + """Fail fast if the strategy doesn't even try to engage CP for this config. + + Note: we do NOT assert the post-split SM guard (total_subseqs * 2 * H <= num_sms). + intracard_fwd_h falls back gracefully to the non-CP path when that guard rejects, + so the test still exercises a valid code path even if CP scheduling itself doesn't + engage. + """ + cu_cpu = cu.cpu() + num_sms = get_device_sm_count(torch.device(DEVICE)) + assert should_use_intracard_cp(cu_cpu, num_sms, H, BT), ( + "should_use_intracard_cp returned False — config does not trigger CP" + ) + max_seq = int(torch.diff(cu_cpu).max().item()) + subseq_len = compute_subseq_len(max_seq, num_sms, H, BT, num_seqs=len(cu_cpu) - 1) + _, split_info, _ = prepare_subseq_cu_seqlens(cu_cpu, subseq_len, BT) + assert split_info, "config must exercise the split path" + + +# ====================== Dispatch path: CP vs no-CP ====================== +# Verifies chunk_gated_delta_rule_fwd_h routes to CP under env+inference_mode, +# and matches the same-kernel no-CP baseline. + +DISPATCH_CONFIGS = [ + ([32768], 4, False), + ([32768], 4, True), + ([65536], 4, True), + ([32768], 8, True), + ([32768, 256, 32768], 4, True), + ([65536, 128], 4, False), + ([32768, 32768, 32768], 4, True), + ([65536, 256, 128, 64], 8, True), +] + + +@pytest.mark.parametrize("seq_lens,H,use_gk", DISPATCH_CONFIGS) +def test_cp_autodispatch_matches_baseline(seq_lens, H, use_gk): + """CP auto-dispatch output equals no-CP baseline (same kernel). + + Tolerance: `torch.testing.assert_close(atol=1e-2, rtol=1e-2)` — + matches the CP block in tests/test_chunk_delta_h.py. + """ + k, w, u, gk, _, cu = make_varlen_inputs(seq_lens, H, use_gk=use_gk) + h_base, v_base, _ = run_cula_no_cp(k, w, u, gk, None, cu) + h_cp, v_cp, _ = run_cula_cp(k, w, u, gk, None, cu) + _assert_same_kernel("h", h_cp, h_base) + _assert_same_kernel("v_new", v_cp, v_base) + + +@pytest.mark.parametrize("seq_lens,H", [([32768], 4), ([32768, 256, 32768], 4)]) +def test_cp_autodispatch_with_h0(seq_lens, H): + """CP path preserves h0 input and ht output.""" + k, w, u, gk, h0, cu = make_varlen_inputs(seq_lens, H, use_gk=True, use_h0=True) + h_base, v_base, ht_base = run_cula_no_cp( + k, + w, + u, + gk, + h0, + cu, + output_final_state=True, + ) + h_cp, v_cp, ht_cp = run_cula_cp( + k, + w, + u, + gk, + h0, + cu, + output_final_state=True, + ) + _assert_same_kernel("h", h_cp, h_base) + _assert_same_kernel("v_new", v_cp, v_base) + _assert_same_kernel("ht", ht_cp, ht_base) + + +@pytest.mark.parametrize("T,H", [(32768, 4), (65536, 4), (32768, 8)]) +def test_cp_autodispatch_vs_fla(T, H): + """CP output matches FLA Triton reference (cross-impl). + + Tolerance: FLA's `assert_close` ratio=0.005 (RMSE/RMS <= 0.5%) — + same as FLA tests/ops/test_gated_delta.py for fwd outputs. + """ + k, w, u, gk, _, cu = make_varlen_inputs([T], H, use_gk=True) + h_fla, _, _ = run_fla(k, w, u, gk, None, cu) + h_cp, _, _ = run_cula_cp(k, w, u, gk, None, cu) + assert_close(f"h (T={T},H={H})", h_fla, h_cp, ratio=RATIO_VS_FLA) + + +# ====================== Accuracy: vs PyTorch fp32 reference ====================== +# Direct entry intracard_fwd_h, ground truth = pure-PyTorch fp32. + +ACCURACY_CONFIGS = [ + ([65536], 4, False, False), + ([65536], 4, True, True), + ([65536, 512], 4, True, True), + ([65536, 256, 32768], 4, True, False), + ([65536, 128], 4, False, True), + ([131072], 4, True, True), + ([65536, 512, 256, 128], 4, True, False), + ([65536, 1024, 8192], 8, True, True), +] + + +@pytest.mark.parametrize("seq_lens,H,use_gk,use_h0", ACCURACY_CONFIGS) +def test_intracard_cp_vs_pytorch_ref(seq_lens, H, use_gk, use_h0): + """CP output (h, v_new, ht) matches PyTorch fp32 reference. + + Tolerance: FLA's `assert_close` ratio=0.005 (RMSE/RMS <= 0.5%). + """ + k, w, u, gk, h0, cu = make_varlen_inputs( + seq_lens, + H, + use_gk=use_gk, + use_h0=use_h0, + seed=20260428, + ) + assert_cp_splits(cu, H, k.shape[1]) + with torch.inference_mode(): + ref_h, ref_v, ref_ht = pytorch_ref( + k, + w, + u, + gk=gk, + initial_state=h0, + cu_seqlens=cu, + ) + cp_h, cp_v, cp_ht = run_intracard_direct(k, w, u, gk, h0, cu) + torch.cuda.synchronize() + assert_close("h", ref_h, cp_h, ratio=RATIO_VS_REF) + assert_close("v_new", ref_v, cp_v, ratio=RATIO_VS_REF) + assert_close("ht", ref_ht, cp_ht, ratio=RATIO_VS_REF) + + +# ====================== Final state ht correctness ====================== +# Per-sequence ht must be independently correct for prefill→decode handoff. + +FINAL_STATE_CONFIGS = [ + ([65536], 4, False, False), + ([65536], 4, True, True), + ([65536], 8, True, True), + ([65536, 16384], 4, True, True), + ([65536, 512, 16384], 4, True, False), +] + + +@pytest.mark.parametrize("seq_lens,H,use_gk,use_h0", FINAL_STATE_CONFIGS) +def test_intracard_cp_final_state_per_seq(seq_lens, H, use_gk, use_h0): + """Each sequence's ht matches PyTorch ref independently (no cross-leakage).""" + k, w, u, gk, h0, cu = make_varlen_inputs( + seq_lens, + H, + use_gk=use_gk, + use_h0=use_h0, + seed=20260430, + ) + assert_cp_splits(cu, H, k.shape[1]) + with torch.inference_mode(): + _, _, ref_ht = pytorch_ref( + k, + w, + u, + gk=gk, + initial_state=h0, + cu_seqlens=cu, + save_new_value=False, + ) + _, _, cp_ht = run_intracard_direct( + k, + w, + u, + gk, + h0, + cu, + save_new_value=False, + ) + torch.cuda.synchronize() + assert cp_ht is not None and cp_ht.shape == ref_ht.shape + for i in range(len(seq_lens)): + assert_close(f"ht[{i}] (len={seq_lens[i]})", ref_ht[i], cp_ht[i], ratio=RATIO_VS_REF) + + +# ====================== Stress: race / non-determinism ====================== +# CP uses dynamic atomicAdd scheduling + multi-sub-seq merge — re-running the +# same inputs must produce the same outputs (no race, no order-dependence). + +STRESS_ITERS = 100 + + +@pytest.mark.parametrize( + "seq_lens,H,use_gk,use_h0", + [ + pytest.param([65536], 4, True, True, id="single-64K-H4-gk-h0"), + pytest.param([65536, 4096], 4, True, True, id="multi-64K+4K-H4-gk-h0"), + ], +) +def test_intracard_cp_stress_repeat(seq_lens, H, use_gk, use_h0): + """Run CP N times; every iter must match the first (race detection). + + Tolerance: ratio=1e-6 — deterministic CP should not drift across runs. + Uses `assert_close`'s atol short-circuit (abs <= 1e-6 → auto-pass). + """ + k, w, u, gk, h0, cu = make_varlen_inputs( + seq_lens, + H, + use_gk=use_gk, + use_h0=use_h0, + seed=20260516, + ) + assert_cp_splits(cu, H, k.shape[1]) + with torch.inference_mode(): + ref_h, ref_v, ref_ht = run_intracard_direct(k, w, u, gk, h0, cu) + torch.cuda.synchronize() + for i in range(STRESS_ITERS): + cp_h, cp_v, cp_ht = run_intracard_direct(k, w, u, gk, h0, cu) + torch.cuda.synchronize() + assert_close(f"iter {i} h", ref_h, cp_h, ratio=RATIO_STRESS) + assert_close(f"iter {i} v", ref_v, cp_v, ratio=RATIO_STRESS) + assert_close(f"iter {i} ht", ref_ht, cp_ht, ratio=RATIO_STRESS) + + +def test_intracard_cp_h0_none_equiv_h0_zeros(): + """h0=None must produce identical ht to h0=zeros (no implicit init).""" + seq_lens, H = [65536, 4096], 4 + k, w, u, gk, _, cu = make_varlen_inputs(seq_lens, H, use_gk=True, seed=20260501) + assert_cp_splits(cu, H, k.shape[1]) + h0_zeros = torch.zeros(len(seq_lens), H, K, V, dtype=torch.float32, device=DEVICE) + with torch.inference_mode(): + _, _, ht_none = run_intracard_direct( + k, + w, + u, + gk, + None, + cu, + save_new_value=False, + ) + _, _, ht_zeros = run_intracard_direct( + k, + w, + u, + gk, + h0_zeros, + cu, + save_new_value=False, + ) + torch.cuda.synchronize() + diff = (ht_none.float() - ht_zeros.float()).abs().max().item() + assert diff < 1e-4, f"h0=None vs h0=zeros diff {diff:.4e}" From 3006957c30a7f18735c81083b1afb03fa98241a4 Mon Sep 17 00:00:00 2001 From: Chaofan Yu <103550325+icavan@users.noreply.github.com> Date: Wed, 10 Jun 2026 11:09:00 +0800 Subject: [PATCH 23/34] Relax dependency minimum versions (#87) --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b70e04fe..ef1a531b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,8 +11,8 @@ readme = "README.md" authors = [ { name = "cula contributors" } ] requires-python = ">=3.10" dependencies = [ - "nvidia-cutlass-dsl==4.4.2", - "apache-tvm-ffi==0.1.9", + "nvidia-cutlass-dsl>=4.4.2", + "apache-tvm-ffi>=0.1.9", ] license = { text = "Apache-2.0" } From 7b1e12789768e312b566284f8c21b4cb918db249 Mon Sep 17 00:00:00 2001 From: tongke <124763920+tongke6@users.noreply.github.com> Date: Tue, 23 Jun 2026 13:49:42 +0800 Subject: [PATCH 24/34] [CI] Add GitHub workflow for building and releasing fat wheels (#91) * Split CUDA extensions by SM architecture for fat-binary wheel builds (#83) Replace the monolithic `cula.cudac` extension with per-arch extensions (`cula._cudac_sm90`, `cula._cudac_sm100`) so that SM90 and SM100/SM103 kernels are compiled independently with their own `-gencode` flags. This enables building fat-binary wheels containing all architectures without needing the target GPU present at build time. Key changes: - Split pybind.cu into per-file PYBIND11_MODULE definitions - Add `cula/cudac.py` proxy module for backwards-compatible imports - Add `CULA_BUILD_ALL_ARCHS=1` env var to enable all SM targets - Add `--fat` flag to build_wheel.sh for CI fat-binary builds - Pin dependency versions and use `no-local-version` scheme for reproducible wheel filenames - Use setuptools_scm for dynamic `__version__` - Document pre-built wheel installation in README * fix ruff lint errors * revert version requirements changes * Make cudac proxy thread-safe and raise on missing extensions Add double-checked locking to _CudacProxy._load() to prevent race conditions in multi-threaded environments. Raise a descriptive ImportError when no CUDA extensions can be loaded instead of silently producing AttributeError later. * Surface per-extension import errors in cudac proxy The blanket `except ImportError: pass` swallowed the actual failure reason, making it impossible to diagnose missing shared libraries or build issues. Collect each extension's ImportError and include them in the raised message. * Fix build-release matrix with DRY expression mapping * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Add README example for building fat wheels * Surface partial-extension failures in cudac proxy Catch (ImportError, AttributeError, OSError) when scanning per-arch extensions: pybind11 modules commonly surface missing-symbol / ABI / libcudart failures as AttributeError or OSError rather than ImportError, so the prior narrow catch silently dropped one extension's failure when another succeeded, leaving its kernels missing without diagnostic. Emit a UserWarning naming each failing extension on partial failure (all-fail still raises ImportError), preserving the c955d47 intent of surfacing per-extension errors. Also document the load-once-per-process semantics in the module docstring. * Build release wheels against manylinux_2_28 * fix python 3.12 GLIBC compat problems on ubi8 * install gcc13 * Load CUDA extension matching current GPU architecture Select the per-architecture CUDA extension from the active device compute capability instead of scanning every built extension. SM100/SM103 now load the SM100 extension, while SM90 loads the SM90 extension. This avoids exposing kernels from mismatched GPU architectures and reports clearer errors when the matching extension is missing or unsupported. --------- Co-authored-by: yz262713 Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .github/workflows/build-release.yml | 141 +++++++++++++++++++++++++ README.md | 18 +++- csrc/api/kda_sm100.cu | 6 ++ csrc/api/kda_sm90.cu | 5 + csrc/api/pybind.cu | 80 --------------- cula/__init__.py | 5 +- cula/cudac.py | 102 +++++++++++++++++++ pyproject.toml | 5 +- scripts/build_wheel.sh | 18 +++- setup.py | 153 +++++++++++++++------------- tests/conftest.py | 7 +- 11 files changed, 373 insertions(+), 167 deletions(-) create mode 100644 .github/workflows/build-release.yml delete mode 100644 csrc/api/pybind.cu create mode 100644 cula/cudac.py diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml new file mode 100644 index 00000000..72e593c5 --- /dev/null +++ b/.github/workflows/build-release.yml @@ -0,0 +1,141 @@ +name: Build & Release Wheels + +on: + push: + tags: + - "v*" + workflow_dispatch: + +concurrency: + group: build-release-${{ github.ref }} + cancel-in-progress: true + +jobs: + build-wheel: + name: "wheel / ${{ matrix.cuda }} / cp312 / ${{ matrix.arch }}" + runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }} + defaults: + run: + shell: bash + strategy: + fail-fast: false + matrix: + cuda: + - cu129 + - cu130 + arch: + - x86_64 + - aarch64 + container: + # UBI 8 provides the glibc 2.28 baseline required by manylinux_2_28. + image: "nvidia/cuda:${{ matrix.cuda == 'cu129' && '12.9.0' || '13.0.0' }}-devel-ubi8" + + steps: + - name: Free disk space + run: | + rm -rf /opt/hostedtoolcache /usr/local/lib/android /usr/share/dotnet \ + /usr/local/share/boost /opt/ghc 2>/dev/null || true + dnf clean all 2>/dev/null || true + df -h / || true + + - name: Install system dependencies + run: | + dnf install -y \ + git \ + gcc-toolset-13-gcc \ + gcc-toolset-13-gcc-c++ \ + python3.12 \ + python3.12-devel \ + python3.12-pip + dnf clean all + + - name: Checkout + uses: actions/checkout@v5 + with: + fetch-depth: 0 + submodules: recursive + + - name: Configure git safe directory + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" + + - name: Install Python dependencies + run: | + python3.12 -m pip install --no-cache-dir --upgrade pip + python3.12 -m pip install --no-cache-dir torch --index-url ${{ matrix.cuda == 'cu129' && 'https://download.pytorch.org/whl/cu129' || 'https://download.pytorch.org/whl/cu130' }} + python3.12 -m pip install --no-cache-dir setuptools wheel "setuptools_scm>=6.0" build ninja auditwheel patchelf + + - name: Compute version + id: version + run: | + if [[ "$GITHUB_REF" == refs/tags/v* ]]; then + BASE="${GITHUB_REF#refs/tags/v}" + else + # Strip any local segment (+gXXX) so we get a clean base + BASE=$(python3.12 -c "from setuptools_scm import get_version; print(get_version().split('+')[0])") + fi + echo "version=${BASE}+${{ matrix.cuda }}" >> "$GITHUB_OUTPUT" + + - name: Build fat-binary wheel + env: + CC: /opt/rh/gcc-toolset-13/root/usr/bin/gcc + CXX: /opt/rh/gcc-toolset-13/root/usr/bin/g++ + CUDAHOSTCXX: /opt/rh/gcc-toolset-13/root/usr/bin/g++ + CULA_BUILD_ALL_ARCHS: "1" + SETUPTOOLS_SCM_PRETEND_VERSION: "${{ steps.version.outputs.version }}" + NVCC_THREADS: "4" + MAX_JOBS: "4" + run: | + "$CC" --version + "$CXX" --version + python3.12 -m build --wheel --no-isolation --outdir dist-raw + + - name: Repair wheel for manylinux_2_28 + run: | + # These libraries are supplied by the NVIDIA driver, PyTorch, or + # PyTorch's CUDA runtime dependency and must remain external. + python3.12 -m auditwheel repair \ + --plat manylinux_2_28_${{ matrix.arch }} \ + --exclude libcuda.so.1 \ + --exclude 'libcudart.so.*' \ + --exclude 'libc10*.so' \ + --exclude 'libtorch*.so' \ + --wheel-dir dist \ + dist-raw/*.whl + + - name: Verify wheel + run: | + echo "Built wheel:" + ls -lh dist/*.whl + ls dist/*.whl | grep -q "+${{ matrix.cuda }}" \ + || { echo "ERROR: wheel name missing +${{ matrix.cuda }} suffix"; exit 1; } + ls dist/*.whl | grep -q "manylinux_2_28_${{ matrix.arch }}" \ + || { echo "ERROR: wheel is not tagged manylinux_2_28_${{ matrix.arch }}"; exit 1; } + python3.12 -m auditwheel show dist/*.whl + + - name: Upload wheel artifact + uses: actions/upload-artifact@v6 + with: + name: wheel-${{ matrix.cuda }}-${{ matrix.arch }} + path: dist/*.whl + + release: + name: Create GitHub Release + needs: [build-wheel] + runs-on: ubuntu-latest + if: startsWith(github.ref, 'refs/tags/v') + permissions: + contents: write + steps: + - name: Download all artifacts + uses: actions/download-artifact@v6 + with: + path: artifacts/ + + - name: Create release + uses: softprops/action-gh-release@v3 + with: + files: | + artifacts/wheel-*/*.whl + generate_release_notes: true + draft: true + prerelease: ${{ contains(github.ref, 'rc') || contains(github.ref, 'beta') || contains(github.ref, 'alpha') }} diff --git a/README.md b/README.md index 7bed61e2..09418811 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,16 @@ cuLA supports both **Hopper (SM90)** and **Blackwell (SM10X)** GPUs. > **Note:** The PyTorch CUDA version must match your system CUDA Toolkit version. Check with `nvcc --version` and `python -c "import torch; print(torch.version.cuda)"`. +### Pre-built Wheels + +Pre-built fat-binary wheels (SM90 + SM100 + SM103) are available on [GitHub Releases](https://github.com/inclusionAI/cuLA/releases). Linux wheels target `manylinux_2_28` and require glibc 2.28 or newer: + + pip install "cuda-linear-attention==+" -f https://github.com/inclusionAI/cuLA/releases/expanded_assets/ + +Replace `` with the release tag (e.g., `v0.2.0`), `` with the base version (e.g., `0.2.0`), and `` with your PyTorch CUDA build tag (e.g., `cu129` or `cu130`). Or download the `.whl` file directly from the [Releases page](https://github.com/inclusionAI/cuLA/releases) and install it with `pip install .whl`. + +### Build from Source + **Clone cuLA & dependencies:** ```bash @@ -47,6 +57,12 @@ pip install -e third_party/flash-linear-attention pip install -e . --no-build-isolation ``` +**Build fat wheel (SM90 + SM100 + SM103):** + +```bash +CULA_BUILD_ALL_ARCHS=1 python -m build --wheel --no-isolation +``` + ## Quick Start ### KDA (Kimi Delta Attention) — Blackwell (SM10X) @@ -239,4 +255,4 @@ No CUDA experience is required as long as you're a quick learner. For Q&A and discussion, you can join us through: - **Slack:** [cuLA Slack Community](https://join.slack.com/t/cula-hq/shared_invite/zt-3uaacvm9y-xJwZyGueeKxZRYQlj7~hxw) -- **WeChat:** The WeChat group has exceeded 200 members and can no longer be joined via QR code. To join, please send your WeChat ID to any of the following emails and we'll invite you: **chaofanyu@gmail.com** / **kevinzz08@foxmail.com** / **yzpag@gmail.com** / **haoc80996@gmail.com**. You can also ask someone already in the group to invite you directly. \ No newline at end of file +- **WeChat:** The WeChat group has exceeded 200 members and can no longer be joined via QR code. To join, please send your WeChat ID to any of the following emails and we'll invite you: **chaofanyu@gmail.com** / **kevinzz08@foxmail.com** / **yzpag@gmail.com** / **haoc80996@gmail.com**. You can also ask someone already in the group to invite you directly. diff --git a/csrc/api/kda_sm100.cu b/csrc/api/kda_sm100.cu index 7edca370..020d90ca 100644 --- a/csrc/api/kda_sm100.cu +++ b/csrc/api/kda_sm100.cu @@ -188,4 +188,10 @@ ChunkKDAFwdRecompWU( StaticPersistentTileScheduler::Params{tile_num, params.h_v, params.heads_per_group, params.num_sm, nullptr}; kda::sm100::run_kda_fwd_recomp_w_u_sm100(params, at::cuda::getCurrentCUDAStream()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "cuLA SM100/SM103 kernels"; + m.def("chunk_kda_fwd_intra_cuda", &ChunkKDAFwdIntra); + m.def("recompute_w_u_cuda", &ChunkKDAFwdRecompWU); } \ No newline at end of file diff --git a/csrc/api/kda_sm90.cu b/csrc/api/kda_sm90.cu index 9e016eb1..d80df7cc 100644 --- a/csrc/api/kda_sm90.cu +++ b/csrc/api/kda_sm90.cu @@ -191,3 +191,8 @@ kda_fwd_prefill( return {output, output_state}; } + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "cuLA SM90 kernels"; + m.def("kda_fwd_prefill", &kda_fwd_prefill); +} diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu deleted file mode 100644 index d14a41c5..00000000 --- a/csrc/api/pybind.cu +++ /dev/null @@ -1,80 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include -#include -#include -#include - -#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) -void -ChunkKDAFwdIntra( - at::Tensor q, - at::Tensor k, - at::Tensor g, - at::Tensor beta, - at::Tensor cu_seqlens, - at::Tensor chunk_indices, - at::Tensor Aqk_out, - at::Tensor Akk_out, - at::Tensor tile_counter, - float scale, - int chunk_size, - bool use_tf32_inverse, - bool unified_gref); -void -ChunkKDAFwdRecompWU( - at::Tensor k, - at::Tensor v, - at::Tensor beta, - at::Tensor A, - at::Tensor g, - at::Tensor cu_seqlens, - at::Tensor chunk_indices, - at::Tensor w_out, - at::Tensor u_out, - at::Tensor kg_out, - int chunk_size, - std::optional q, - std::optional qg_out); -#endif - -#if defined(CULA_SM90A_ENABLED) -std::tuple> -kda_fwd_prefill( - std::optional output_, - std::optional output_state_, - torch::Tensor const& q, - torch::Tensor const& k, - torch::Tensor const& v, - std::optional input_state_, - std::optional alpha_, - std::optional beta_, - torch::Tensor const& cu_seqlens, - torch::Tensor workspace_buffer, - float scale, - bool output_final_state, - bool safe_gate); -#endif - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "cuLA"; -#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) - m.def("chunk_kda_fwd_intra_cuda", &ChunkKDAFwdIntra); - m.def("recompute_w_u_cuda", &ChunkKDAFwdRecompWU); -#endif -#if defined(CULA_SM90A_ENABLED) - m.def("kda_fwd_prefill", &kda_fwd_prefill); -#endif -} diff --git a/cula/__init__.py b/cula/__init__.py index 7272e289..6e13aa13 100644 --- a/cula/__init__.py +++ b/cula/__init__.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "0.1.0" +try: + from cula._version import version as __version__ +except ImportError: + __version__ = "0.1.0" from cula.ops.lightning_attn_sm100 import LinearAttentionChunkwiseDecay diff --git a/cula/cudac.py b/cula/cudac.py new file mode 100644 index 00000000..28fb5f38 --- /dev/null +++ b/cula/cudac.py @@ -0,0 +1,102 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unified interface to per-architecture CUDA extensions. + +Downstream code can continue to use ``import cula.cudac as cula_cuda`` +and call ``cula_cuda.kda_fwd_prefill(...)`` or +``cula_cuda.chunk_kda_fwd_intra_cuda(...)`` without knowing which +extension provides the function. + +Loading is **once per process**: the first attribute access checks the +currently active CUDA device, imports the matching ``cula._cudac_sm*`` +extension, and caches the discovered callables on the module instance. +Changing the active CUDA device to a different architecture after a +process has already loaded ``cula.cudac`` will therefore not be picked +up -- callers that need a different extension must restart Python. +""" + +import importlib +import sys +import threading +from types import ModuleType + + +def _current_device_extension() -> tuple[str, str]: + try: + import torch + except ImportError as exc: + raise ImportError("cuLA CUDA extensions require PyTorch to detect the current GPU.") from exc + + if not torch.cuda.is_available(): + raise RuntimeError("cuLA CUDA extensions require a visible CUDA GPU, but torch.cuda.is_available() is False.") + + device = torch.cuda.current_device() + prop = torch.cuda.get_device_properties(device) + sm_label = f"sm_{prop.major}{prop.minor}" + if prop.major == 10 and prop.minor in (0, 3): + return "cula._cudac_sm100", sm_label + if prop.major == 9 and prop.minor == 0: + return "cula._cudac_sm90", sm_label + raise RuntimeError(f"Unsupported CUDA compute capability {sm_label}. Supported architectures: sm_100, sm_103, sm_90.") + + +class _CudacProxy(ModuleType): + """Lazy proxy that exposes functions from the current GPU arch extension.""" + + def __init__(self): + super().__init__(__name__) + self.__path__ = [] + self._modules_loaded = False + self._funcs: dict[str, object] = {} + self._lock = threading.Lock() + + def _load(self): + if self._modules_loaded: + return + with self._lock: + if self._modules_loaded: + return + ext_name, sm_label = _current_device_extension() + try: + mod = importlib.import_module(ext_name) + for attr in dir(mod): + if not attr.startswith("_"): + self._funcs[attr] = getattr(mod, attr) + except (ImportError, AttributeError, OSError) as exc: + raise ImportError( + f"The cuLA CUDA extension for the current GPU ({sm_label}) could not be imported. " + f"Extension {ext_name} failed with: {exc}. " + "Please make sure cuLA is compiled correctly." + ) from exc + self.__dict__.update(self._funcs) + self._modules_loaded = True + + def __getattr__(self, name: str): + if name.startswith("_"): + raise AttributeError(name) + self._load() + try: + return self._funcs[name] + except KeyError: + raise AttributeError(f"module 'cula.cudac' has no attribute '{name}'") from None + + def __dir__(self): + self._load() + return list(self._funcs.keys()) + + +_proxy = _CudacProxy() +_proxy.__dict__.update({k: globals().get(k) for k in ("__spec__", "__file__", "__package__", "__loader__")}) +sys.modules[__name__] = _proxy diff --git a/pyproject.toml b/pyproject.toml index ef1a531b..fe93e562 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,9 +84,6 @@ force-sort-within-sections = false "cula/kda/blackwell_fused_fwd.py" = ["F821"] [tool.setuptools_scm] -# write generated version into package for runtime access write_to = "cula/_version.py" -# add a date-based local suffix when needed -local_scheme = "node-and-date" -# fallback for non-git sources +local_scheme = "no-local-version" fallback_version = "0.1.0" diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 42b35665..79ac3305 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -18,10 +18,19 @@ cd "$REPO_ROOT" # Parse args ISOLATION_FLAG="--no-isolation" -if [[ "${1:-}" == "--isolated" ]]; then - ISOLATION_FLAG="" - echo "[build_wheel] Using isolated build environment" -else +for arg in "$@"; do + case "$arg" in + --isolated) + ISOLATION_FLAG="" + echo "[build_wheel] Using isolated build environment" + ;; + --fat) + export CULA_BUILD_ALL_ARCHS=1 + echo "[build_wheel] Fat binary: building for all SM architectures" + ;; + esac +done +if [[ "$ISOLATION_FLAG" == "--no-isolation" ]]; then echo "[build_wheel] Using current environment (--no-isolation)" fi @@ -33,6 +42,7 @@ rm -rf dist build *.egg-info echo "[build_wheel] Python: $(python -V 2>&1)" echo "[build_wheel] torch: $(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo 'not installed')" echo "[build_wheel] CUDA: $(nvcc --version 2>/dev/null | grep 'release' | sed 's/.*release //' | sed 's/,.*//' || echo 'not found')" +echo "[build_wheel] Fat binary: ${CULA_BUILD_ALL_ARCHS:-0}" # Build wheel echo "[build_wheel] Building wheel..." diff --git a/setup.py b/setup.py index f7b11b95..78c61e5c 100644 --- a/setup.py +++ b/setup.py @@ -46,13 +46,15 @@ def detect_gpu_archs() -> tuple[bool, bool, bool]: def resolve_disable_flag(env_name: str, detected: bool) -> bool: """ Resolve whether to disable a given SM target. + - If CULA_BUILD_ALL_ARCHS is set, all targets are enabled unconditionally. - If the environment variable is explicitly set, honour it. - Otherwise, disable the target when no matching GPU is detected. """ + if os.getenv("CULA_BUILD_ALL_ARCHS", "0") == "1": + return False env_val = os.getenv(env_name) if env_val is not None: return env_val.lower() in ["true", "1", "y", "yes"] - # Auto-detect: disable if no matching device found disable = not detected if disable: print(f" No matching GPU detected; auto-setting {env_name}=1 (disable). Set {env_name}=0 to override.") @@ -66,7 +68,11 @@ def get_features_args(): USE_FAST_MATH = os.getenv("CULA_USE_FAST_MATH", "1") == "1" -print("Detecting GPU architectures...") +if os.getenv("CULA_BUILD_ALL_ARCHS", "0") == "1": + print("CULA_BUILD_ALL_ARCHS=1: enabling all SM targets (sm90a, sm100a, sm103a)") +else: + print("Detecting GPU architectures...") + _has_sm100, _has_sm103, _has_sm90 = detect_gpu_archs() DISABLE_SM100 = resolve_disable_flag("CULA_DISABLE_SM100", _has_sm100) DISABLE_SM103 = resolve_disable_flag("CULA_DISABLE_SM103", _has_sm103) @@ -111,26 +117,6 @@ def assert_blackwell_build_env() -> None: ) -def get_arch_flags(): - major, minor = get_nvcc_version() - print(f"Compiling using NVCC {major}.{minor}") - - # Validate Blackwell build environment - assert_blackwell_build_env() - - arch_flags = [] - if not DISABLE_SM100: - arch_flags.extend(["-gencode", "arch=compute_100a,code=sm_100a"]) - arch_flags.extend(["-DCULA_SM100_ENABLED"]) - if not DISABLE_SM103: - arch_flags.extend(["-gencode", "arch=compute_103a,code=sm_103a"]) - arch_flags.extend(["-DCULA_SM103_ENABLED"]) - if not DISABLE_SM90: - arch_flags.extend(["-gencode", "arch=compute_90a,code=sm_90a"]) - arch_flags.extend(["-DCULA_SM90A_ENABLED"]) - return arch_flags - - def get_nvcc_thread_args(): nvcc_threads = os.getenv("NVCC_THREADS") or "32" return ["--threads", nvcc_threads] @@ -145,61 +131,84 @@ def get_nvcc_thread_args(): else: cxx_args = ["-O3", "-std=c++20", "-DNDEBUG", "-Wno-deprecated-declarations"] -cuda_sources = [ - "csrc/api/pybind.cu", +nvcc_common_args = [ + "-O3", + "-std=c++20", + "-DNDEBUG", + # "-D_USE_MATH_DEFINES", + "-Wno-deprecated-declarations", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "-lineinfo", + "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage", + "-diag-suppress=3189", ] + +include_dirs = [ + Path(this_dir) / "csrc", + Path(this_dir) / "csrc" / "kerutils" / "include", + Path(this_dir) / "csrc" / "cutlass" / "include", + Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", +] + +major, minor = get_nvcc_version() +print(f"Compiling using NVCC {major}.{minor}") +assert_blackwell_build_env() + +ext_modules = [] + if not DISABLE_SM100 or not DISABLE_SM103: - cuda_sources.extend( - [ - "csrc/api/kda_sm100.cu", - "csrc/kda/sm100/kda_fwd_sm100.cu", - ] - ) -if not DISABLE_SM90: - cuda_sources.extend( - [ - "csrc/api/kda_sm90.cu", - "csrc/kda/sm90/kda_fwd_sm90.cu", - "csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu", - ] + sm100_arch_flags = [] + if not DISABLE_SM100: + sm100_arch_flags.extend(["-gencode", "arch=compute_100a,code=sm_100a"]) + if not DISABLE_SM103: + sm100_arch_flags.extend(["-gencode", "arch=compute_103a,code=sm_103a"]) + + ext_modules.append( + CUDAExtension( + name="cula._cudac_sm100", + sources=[ + "csrc/api/kda_sm100.cu", + "csrc/kda/sm100/kda_fwd_sm100.cu", + ], + extra_compile_args={ + "cxx": cxx_args + get_features_args(), + "nvcc": nvcc_common_args + + get_features_args() + + sm100_arch_flags + + get_nvcc_thread_args() + + (["--use_fast_math"] if USE_FAST_MATH else []), + }, + include_dirs=include_dirs, + ) ) -ext_modules = [] -ext_modules.append( - CUDAExtension( - name="cula.cudac", - sources=cuda_sources, - extra_compile_args={ - "cxx": cxx_args + get_features_args(), - "nvcc": [ - "-O3", - "-std=c++20", - "-DNDEBUG", - # "-D_USE_MATH_DEFINES", - "-Wno-deprecated-declarations", - "-U__CUDA_NO_HALF_OPERATORS__", - "-U__CUDA_NO_HALF_CONVERSIONS__", - "-U__CUDA_NO_HALF2_OPERATORS__", - "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", - "--expt-relaxed-constexpr", - "--expt-extended-lambda", - "-lineinfo", - "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage", - "-diag-suppress=3189", # suppress the warning of torch in C++ 20 - ] - + get_features_args() - + get_arch_flags() - + get_nvcc_thread_args() - + (["--use_fast_math"] if USE_FAST_MATH else []), - }, - include_dirs=[ - Path(this_dir) / "csrc", - Path(this_dir) / "csrc" / "kerutils" / "include", - Path(this_dir) / "csrc" / "cutlass" / "include", - Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", - ], +if not DISABLE_SM90: + sm90_arch_flags = ["-gencode", "arch=compute_90a,code=sm_90a", "-DCULA_SM90A_ENABLED"] + + ext_modules.append( + CUDAExtension( + name="cula._cudac_sm90", + sources=[ + "csrc/api/kda_sm90.cu", + "csrc/kda/sm90/kda_fwd_sm90.cu", + "csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu", + ], + extra_compile_args={ + "cxx": cxx_args + get_features_args(), + "nvcc": nvcc_common_args + + get_features_args() + + sm90_arch_flags + + get_nvcc_thread_args() + + (["--use_fast_math"] if USE_FAST_MATH else []), + }, + include_dirs=include_dirs, + ) ) -) setup( name="cuda-linear-attention", diff --git a/tests/conftest.py b/tests/conftest.py index f144c10b..a9338aca 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ import re + import pytest import torch @@ -56,9 +57,5 @@ def pytest_collection_modifyitems(config, items): item.add_marker(skip_slow) continue callspec = getattr(item, "callspec", None) - if ( - callspec is not None - and callspec.params.get("disable_recompute") - and "kda_fast_norecomp" not in item.keywords - ): + if callspec is not None and callspec.params.get("disable_recompute") and "kda_fast_norecomp" not in item.keywords: item.add_marker(skip_fast_norecomp) From ff1b2361cadedf652a79fc6600db02260445bb06 Mon Sep 17 00:00:00 2001 From: tongke <124763920+tongke6@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:08:49 +0800 Subject: [PATCH 25/34] ci: add prek lint workflow (#98) --- .github/workflows/lint.yml | 40 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 .github/workflows/lint.yml diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000..82daf202 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,40 @@ +name: Lint + +on: + pull_request: + branches: + - main + push: + branches: + - main + +concurrency: + group: lint-${{ github.ref }} + cancel-in-progress: true + +jobs: + clang-format: + name: clang-format + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: j178/prek-action@v2 + with: + install-only: true + - name: Check clang-format + run: prek run clang-format --all-files --show-diff-on-failure --color=always + + ruff: + name: Ruff + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: j178/prek-action@v2 + with: + install-only: true + + - name: Check Ruff lint + run: prek run ruff --all-files --show-diff-on-failure --color=always + + - name: Check Ruff format + run: prek run ruff-format --all-files --show-diff-on-failure --color=always From 6cacc37fd420e72c859c1dec6c870a13dc2a0e9a Mon Sep 17 00:00:00 2001 From: tongke <124763920+tongke6@users.noreply.github.com> Date: Thu, 25 Jun 2026 10:50:13 +0800 Subject: [PATCH 26/34] ci: align pytorch version requirements in wheel build (#99) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ci: pin torch=2.9.1 * ci: add LD_LIBRARY_PATH for pip-installed CUDA libraries The nvidia/cuda devel container images include the CUDA toolkit but not cuDNN. When torch is installed via pip, nvidia-cudnn-cu12 (and friends) land in site-packages/nvidia/*/lib/, which is not on the system library search path. This causes an ImportError for libcudnn.so.9 when the build backend tries to import torch. Add a step that discovers all PyTorch and NVIDIA pip package library directories and appends them to LD_LIBRARY_PATH via GITHUB_ENV, with a validation check that libcudnn.so.9 is actually found. * ci: simplify LD_LIBRARY_PATH step to 1-liner Verified working in CI — drop debug logging, redundant validation, and heredoc in favor of a compact inline python -c one-liner. * ci: fix shell quoting in LD_LIBRARY_PATH step The previous 1-liner broke because single quotes inside $(python3.12 -c '...') conflicted with the outer shell quoting. Switch to a heredoc (<<'EOF') for the Python snippet, which avoids nested quoting entirely and keeps the code readable. --- .github/workflows/build-release.yml | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml index 72e593c5..73ea4929 100644 --- a/.github/workflows/build-release.yml +++ b/.github/workflows/build-release.yml @@ -61,9 +61,26 @@ jobs: - name: Install Python dependencies run: | python3.12 -m pip install --no-cache-dir --upgrade pip - python3.12 -m pip install --no-cache-dir torch --index-url ${{ matrix.cuda == 'cu129' && 'https://download.pytorch.org/whl/cu129' || 'https://download.pytorch.org/whl/cu130' }} + python3.12 -m pip install --no-cache-dir torch==2.9.1 --index-url ${{ matrix.cuda == 'cu129' && 'https://download.pytorch.org/whl/cu129' || 'https://download.pytorch.org/whl/cu130' }} python3.12 -m pip install --no-cache-dir setuptools wheel "setuptools_scm>=6.0" build ninja auditwheel patchelf + - name: Expose PyTorch CUDA libraries + # The nvidia/cuda devel image ships the CUDA toolkit but not cuDNN. + # torch's pip wheels bundle cuDNN inside site-packages/nvidia/*/lib/, + # which is not on the linker search path — add it so torch._C can load. + run: | + LIB_DIRS="$(python3.12 <<'EOF' + import glob, os, site + dirs = [] + for sp in site.getsitepackages(): + for p in [os.path.join(sp, "torch", "lib")] + glob.glob(os.path.join(sp, "nvidia", "*", "lib")): + if os.path.isdir(p): + dirs.append(p) + print(":".join(dirs)) + EOF + )" + echo "LD_LIBRARY_PATH=${LIB_DIRS}${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" >> "$GITHUB_ENV" + - name: Compute version id: version run: | From f8d7db20d8ba2cca8aca04a28da05b6ad3182209 Mon Sep 17 00:00:00 2001 From: cher <117337477+cherhh@users.noreply.github.com> Date: Tue, 7 Jul 2026 23:13:35 +0800 Subject: [PATCH 27/34] refactor(kda): reorganize KDA backends into arch-first layout and add lazy imports (#100) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(kda): reorganize KDA backends into an arch-first cula/ops/kda/ layout - Move SM100 (Blackwell) modular-chunk backends, decode, and the unwired fully-fused WIP from flat cula/ops/*.py into cula/ops/kda/{sm100,decode,experimental}/. - Move the non-KDA lightning/linear prototypes under cula/ops/. - Add a central CP dispatch policy at cula/ops/kda/policy.py. - Make cula / cula.ops / cula.kda imports lazy (PEP 562) so `import cula` no longer eagerly pulls the CuTeDSL/CUDA-heavy modules. - Repoint all in-repo imports, benchmarks, tests, and docs. Pure reorganization, no kernel behavior change. The SM90 (Hopper) prefill stays the existing C++ kernel under csrc/kda/sm90. * refactor(kda): mark the fused Blackwell prefill as not-yet-available - Drop kda_prefill_blackwell from the cula.kda public exports; the fully-fused Blackwell prefill (cula/ops/kda/experimental/sm100_fused/) is unwired WIP. - get_kda_fused_fwd now raises NotImplementedError on SM100/SM103 instead of returning that experimental kernel. - Production Blackwell prefill stays the modular chunk_kda path. * refactor(kda): make SM100 intracard_fwd_h a pure split-or-raise executor - intracard_fwd_h now raises NotSplittableError when the shape cannot be meaningfully split, instead of silently falling back. - Drop the allow_fallback / skip_precheck flags, the two duplicated _no_cp fallback blocks, and the redundant pre-split heuristic recheck that the dispatch policy already performed. - chunk_gated_delta_rule_fwd_h now owns the fallback: re-raise for forced CP, fall through to the serial body for auto. - NotSplittableError subclasses ValueError for backward compatibility. Behavior-preserving: force -> raise and auto -> serial fallback are unchanged. * refactor(kda): remove dead LinearAttentionChunkwiseDecay re-export and stale REPO_LAYOUT sections * doc: clean up REPO_LAYOUT.md — remove stale sections and fix descriptions * refactor(kda): annotate use_cp compat shim and clean up kda package docstring * refactor(kda): trim verbose comments in policy.py, delta_h.py, and test_intracard_cp.py * style: apply ruff-format - cula/__init__.py: drop trailing blank line - cula/ops/kda/__init__.py: dedent module docstring, add final newline - cula/ops/kda/sm100/delta_h.py: drop extra blank lines * fix(kda): route get_kda_fused_fwd SM100 to the experimental fused wrapper The arch-first move left the SM100 branch as a TODO raise, so importing bench_kda_fused_fwd.py failed on Blackwell (module-level dispatch). Restore the pre-refactor routing to flash_kda_prefill at its new home. * chore: fix whitespace nits flagged by git diff --check --------- Co-authored-by: cheheng.ch --- REPO_LAYOUT.md | 146 ++- USAGE.md | 2 +- benchmarks/bench_chunk_delta_h.py | 2 +- benchmarks/bench_fwd_o.py | 2 +- benchmarks/bench_intracard_cp.py | 2 +- benchmarks/bench_kda_bwd_wy_dqkg_sm100.py | 2 +- benchmarks/bench_kda_decode.py | 2 +- benchmarks/bench_kda_fused_fwd.py | 2 +- benchmarks/bench_la_decode_vs_fla.py | 2 +- benchmarks/bench_lightning_attn.py | 2 +- benchmarks/bench_linear_attn.py | 2 +- cula/__init__.py | 6 - cula/kda/__init__.py | 29 +- cula/kda/chunk.py | 23 + cula/kda/chunk_bwd.py | 6 +- cula/kda/chunk_fwd.py | 13 +- cula/kda/chunk_intra.py | 1 + cula/lightning/__init__.py | 4 +- cula/ops/__init__.py | 25 +- cula/ops/cp/__init__.py | 3 - cula/ops/experimental/__init__.py | 4 + .../linear_attn_prototype.py} | 2 +- cula/ops/intrinsics_sm100.py | 418 -------- cula/ops/kda/__init__.py | 14 + cula/ops/kda/decode/__init__.py | 8 + .../ops/{kda_decode.py => kda/decode/cute.py} | 0 .../decode/reference_fla.py} | 0 cula/ops/kda/experimental/__init__.py | 4 + .../kda/experimental/sm100_fused/__init__.py | 4 + .../sm100_fused/kda_fully_fused_wip.py} | 0 .../kda/experimental/sm100_fused/wrapper.py} | 15 +- cula/ops/kda/policy.py | 98 ++ cula/ops/kda/sm100/__init__.py | 4 + .../sm100/bwd_wy_dqkg.py} | 12 +- cula/ops/kda/sm100/cp/__init__.py | 8 + cula/ops/{ => kda/sm100}/cp/chunk_delta_h.py | 49 +- cula/ops/{ => kda/sm100}/cp/merge.py | 100 +- cula/ops/{ => kda/sm100}/cp/pre_scan.py | 0 .../sm100/delta_h.py} | 51 +- .../{fwd_o_sm100.py => kda/sm100/fwd_o.py} | 0 cula/ops/lightning/__init__.py | 4 + .../ops/{la_decode.py => lightning/decode.py} | 0 .../prefill_sm100.py} | 0 cula/ops/ptx.py | 152 +++ cula/ops/ptx_umma_ext.py | 961 ------------------ cula/ops/sm100/__init__.py | 2 + cula/ops/sm100/ptx.py | 788 ++++++++++++++ cula/utils.py | 37 +- docs/chunk_delta_h_pipeline.md | 4 +- pyproject.toml | 4 +- tests/test_chunk_delta_h.py | 4 +- tests/test_compare_with_fla.py | 2 +- tests/test_fwd_o.py | 4 +- tests/test_intracard_cp.py | 64 +- tests/test_la_decode.py | 5 +- tests/test_la_decode_pool.py | 2 +- tests/test_lightning_attn.py | 2 +- tests/test_ptx_umma_masked.py | 9 +- tests/test_ptx_umma_ws.py | 14 +- 59 files changed, 1390 insertions(+), 1735 deletions(-) delete mode 100644 cula/ops/cp/__init__.py create mode 100644 cula/ops/experimental/__init__.py rename cula/ops/{linear_attn_sm100.py => experimental/linear_attn_prototype.py} (99%) delete mode 100644 cula/ops/intrinsics_sm100.py create mode 100644 cula/ops/kda/__init__.py create mode 100644 cula/ops/kda/decode/__init__.py rename cula/ops/{kda_decode.py => kda/decode/cute.py} (100%) rename cula/ops/{kda_decode_fla.py => kda/decode/reference_fla.py} (100%) create mode 100644 cula/ops/kda/experimental/__init__.py create mode 100644 cula/ops/kda/experimental/sm100_fused/__init__.py rename cula/ops/{kda_fully_fused_sm100_wip.py => kda/experimental/sm100_fused/kda_fully_fused_wip.py} (100%) rename cula/{kda/blackwell_fused_fwd.py => ops/kda/experimental/sm100_fused/wrapper.py} (97%) create mode 100644 cula/ops/kda/policy.py create mode 100644 cula/ops/kda/sm100/__init__.py rename cula/ops/{chunk_wy_dqkg_sm100.py => kda/sm100/bwd_wy_dqkg.py} (99%) create mode 100644 cula/ops/kda/sm100/cp/__init__.py rename cula/ops/{ => kda/sm100}/cp/chunk_delta_h.py (93%) rename cula/ops/{ => kda/sm100}/cp/merge.py (83%) rename cula/ops/{ => kda/sm100}/cp/pre_scan.py (100%) rename cula/ops/{chunk_delta_h_sm100.py => kda/sm100/delta_h.py} (99%) rename cula/ops/{fwd_o_sm100.py => kda/sm100/fwd_o.py} (100%) create mode 100644 cula/ops/lightning/__init__.py rename cula/ops/{la_decode.py => lightning/decode.py} (100%) rename cula/ops/{lightning_attn_sm100.py => lightning/prefill_sm100.py} (100%) create mode 100644 cula/ops/ptx.py delete mode 100644 cula/ops/ptx_umma_ext.py create mode 100644 cula/ops/sm100/__init__.py create mode 100644 cula/ops/sm100/ptx.py diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 6a9d5866..51f44511 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -1,101 +1,71 @@ # Repository Layout +Legend: `[exp]` experimental / unwired · `[non-KDA]` other operator. + ``` cuLA/ ├── cula/ # Python package (pip install -e .) -│ ├── kda/ # KDA (Kimi Delta Attention) operators -│ │ ├── chunk.py # End-to-end chunk KDA (fwd + bwd entry point) -│ │ ├── chunk_fwd.py # Chunk forward dispatch -│ │ ├── chunk_intra.py # Intra-chunk forward logic -│ │ ├── blackwell_fused_fwd.py # Fused KDA forward (SM100) -│ │ └── hopper_fused_fwd.py # Fused KDA forward (SM90) -│ ├── lightning/ # Lightning Attention operators -│ │ └── la_decode.py # Single-token decode kernel (CuTe DSL) -│ ├── ops/ # CuTe DSL kernel implementations -│ │ ├── chunk_delta_h_sm100.py # Chunk delta-H kernel (SM100) -│ │ ├── fwd_o_sm100.py # Forward output kernel (SM100) -│ │ ├── lightning_attn_sm100.py # Lightning Attention prefill kernel (SM100) -│ │ ├── linear_attn_sm100.py # Generic linear attention kernel (SM100) -│ │ ├── kda_fully_fused_sm100_wip.py # WIP fully fused KDA kernel (SM100) -│ └── utils.py # Shared utilities -│ -├── csrc/ # CUDA C++ / CUTLASS kernels -│ ├── api/ # PyBind11 bindings -│ │ ├── pybind.cu # Python ↔ CUDA binding entry -│ │ ├── kda_sm90.cu # SM90 API wrappers -│ │ └── kda_sm100.cu # SM100 API wrappers -│ ├── kda/ -│ │ ├── sm90/ # Hopper KDA kernels (CUTLASS 3.x) -│ │ │ ├── kda_fwd_sm90.cu -│ │ │ ├── kda_fwd_sm90_safe_gate.cu -│ │ │ ├── prefill_kernel.hpp -│ │ │ ├── collective/ # CUTLASS collective mainloop -│ │ │ ├── device/ # Device-level kernel wrappers -│ │ │ ├── kernel/ # Kernel-level logic -│ │ │ └── utils/ # SM90-specific helpers -│ │ └── sm100/ # Blackwell KDA kernels (CUTLASS 3.x) -│ │ ├── kda_fwd_sm100.cu -│ │ ├── kda_fwd_common.cuh -│ │ ├── kda_fwd_intra_kernel_sm100.hpp -│ │ ├── kda_fwd_intra_mainloop_sm100.hpp # Chunk intra mainloop -│ │ ├── kda_fwd_recomp_w_u_kernel_sm100.hpp -│ │ ├── kda_fwd_recomp_w_u_mainloop_sm100.hpp # Recompute W&U mainloop -│ │ ├── kda_config.hpp -│ │ ├── fwd_helpers.hpp -│ │ ├── sm100_umma_ext.hpp -│ │ └── tile_scheduler.hpp -│ └── kerutils/ -│ └── include/ # Shared C++ header utilities -│ -├── benchmarks/ # Performance benchmarks -│ ├── bench_kda.py # KDA fixed + varlen benchmark -│ ├── bench_lightning_attn.py # Lightning Attention prefill + varlen -│ ├── bench_la_decode_vs_fla.py # Decode: la_decode vs fla fused_recurrent -│ ├── bench_kda_fused_fwd.py # KDA fused forward benchmark -│ ├── bench_kda_chunk_intra.py # KDA chunk intra benchmark -│ ├── bench_chunk_delta_h.py # Chunk delta-H benchmark -│ ├── bench_fwd_o.py # Forward output benchmark -│ ├── bench_linear_attn.py # Linear attention benchmark -│ ├── generate_benchmark_md.py # Auto-generate BENCHMARK_GB200.md (Blackwell) -│ ├── generate_benchmark_hopper_md.py # Auto-generate BENCHMARK_H200.md (Hopper) -│ └── utils.py # Benchmark utilities -│ -├── tests/ # Unit / integration tests -│ ├── test_kda_compare_fla.py # Modular KDA forward vs FLA Triton -│ ├── test_kda.py # Modular KDA forward vs naive reference -│ ├── test_kda_fused_fwd.py # Fused KDA forward tests -│ ├── test_chunk_delta_h.py # Chunk delta-H tests -│ ├── test_fwd_o.py # Forward output tests -│ ├── test_compare_with_fla.py # General FLA comparison -│ ├── test_lightning_attn.py # Lightning Attention tests -│ └── test_la_decode.py # Decode kernel tests -│ -├── docs/ # Design documents -│ ├── chunk_delta_h_pipeline.md -│ ├── fwd_o_pipeline.md -│ └── lightning_attn_pipeline.md +│ ├── __init__.py +│ ├── utils.py # arch asserts, get_pre_scan, cu_seqlens helpers, ... +│ ├── cudac.py # re-export shim over the compiled C++ extension(s) +│ │ +│ ├── kda/ # KDA PUBLIC API + autograd + dispatch (NO kernels) +│ │ ├── __init__.py # lazy PUBLIC API: chunk_kda, kda_prefill_hopper, +│ │ │ # kda_decode, fused_sigmoid_gating_delta_rule_update +│ │ ├── chunk.py # chunk_kda + autograd — SM100 modular path (train + Blackwell prefill) +│ │ ├── chunk_fwd.py # chunk_kda_fwd — fwd orchestration (lazy-imports kernels) +│ │ ├── chunk_intra.py # fwd intra (C++ ext) + bwd intra (Triton) +│ │ ├── chunk_bwd.py # chunk_kda_bwd — Triton + FLA + CuTeDSL + C++ mix +│ │ └── hopper_fused_fwd.py # cula_kda_prefill (=kda_prefill_hopper) — SM90 prefill via the C++ kernel (cula.cudac) +│ │ +│ ├── lightning/ # [non-KDA] Lightning Attention operator (LinearAttentionChunkwiseDecay, lightning_attn_fwd, linear_attention_decode) +│ │ └── __init__.py +│ │ +│ └── ops/ # backend kernels (CuTe DSL / TVM-FFI) + shared helpers +│ ├── __init__.py # exports kda_decode, fused_sigmoid_..., linear_attention_decode +│ ├── inv.py / ptx.py # shared low-level helpers +│ ├── sm100/ # SM100 shared helper only +│ │ └── ptx.py # shared PTX helpers (used by KDA + lightning kernels) +│ │ +│ ├── kda/ # ★ KDA Python backends — by arch (sm100 today) +│ │ ├── policy.py # SM100 CP dispatch policy: use_intracard_cp:"auto"|bool +│ │ ├── sm100/ # SM100 (Blackwell) modular-chunk kernels +│ │ │ ├── delta_h.py # recurrence (chunk_gated_delta_rule_fwd_h) +│ │ │ ├── fwd_o.py # output (chunk_gla_fwd_o) +│ │ │ ├── bwd_wy_dqkg.py# backward wy/dqkg fused (used by chunk_bwd) +│ │ │ └── cp/ # SM100 intracard-CP: chunk_delta_h, pre_scan, merge +│ │ ├── decode/ # single-token decode +│ │ │ ├── cute.py # kda_decode / fused_sigmoid_gating_delta_rule_update (CuTe DSL) +│ │ │ └── reference_fla.py +│ │ └── experimental/sm100_fused/ # [exp] unwired fully-fused +│ │ ├── kda_fully_fused_wip.py # KDAChunkwise (~6k lines) +│ │ └── wrapper.py # flash_kda_prefill (dead path; raises on SM100 dispatch) +│ │ +│ ├── lightning/ # [non-KDA] Lightning/linear attention kernels +│ │ ├── prefill_sm100.py # Lightning Attn prefill (LinearAttentionChunkwiseDecay, lightning_attn_fwd[_varlen]) +│ │ └── decode.py # linear_attention_decode +│ └── experimental/ +│ └── linear_attn_prototype.py # [non-KDA] unwired normalized-linear-attn prototype │ -├── third_party/ -│ └── flash-linear-attention/ # FLA submodule (baseline) +├── csrc/ # CUDA C++ / CUTLASS +│ ├── api/{kda_sm90.cu, kda_sm100.cu} # PyBind11 (cula.cudac): SM90 prefill + SM100 chunk intra/recompute_w_u +│ ├── kda/sm90/ # SM90 (Hopper) KDA C++ kernels (CUTLASS 3.x, TMA/wgmma) +│ ├── kda/sm100/ # Blackwell KDA C++ kernels (CUTLASS 3.x + UMMA) +│ └── kerutils/include/ # shared C++ headers (generic device helpers sm80/sm90/sm100, host) │ -├── BENCHMARK_GB200.md # Auto-generated Blackwell benchmark results -├── BENCHMARK_H200.md # Auto-generated Hopper benchmark results -├── README.md # Project overview -├── setup.py # Build configuration -├── pyproject.toml # Project metadata -└── LICENSE +├── benchmarks/ tests/ docs/ +├── scripts/build_wheel.sh +├── third_party/flash-linear-attention/ # FLA submodule (baseline + reused gate/CP ops) +├── README.md USAGE.md REPO_LAYOUT.md RECOMMENDED_CODING_STYLE.md +└── setup.py pyproject.toml LICENSE ``` ## Key Directories | Directory | Language | Description | |-----------|----------|-------------| -| `cula/ops/` | Python (CuTe DSL) | Warp-specialized GPU kernels written in CuTe DSL — compiled to CUDA at import time | -| `cula/kda/` | Python | KDA operator dispatch — selects SM90 or SM100 path, handles chunking and autograd | -| `cula/lightning/` | Python (CuTe DSL) | Lightning Attention decode kernel | -| `csrc/kda/sm90/` | CUDA C++ | Hopper KDA kernels using CUTLASS 3.x collective API | -| `csrc/kda/sm100/` | CUDA C++ | Blackwell KDA kernels using CUTLASS 3.x + UMMA extensions | -| `csrc/api/` | CUDA C++ | PyBind11 entry points exposing C++ kernels to Python | -| `benchmarks/` | Python | Performance benchmarks vs FLA Triton baselines | -| `tests/` | Python | Correctness tests (pytest) | -| `docs/` | Markdown | Internal pipeline design notes | +| `cula/kda/` | Python | KDA **public API only** — autograd + dispatch, no kernels. Two prefill entries: modular chunk `chunk_kda` (SM100) and `kda_prefill_hopper` (SM90, driving the C++ kernel). | +| `cula/ops/kda/` | Python (CuTe DSL) | **KDA Python backends**, by arch: `sm100/` (+cp), `decode/`, `experimental/`, plus `policy.py` (CP dispatch). | +| `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTe DSL) | `[non-KDA]` Lightning/linear attention kernels. | +| `cula/ops/{inv,ptx}.py`, `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers (kept in place; not KDA-specific). | +| `csrc/kda/{sm90,sm100}/` · `csrc/api/` | CUDA C++ | Hopper SM90 prefill + Blackwell SM100 (chunk intra + recompute_w_u), exposed as `cula.cudac`. | diff --git a/USAGE.md b/USAGE.md index 86d2813c..6a4c1d8e 100644 --- a/USAGE.md +++ b/USAGE.md @@ -137,7 +137,7 @@ import os os.environ["CULA_INTRACARD_CP"] = "1" import torch -from cula.ops.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h B, T, H, K, V = 1, 65536, 8, 128, 128 device = 'cuda' diff --git a/benchmarks/bench_chunk_delta_h.py b/benchmarks/bench_chunk_delta_h.py index ad26a9b3..0ecad3c1 100644 --- a/benchmarks/bench_chunk_delta_h.py +++ b/benchmarks/bench_chunk_delta_h.py @@ -50,7 +50,7 @@ from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_mean_abs # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── -_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") +_delta_h_mod = importlib.import_module("cula.ops.kda.sm100.delta_h") chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h # ─── FLA baseline imports ─── diff --git a/benchmarks/bench_fwd_o.py b/benchmarks/bench_fwd_o.py index d9ab9bbe..4e77bc7e 100644 --- a/benchmarks/bench_fwd_o.py +++ b/benchmarks/bench_fwd_o.py @@ -47,7 +47,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) # ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── -_fwd_o_mod = importlib.import_module("cula.ops.fwd_o_sm100") +_fwd_o_mod = importlib.import_module("cula.ops.kda.sm100.fwd_o") chunk_gla_fwd_o = _fwd_o_mod.chunk_gla_fwd_o build_chunk_indices = _fwd_o_mod.build_chunk_indices diff --git a/benchmarks/bench_intracard_cp.py b/benchmarks/bench_intracard_cp.py index f24fb7bf..4eba5a48 100644 --- a/benchmarks/bench_intracard_cp.py +++ b/benchmarks/bench_intracard_cp.py @@ -47,7 +47,7 @@ set_seed, ) from cula.kda.chunk_fwd import chunk_kda_fwd -from cula.ops.cp.chunk_delta_h import ( +from cula.ops.kda.sm100.cp.chunk_delta_h import ( compute_subseq_len, prepare_subseq_cu_seqlens, should_use_intracard_cp, diff --git a/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py b/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py index 77aeb779..b938c6e0 100644 --- a/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py +++ b/benchmarks/bench_kda_bwd_wy_dqkg_sm100.py @@ -50,7 +50,7 @@ relative_rms_error_rel_max_mean_abs, set_seed, ) -from cula.ops.chunk_wy_dqkg_sm100 import chunk_kda_bwd_wy_dqkg_fused as cula_chunk_kda_bwd_wy_dqkg_fused +from cula.ops.kda.sm100.bwd_wy_dqkg import chunk_kda_bwd_wy_dqkg_fused as cula_chunk_kda_bwd_wy_dqkg_fused torch.backends.cuda.matmul.allow_tf32 = True diff --git a/benchmarks/bench_kda_decode.py b/benchmarks/bench_kda_decode.py index cf68b3a6..f1d19909 100644 --- a/benchmarks/bench_kda_decode.py +++ b/benchmarks/bench_kda_decode.py @@ -55,7 +55,7 @@ from benchmarks.utils import benchmark_cuda_fn, relative_rms_error_rel_max from cula.kda import fused_sigmoid_gating_delta_rule_update as cula_fused -from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as fla_fused +from cula.ops.kda.decode.reference_fla import fused_sigmoid_gating_delta_rule_update as fla_fused # ────────────────────────────────────────────────────────────────────── diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index b42a0f7e..e1443f49 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -18,7 +18,7 @@ Automatically selects the cuLA fully-fused implementation based on the current GPU architecture: - - sm100 (Blackwell) → cula.kda.blackwell_fused_fwd.flash_kda_prefill + - sm100 (Blackwell) → cula.ops.kda.experimental.sm100_fused.wrapper.flash_kda_prefill - sm90 (Hopper) → cula.kda.hopper_fused_fwd.cula_kda_prefill Compares: diff --git a/benchmarks/bench_la_decode_vs_fla.py b/benchmarks/bench_la_decode_vs_fla.py index 27eafbea..75f1dc79 100644 --- a/benchmarks/bench_la_decode_vs_fla.py +++ b/benchmarks/bench_la_decode_vs_fla.py @@ -56,7 +56,7 @@ from fla.ops.common.fused_recurrent import fused_recurrent_fwd, fused_recurrent_fwd_kernel from benchmarks.utils import benchmark_cuda_fn, relative_rms_error -from cula.ops.la_decode import _get_compiled_kernel, linear_attention_decode +from cula.ops.lightning.decode import _get_compiled_kernel, linear_attention_decode from cula.utils import USE_FAST_MATH diff --git a/benchmarks/bench_lightning_attn.py b/benchmarks/bench_lightning_attn.py index 17001446..bf8db7a3 100644 --- a/benchmarks/bench_lightning_attn.py +++ b/benchmarks/bench_lightning_attn.py @@ -53,7 +53,7 @@ from fla.ops.simple_gla.chunk import chunk_simple_gla_fwd from benchmarks.utils import gen_random, gen_skewed, gen_uniform, relative_rms_error, time_cuda_fn -from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen +from cula.ops.lightning.prefill_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen # ============================================================================= # Constants diff --git a/benchmarks/bench_linear_attn.py b/benchmarks/bench_linear_attn.py index 5079b7ce..5d4f70da 100644 --- a/benchmarks/bench_linear_attn.py +++ b/benchmarks/bench_linear_attn.py @@ -25,7 +25,7 @@ # from fla.ops.linear_attn.naive import naive_recurrent_linear_attn from fla.utils import assert_close, device -from cula.ops.linear_attn_sm100 import LinearAttentionChunkwise +from cula.ops.experimental.linear_attn_prototype import LinearAttentionChunkwise os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison diff --git a/cula/__init__.py b/cula/__init__.py index 6e13aa13..596f66ef 100644 --- a/cula/__init__.py +++ b/cula/__init__.py @@ -16,9 +16,3 @@ from cula._version import version as __version__ except ImportError: __version__ = "0.1.0" - -from cula.ops.lightning_attn_sm100 import LinearAttentionChunkwiseDecay - -__all__ = [ - "LinearAttentionChunkwiseDecay", -] diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index ee1a2bb9..8baa41e2 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -12,15 +12,34 @@ # See the License for the specific language governing permissions and # limitations under the License. -from cula.kda.blackwell_fused_fwd import flash_kda_prefill as kda_prefill_blackwell -from cula.kda.chunk import chunk_kda -from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper -from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode +"""Public KDA API exports for chunk, prefill, and decode""" __all__ = [ "chunk_kda", - "kda_prefill_blackwell", "kda_decode", "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", ] + +_LAZY = { + "chunk_kda": ("cula.kda.chunk", "chunk_kda"), + "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"), + "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), + "fused_sigmoid_gating_delta_rule_update": ( + "cula.ops.kda.decode.cute", + "fused_sigmoid_gating_delta_rule_update", + ), +} + + +def __getattr__(name): + target = _LAZY.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + + return getattr(importlib.import_module(target[0]), target[1]) + + +def __dir__(): + return sorted(__all__) diff --git a/cula/kda/chunk.py b/cula/kda/chunk.py index cb2df476..2c9daad6 100644 --- a/cula/kda/chunk.py +++ b/cula/kda/chunk.py @@ -15,6 +15,10 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # Related files are modified and supported by the Moonshot AI Team +"""SM100 modular chunk KDA public API and autograd wrapper""" + +from typing import Literal + import torch from fla.modules.l2norm import l2norm_bwd, l2norm_fwd from fla.ops.cp import FLACPContext @@ -23,6 +27,7 @@ from cula.kda.chunk_bwd import chunk_kda_bwd from cula.kda.chunk_fwd import chunk_kda_fwd +from cula.ops.kda.policy import IntracardCPMode, resolve_intracard_cp_mode class ChunkKDAFunction(torch.autograd.Function): @@ -50,6 +55,7 @@ def forward( disable_recompute: bool = False, return_intermediate_states: bool = False, cp_context: FLACPContext | None = None, + use_intracard_cp: IntracardCPMode | None = None, ): chunk_size = 64 @@ -85,6 +91,7 @@ def forward( disable_recompute=disable_recompute, return_intermediate_states=return_intermediate_states, cp_context=cp_context, + use_intracard_cp=use_intracard_cp, ) if return_intermediate_states: @@ -211,6 +218,7 @@ def backward( None, None, None, + None, ) @@ -233,6 +241,7 @@ def chunk_kda( disable_recompute: bool = False, return_intermediate_states: bool = False, cp_context: FLACPContext = None, + use_intracard_cp: Literal["auto"] | bool | None = None, **kwargs, ): r""" @@ -345,7 +354,15 @@ def chunk_kda( ) """ + # just for backward compatibility, resolve the deprecated `use_cp` argument + # TODO: maybe we can remove this in the future + use_cp_alias = kwargs.pop("use_cp", None) + use_intracard_cp = resolve_intracard_cp_mode(use_intracard_cp, use_cp_alias) + if cp_context is not None: + if use_intracard_cp is True: + raise ValueError("use_intracard_cp=True cannot be combined with FLA cp_context.") + use_intracard_cp = False assert initial_state is None, "Initial state is not supported for CP" assert output_final_state is False, "Output final state is not supported for CP" assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" @@ -354,6 +371,11 @@ def chunk_kda( if cp_context.cu_seqlens_cpu is not None: cu_seqlens_cpu = cp_context.cu_seqlens_cpu + if return_intermediate_states: + if use_intracard_cp is True: + raise ValueError("use_intracard_cp=True is not supported with return_intermediate_states=True.") + use_intracard_cp = False + if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -414,4 +436,5 @@ def chunk_kda( disable_recompute, return_intermediate_states, cp_context, + use_intracard_cp, ) diff --git a/cula/kda/chunk_bwd.py b/cula/kda/chunk_bwd.py index d4cf6ee7..40e56a86 100644 --- a/cula/kda/chunk_bwd.py +++ b/cula/kda/chunk_bwd.py @@ -14,6 +14,8 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +"""SM100 modular chunk KDA backward orchestration""" + import importlib import torch @@ -37,10 +39,10 @@ import cula.cudac as cula_cuda from cula.kda.chunk_intra import chunk_kda_bwd_intra -from cula.ops.chunk_wy_dqkg_sm100 import chunk_kda_bwd_wy_dqkg_fused as chunk_kda_bwd_wy_dqkg_fused_cutedsl +from cula.ops.kda.sm100.bwd_wy_dqkg import chunk_kda_bwd_wy_dqkg_fused as chunk_kda_bwd_wy_dqkg_fused_cutedsl from cula.utils import prepare_uniform_cu_seqlens -_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") +_delta_h_mod = importlib.import_module("cula.ops.kda.sm100.delta_h") chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h BK_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/cula/kda/chunk_fwd.py b/cula/kda/chunk_fwd.py index d0a93fac..f73a03c1 100644 --- a/cula/kda/chunk_fwd.py +++ b/cula/kda/chunk_fwd.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import importlib +"""SM100 modular chunk KDA forward orchestration""" import torch @@ -29,12 +29,6 @@ from cula.kda.chunk_intra import chunk_kda_fwd_intra from cula.utils import assert_blackwell -# ─── CuTe DSL wrapper (TVM-FFI compile cache) ─── -_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100") -chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h -_fwd_o_mod = importlib.import_module("cula.ops.fwd_o_sm100") -chunk_gla_fwd_o = _fwd_o_mod.chunk_gla_fwd_o - def chunk_kda_fwd( q: torch.Tensor, @@ -57,11 +51,15 @@ def chunk_kda_fwd( disable_recompute: bool = False, return_intermediate_states: bool = False, cp_context: FLACPContext | None = None, + use_intracard_cp=None, use_tf32_inverse: bool = True, unified_gref: bool = False, # Set True for ~5% extra perf (slightly lower precision) ): assert_blackwell(q.device) + from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h + from cula.ops.kda.sm100.fwd_o import chunk_gla_fwd_o + # Apply gate activation g_org = None if use_gate_in_kernel: @@ -118,6 +116,7 @@ def chunk_kda_fwd( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, cu_seqlens_cpu=cu_seqlens_cpu, + use_intracard_cp=use_intracard_cp, ) if cp_context is not None: diff --git a/cula/kda/chunk_intra.py b/cula/kda/chunk_intra.py index afc02a50..f864849c 100644 --- a/cula/kda/chunk_intra.py +++ b/cula/kda/chunk_intra.py @@ -14,6 +14,7 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +"""SM100 modular chunk KDA intra-chunk wrapper and helpers""" import torch import triton diff --git a/cula/lightning/__init__.py b/cula/lightning/__init__.py index fb5e5635..af2e3dcf 100644 --- a/cula/lightning/__init__.py +++ b/cula/lightning/__init__.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from cula.ops.la_decode import linear_attention_decode -from cula.ops.lightning_attn_sm100 import ( +from cula.ops.lightning.decode import linear_attention_decode +from cula.ops.lightning.prefill_sm100 import ( LinearAttentionChunkwiseDecay, lightning_attn_fwd, lightning_attn_fwd_varlen, diff --git a/cula/ops/__init__.py b/cula/ops/__init__.py index 6450488b..8332f620 100644 --- a/cula/ops/__init__.py +++ b/cula/ops/__init__.py @@ -12,11 +12,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode -from cula.ops.la_decode import linear_attention_decode - __all__ = [ "kda_decode", "fused_sigmoid_gating_delta_rule_update", "linear_attention_decode", ] + +_LAZY = { + "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), + "fused_sigmoid_gating_delta_rule_update": ( + "cula.ops.kda.decode.cute", + "fused_sigmoid_gating_delta_rule_update", + ), + "linear_attention_decode": ("cula.ops.lightning.decode", "linear_attention_decode"), +} + + +def __getattr__(name): + target = _LAZY.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + + return getattr(importlib.import_module(target[0]), target[1]) + + +def __dir__(): + return sorted(__all__) diff --git a/cula/ops/cp/__init__.py b/cula/ops/cp/__init__.py deleted file mode 100644 index 5727f9d4..00000000 --- a/cula/ops/cp/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from cula.ops.cp.chunk_delta_h import intracard_fwd_h - -__all__ = ["intracard_fwd_h"] diff --git a/cula/ops/experimental/__init__.py b/cula/ops/experimental/__init__.py new file mode 100644 index 00000000..cda4f1d9 --- /dev/null +++ b/cula/ops/experimental/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental / unwired non-KDA kernels.""" diff --git a/cula/ops/linear_attn_sm100.py b/cula/ops/experimental/linear_attn_prototype.py similarity index 99% rename from cula/ops/linear_attn_sm100.py rename to cula/ops/experimental/linear_attn_prototype.py index 64a11755..5cfa0248 100644 --- a/cula/ops/linear_attn_sm100.py +++ b/cula/ops/experimental/linear_attn_prototype.py @@ -42,7 +42,7 @@ .. code-block:: bash - python examples/blackwell/linear_attn_sm100.py \\ + python examples/blackwell/linear_attn.py \\ --batch_size 4 --seq_len 1024 --num_heads 8 --head_dim 64 \\ --chunk_size 64 --decay 0.95 diff --git a/cula/ops/intrinsics_sm100.py b/cula/ops/intrinsics_sm100.py deleted file mode 100644 index e7c1ce5b..00000000 --- a/cula/ops/intrinsics_sm100.py +++ /dev/null @@ -1,418 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""NVVM wrappers for SM100 (Blackwell) Tensor Memory intrinsics. - -Provides low-level, CuteDSL-compatible helpers that move data between -Tensor Memory (TMEM) and registers / shared memory via the native -``nvvm.tcgen05.*`` MLIR ops. - -**T2R / R2T** – ``tcgen05.ld`` / ``tcgen05.st`` with ``.32x32b`` shape. -**S2T** – ``tcgen05.cp`` with ``.128x256b`` shape (SMEM → TMEM) -PTX reference -------------- - tcgen05.ld.sync.aligned.32x32b.xN.b32 {r0, ..., rN-1}, [taddr]; - tcgen05.st.sync.aligned.32x32b.xN.b32 [taddr], {r0, ..., rN-1}; - -where ``N ∈ {2, 4, 8, 16, 32, 64, 128}`` and each ``r`` is a 32-bit -register. ``taddr`` encodes both the TMEM column index (bits [15:0]) -and the lane index (bits [31:16]). - -See https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld - -Usage inside a ``@cute.kernel`` or ``@cute.jit`` function:: - - from cula.ops.intrinsics_sm100 import ( - tcgen05_ld_32x32b, tcgen05_st_32x32b, - reinterpret_cast, subvec, store_256b, - ) - from cutlass.cute.typing import Float32, Int32 - - # Load 32 × 32-bit values from TMEM → opaque vector<32 x i32> - vec_i32 = tcgen05_ld_32x32b(32, taddr) - - # Zero-cost reinterpret as f32 (single vector.bitcast, no instructions) - vec_f32 = reinterpret_cast(vec_i32, Int32, 32, Float32) - - # Store to global via store_256b (4 × 256-bit stores) - # store_256b takes vector<8 x i32>, so reinterpret back and slice - vec_i32_back = reinterpret_cast(vec_f32, Float32, 32, Int32) - for chunk in range(4): # 32 / 8 = 4 chunks - store_256b(gmem_addr + chunk * 32, subvec(vec_i32_back, chunk * 8, 8)) - - # Store back to TMEM - tcgen05_st_32x32b(32, taddr, vec_i32_back) -""" - -__all__ = [ - "tcgen05_ld_32x32b", - "tcgen05_st_32x32b", - "tcgen05_cp_128x256b", - "reinterpret_cast", - "subvec", - "store_256b", - "umma_arrive", - "umma_arrive_noelect", -] - -import cutlass.cute as cute -from cutlass._mlir import ir as _ir_mod -from cutlass._mlir.dialects import arith as _arith -from cutlass._mlir.dialects import llvm -from cutlass._mlir.dialects import nvvm as _nvvm -from cutlass._mlir.dialects import vector as _vector -from cutlass.cute.arch import elect_one -from cutlass.cute.nvgpu import tcgen05 -from cutlass.cute.typing import Int32 -from cutlass.cutlass_dsl import dsl_user_op - -from cula.ops.ptx_umma_ext import Tcgen05SmemDescriptor - - -def _to_ir(val, loc=None, ip=None): - """Extract raw MLIR IR value from a CuteDSL wrapper.""" - return val.ir_value(loc=loc, ip=ip) if hasattr(val, "ir_value") else val - - -# --------------------------------------------------------------------------- -# tcgen05.ld.sync.aligned.32x32b.xN.b32 (via nvvm.tcgen05.ld) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05_ld_32x32b(num: int, taddr: int): - """Load *num* × 32-bit values from TMEM → an opaque ``vector``. - - ``num`` must be a **compile-time constant** in {2, 4, 8, 16, 32, 64, 128}. - Returns a single opaque MLIR vector value (``vector``). - - Use :func:`reinterpret_cast` to reinterpret the element type (zero-cost), - and :func:`subvec` to slice a contiguous sub-vector. - - Parameters - ---------- - num : int - Number of 32-bit registers to load. Must be a compile-time constant. - taddr : int - TMEM address (bits [31:16] = lane, bits [15:0] = column). - """ - - @dsl_user_op - def _do(addr_val, *, loc=None, ip=None): - i32_ty = _ir_mod.IntegerType.get_signless(32) - ptr6_ty = llvm.PointerType.get(address_space=6) - tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) - vec_i32_ty = _ir_mod.VectorType.get([num], i32_ty) - return _nvvm.tcgen05_ld( - res=vec_i32_ty, - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - num=num, - tmem_addr=tmem_ptr, - loc=loc, - ip=ip, - ) - - return _do(Int32(taddr)) - - -# --------------------------------------------------------------------------- -# tcgen05.st.sync.aligned.32x32b.xN.b32 (via nvvm.tcgen05.st) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05_st_32x32b(num: int, taddr: int, vec): - """Store *num* × 32-bit values from an opaque vector → TMEM. - - ``num`` must be a **compile-time constant** in {2, 4, 8, 16, 32, 64, 128}. - - Parameters - ---------- - num : int - Number of 32-bit registers to store. Must be a compile-time constant. - taddr : int - TMEM address (bits [31:16] = lane, bits [15:0] = column). - vec : opaque vector - An opaque ``vector`` value (from :func:`tcgen05_ld_32x32b` - or :func:`reinterpret_cast`). - """ - - @dsl_user_op - def _do(addr_val, vec_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) - _nvvm.tcgen05_st( - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - num=num, - tmem_addr=tmem_ptr, - r=_to_ir(vec_val, loc, ip), - loc=loc, - ip=ip, - ) - - _do(Int32(taddr), vec) - - -# --------------------------------------------------------------------------- -# reinterpret_cast (zero-cost vector.bitcast) -# --------------------------------------------------------------------------- - - -@cute.jit -def reinterpret_cast(vec, src_type, src_num, tgt_type): - """Zero-cost reinterpret of a vector's element type (single ``vector.bitcast``). - - Analogous to C++ ``reinterpret_cast``: no instructions emitted, just - re-labels the bits. The total bit-width is preserved: - ``src_num * src_type.width == tgt_num * tgt_type.width``. - - Parameters - ---------- - vec : opaque vector - Source vector (e.g. ``vector`` from :func:`tcgen05_ld_32x32b`). - src_type : CuTeDSL type - Element type of *vec* (e.g. ``Int32``). - src_num : int - Number of elements in *vec* (compile-time constant). - tgt_type : CuTeDSL type - Desired element type (e.g. ``Float32``, ``BFloat16``, ``Float16``). - - Returns - ------- - opaque vector - ``vector`` where ``M = src_num * src_type.width // tgt_type.width``. - - Examples - -------- - :: - - vec_i32 = tcgen05_ld_32x32b(8, taddr) # vector<8 x i32> - vec_f32 = reinterpret_cast(vec_i32, Int32, 8, Float32) # vector<8 x f32> - vec_bf16 = reinterpret_cast(vec_i32, Int32, 8, BFloat16) # vector<16 x bf16> - vec_back = reinterpret_cast(vec_bf16, BFloat16, 16, Int32) # vector<8 x i32> - """ - tgt_num = src_num * src_type.width // tgt_type.width - - @dsl_user_op - def _do(v, *, loc=None, ip=None): - tgt_vec_ty = _ir_mod.VectorType.get([tgt_num], tgt_type.mlir_type) - return _vector.bitcast(tgt_vec_ty, _to_ir(v, loc, ip), loc=loc, ip=ip) - - return _do(vec) - - -# --------------------------------------------------------------------------- -# subvec (extract a contiguous sub-vector) -# --------------------------------------------------------------------------- - - -@cute.jit -def subvec(vec, offset, size): - """Extract a contiguous sub-vector (``vector.extract_strided_slice``). - - Parameters - ---------- - vec : opaque vector - Source vector. - offset : int - Starting element index (compile-time constant). - size : int - Number of elements to extract (compile-time constant). - - Returns - ------- - opaque vector - ``vector``. - """ - - @dsl_user_op - def _do(v, *, loc=None, ip=None): - ir_v = _to_ir(v, loc, ip) - elem_ty = _ir_mod.VectorType(ir_v.type).element_type - res_ty = _ir_mod.VectorType.get([size], elem_ty) - return _vector.extract_strided_slice( - res_ty, - ir_v, - offsets=[offset], - sizes=[size], - strides=[1], - loc=loc, - ip=ip, - ) - - return _do(vec) - - -# --------------------------------------------------------------------------- -# st.global.L1::no_allocate.v8.f32 (256-bit direct R2G store) -# --------------------------------------------------------------------------- - -_STORE_256B_ASM = "st.global.L1::no_allocate.v8.f32 [$0], {$1, $2, $3, $4, $5, $6, $7, $8};" -_STORE_256B_CONSTRAINTS = "l,r,r,r,r,r,r,r,r" - - -@cute.jit -def store_256b(gmem_ptr, vec): - """Store 256 bits (8 × 32-bit) to global memory, bypassing L1 allocation. - - Issues ``st.global.L1::no_allocate.v8.f32`` with ``"r"`` (integer register) - constraints — type-agnostic, just like C++ ``reinterpret_cast``. - - Parameters - ---------- - gmem_ptr : pointer - Global-memory destination address (must be 32-byte aligned). - vec : opaque vector - A ``vector<8 x i32>`` (use :func:`subvec` to slice from a larger vector). - """ - - @dsl_user_op - def _do(addr, v, *, loc=None, ip=None): - i32_ty = _ir_mod.IntegerType.get_signless(32) - ir_v = _to_ir(v, loc, ip) - elems = [ - _vector.extractelement( - ir_v, - position=_arith.constant(i32_ty, i, loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - for i in range(8) - ] - operands = [_to_ir(addr, loc, ip)] + elems - llvm.inline_asm( - _ir_mod.Type.parse("!llvm.void"), - operands, - _STORE_256B_ASM, - _STORE_256B_CONSTRAINTS, - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - _do(gmem_ptr, vec) - - -# --------------------------------------------------------------------------- -# tcgen05.cp.cta_group::1.128x256b (via nvvm.tcgen05.cp) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05_cp_128x256b(taddr: int, smem_desc: Tcgen05SmemDescriptor): - """Async copy SMEM → TMEM with shape ``128x256b`` (``cta_group::1``). - - Issues ``tcgen05.cp.cta_group::1.128x256b [taddr], s-desc;`` - via the native ``nvvm.tcgen05.cp`` MLIR op. - - The instruction copies a 128-row × 256-bit tile from shared memory - (described by *smem_desc*) into Tensor Memory at *taddr*. The copy - is **asynchronous** — use ``tcgen05.commit`` + ``mbarrier.wait`` to - synchronize. - - PTX reference - ------------- - tcgen05.cp.cta_group::1.128x256b [taddr], s-desc; - - Parameters - ---------- - taddr : int - TMEM destination address (uint32, passed as ``!llvm.ptr<6>``). - smem_desc : Tcgen05SmemDescriptor - 64-bit SMEM matrix descriptor (same format as ``tcgen05.mma`` - descriptors — see ``Tcgen05SmemDescriptor``). - """ - - @dsl_user_op - def _do(addr_val, desc_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) - _nvvm.tcgen05_cp( - shape=_nvvm.Tcgen05CpShape.SHAPE_128x256b, - taddr=tmem_ptr, - smem_desc=_to_ir(desc_val, loc, ip), - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, - loc=loc, - ip=ip, - ) - - _do(Int32(taddr), smem_desc.desc_i64[0]) - - -@cute.jit -def tcgen05_cp_128x128b(taddr: int, smem_desc: Tcgen05SmemDescriptor): - """Async copy SMEM → TMEM with shape ``128x128b`` (``cta_group::1``). - - Issues ``tcgen05.cp.cta_group::1.128x128b [taddr], s-desc;`` - via the native ``nvvm.tcgen05.cp`` MLIR op. - - The instruction copies a 128-row × 128-bit tile from shared memory - (described by *smem_desc*) into Tensor Memory at *taddr*. The copy - is **asynchronous** — use ``tcgen05.commit`` + ``mbarrier.wait`` to - synchronize. - - PTX reference - ------------- - tcgen05.cp.cta_group::1.128x128b [taddr], s-desc; - - Parameters - ---------- - taddr : int - TMEM destination address (uint32, passed as ``!llvm.ptr<6>``). - smem_desc : Tcgen05SmemDescriptor - 64-bit SMEM matrix descriptor (same format as ``tcgen05.mma`` - descriptors — see ``Tcgen05SmemDescriptor``). - """ - - @dsl_user_op - def _do(addr_val, desc_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) - _nvvm.tcgen05_cp( - shape=_nvvm.Tcgen05CpShape.SHAPE_128x128b, - taddr=tmem_ptr, - smem_desc=_to_ir(desc_val, loc, ip), - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, - loc=loc, - ip=ip, - ) - - _do(Int32(taddr), smem_desc.desc_i64[0]) - - -@cute.jit -def tcgen05_fence_before(): - """tcgen05.fence::before_thread_sync — non-blocking ordering fence.""" - _nvvm.tcgen05_fence(kind=_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) - - -@cute.jit -def tcgen05_fence_after(): - """tcgen05.fence::after_thread_sync — non-blocking ordering fence.""" - _nvvm.tcgen05_fence(kind=_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC) - - -@cute.jit -def umma_arrive(mbar_ptr: cute.Pointer): - """tcgen05.commit.cta_group::1.mbarrier::arrive::one — signal MMA done.""" - with elect_one(): - tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) - - -@cute.jit -def umma_arrive_noelect(mbar_ptr: cute.Pointer): - """tcgen05.commit.cta_group::1.mbarrier::arrive::one — signal MMA done.""" - tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) diff --git a/cula/ops/kda/__init__.py b/cula/ops/kda/__init__.py new file mode 100644 index 00000000..a7d2cbd2 --- /dev/null +++ b/cula/ops/kda/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""KDA backend kernels migrated to the arch-first layout. + +sm100/ SM100 (Blackwell) modular-chunk recurrence/output/bwd kernels +decode/ single-token decode (CuTe DSL + FLA reference) +experimental/ unwired fully-fused WIP +policy.py CP dispatch policy (use_cp / use_intracard_cp) + +""" + +# TODO: The SM90 (Hopper) prefill is still the C++ kernel under csrc/kda/sm90 (CuTeDSL +# port pending); it is not yet part of this package. diff --git a/cula/ops/kda/decode/__init__.py b/cula/ops/kda/decode/__init__.py new file mode 100644 index 00000000..ee7cedbc --- /dev/null +++ b/cula/ops/kda/decode/__init__.py @@ -0,0 +1,8 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""KDA single-token decode backend (CuTe DSL) and its FLA reference.""" + +from cula.ops.kda.decode.cute import fused_sigmoid_gating_delta_rule_update, kda_decode + +__all__ = ["kda_decode", "fused_sigmoid_gating_delta_rule_update"] diff --git a/cula/ops/kda_decode.py b/cula/ops/kda/decode/cute.py similarity index 100% rename from cula/ops/kda_decode.py rename to cula/ops/kda/decode/cute.py diff --git a/cula/ops/kda_decode_fla.py b/cula/ops/kda/decode/reference_fla.py similarity index 100% rename from cula/ops/kda_decode_fla.py rename to cula/ops/kda/decode/reference_fla.py diff --git a/cula/ops/kda/experimental/__init__.py b/cula/ops/kda/experimental/__init__.py new file mode 100644 index 00000000..1badba80 --- /dev/null +++ b/cula/ops/kda/experimental/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental / unwired KDA kernels — not part of any production path.""" diff --git a/cula/ops/kda/experimental/sm100_fused/__init__.py b/cula/ops/kda/experimental/sm100_fused/__init__.py new file mode 100644 index 00000000..303f715b --- /dev/null +++ b/cula/ops/kda/experimental/sm100_fused/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""[experimental] SM100 fully-fused KDA prefill (WIP, unwired).""" diff --git a/cula/ops/kda_fully_fused_sm100_wip.py b/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py similarity index 100% rename from cula/ops/kda_fully_fused_sm100_wip.py rename to cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py diff --git a/cula/kda/blackwell_fused_fwd.py b/cula/ops/kda/experimental/sm100_fused/wrapper.py similarity index 97% rename from cula/kda/blackwell_fused_fwd.py rename to cula/ops/kda/experimental/sm100_fused/wrapper.py index c7dec95c..deec324a 100644 --- a/cula/kda/blackwell_fused_fwd.py +++ b/cula/ops/kda/experimental/sm100_fused/wrapper.py @@ -12,17 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pathlib -import sys -import warnings - -import torch +"""[experimental] Unwired SM100 fully fused KDA prefill dead-path wrapper;arch=SM100""" -sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) +import warnings import cutlass import cutlass.cute as cute import cutlass.torch as cutlass_torch +import torch from cutlass.cute.runtime import from_dlpack from fla.modules.l2norm import l2norm_fwd @@ -32,7 +29,7 @@ from fla.ops.utils.constant import RCP_LN2 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard -from cula.ops.kda_fully_fused_sm100_wip import KDAChunkwise +from cula.ops.kda.experimental.sm100_fused.kda_fully_fused_wip import KDAChunkwise from cula.utils import USE_FAST_MATH, assert_blackwell # Global kernel cache @@ -44,7 +41,7 @@ _dummy_cache = {} -class ChunkKDAFunction(torch.autograd.Function): +class BlackwellFusedKDAFunction(torch.autograd.Function): @staticmethod @input_guard @autocast_custom_fwd @@ -311,7 +308,7 @@ def flash_kda_prefill( if scale is None: scale = k.shape[-1] ** -0.5 - o, final_state = ChunkKDAFunction.apply( + o, final_state = BlackwellFusedKDAFunction.apply( q, k, v, diff --git a/cula/ops/kda/policy.py b/cula/ops/kda/policy.py new file mode 100644 index 00000000..bbd6e9b2 --- /dev/null +++ b/cula/ops/kda/policy.py @@ -0,0 +1,98 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Context-parallel dispatch policy for KDA wrappers.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass +from typing import Literal + +import torch + +IntracardCPMode = Literal["auto"] | bool + + +@dataclass(frozen=True) +class IntracardCPDecision: + enabled: bool + reason: str | None = None + force: bool = False + + +class NotSplittableError(ValueError): + """Raised when intracard CP cannot meaningfully split the given shape. + + Subclasses ValueError so existing ``except ValueError`` callers keep working, + while new code can catch it narrowly and fall back to the serial path. + """ + + +def normalize_intracard_cp_mode(mode: IntracardCPMode) -> IntracardCPMode: + # Identity checks (not `in`): `1 == True` / `0 == False` would match stray ints. + if mode != "auto" and mode is not True and mode is not False: + raise ValueError(f'use_intracard_cp must be "auto", True, or False, got {mode!r}') + return mode + + +def resolve_intracard_cp_mode( + use_intracard_cp: IntracardCPMode | None, + use_cp_alias: IntracardCPMode | None, +) -> IntracardCPMode | None: + if use_intracard_cp is not None and use_cp_alias is not None: + raise TypeError("Pass only one of use_intracard_cp or use_cp.") + mode = use_intracard_cp if use_intracard_cp is not None else use_cp_alias + if mode is None: + return None + return normalize_intracard_cp_mode(mode) + + +def _reject_or_disable(mode: IntracardCPMode, reason: str) -> IntracardCPDecision: + if mode is True: + raise ValueError(reason) + return IntracardCPDecision(False, reason) + + +def _sm100_env_cp_enabled() -> bool: + return os.environ.get("CULA_INTRACARD_CP", "0") != "0" + + +def sm100_intracard_cp_decision( + *, + mode: IntracardCPMode | None, + cu_seqlens: torch.Tensor | None, + cu_seqlens_cpu: torch.Tensor | None, + g: torch.Tensor | None, + num_qk_heads: int, + chunk_size: int, + is_inference: bool, + sm_count_provider: Callable[[], int], + no_cp: bool = False, +) -> IntracardCPDecision: + if mode is None: + mode = "auto" if _sm100_env_cp_enabled() else False + mode = normalize_intracard_cp_mode(mode) + # no_cp is the recursion guard: intracard_fwd_h re-invokes fwd_h with _no_cp=True + # so sub-sequences do not recursively re-trigger CP. + if mode is False or no_cp: + return IntracardCPDecision(False, "disabled") + + if cu_seqlens is None: + return _reject_or_disable(mode, "SM100 intracard CP requires varlen cu_seqlens.") + if g is not None: + return _reject_or_disable(mode, "SM100 intracard CP requires g is None; pass gate through gk.") + if not is_inference: + return _reject_or_disable(mode, "SM100 intracard CP is inference-only.") + + if mode is True: + return IntracardCPDecision(True, force=True) + + # auto: consult the CPU-only perf heuristic. + from cula.ops.kda.sm100.cp.chunk_delta_h import should_use_intracard_cp + + cpu = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.cpu() + if should_use_intracard_cp(cpu, sm_count_provider(), num_qk_heads, chunk_size): + return IntracardCPDecision(True) + return IntracardCPDecision(False, "SM100 intracard CP heuristic declined for this shape.") diff --git a/cula/ops/kda/sm100/__init__.py b/cula/ops/kda/sm100/__init__.py new file mode 100644 index 00000000..10e7606d --- /dev/null +++ b/cula/ops/kda/sm100/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""SM100 (Blackwell) modular-chunk KDA backend kernels.""" diff --git a/cula/ops/chunk_wy_dqkg_sm100.py b/cula/ops/kda/sm100/bwd_wy_dqkg.py similarity index 99% rename from cula/ops/chunk_wy_dqkg_sm100.py rename to cula/ops/kda/sm100/bwd_wy_dqkg.py index cafbb247..4ab6bf4c 100644 --- a/cula/ops/chunk_wy_dqkg_sm100.py +++ b/cula/ops/kda/sm100/bwd_wy_dqkg.py @@ -24,19 +24,15 @@ from cutlass.cute.typing import BFloat16, Float32, Int32, Int64 from fla.ops.utils import prepare_chunk_indices -from cula.ops.intrinsics_sm100 import ( - reinterpret_cast, - store_256b, - subvec, +from cula.ops.ptx import reinterpret_cast, store_256b, subvec +from cula.ops.sm100.ptx import ( + Tcgen05SmemDescriptor, tcgen05_fence_after, tcgen05_fence_before, tcgen05_ld_32x32b, tcgen05_st_32x32b, - umma_arrive, -) -from cula.ops.ptx_umma_ext import ( - Tcgen05SmemDescriptor, tcgen05mma_ws_ss_f16, + umma_arrive, ) from cula.utils import USE_FAST_MATH, assert_blackwell, prepare_uniform_cu_seqlens diff --git a/cula/ops/kda/sm100/cp/__init__.py b/cula/ops/kda/sm100/cp/__init__.py new file mode 100644 index 00000000..1af1a11e --- /dev/null +++ b/cula/ops/kda/sm100/cp/__init__.py @@ -0,0 +1,8 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""SM100 intracard context-parallel backend for the chunk delta-rule recurrence.""" + +from cula.ops.kda.sm100.cp.chunk_delta_h import intracard_fwd_h + +__all__ = ["intracard_fwd_h"] diff --git a/cula/ops/cp/chunk_delta_h.py b/cula/ops/kda/sm100/cp/chunk_delta_h.py similarity index 93% rename from cula/ops/cp/chunk_delta_h.py rename to cula/ops/kda/sm100/cp/chunk_delta_h.py index 6762476e..74937cd0 100644 --- a/cula/ops/cp/chunk_delta_h.py +++ b/cula/ops/kda/sm100/cp/chunk_delta_h.py @@ -26,7 +26,7 @@ Reference: - FLA intra-card CP: fla/ops/common/intracard_cp.py - FLA CP kernels: fla/ops/cp/chunk_delta_h.py - - cuLA chunk_delta_h: cula/ops/chunk_delta_h.py + - cuLA chunk_delta_h: cula/ops/kda/sm100/delta_h.py """ from __future__ import annotations @@ -40,14 +40,14 @@ from cula.utils import get_device_sm_count, get_pre_scan -# Lazy import to avoid circular dependency with cula.ops.chunk_delta_h +# Lazy import to avoid circular dependency with cula.ops.kda.sm100.delta_h _chunk_gated_delta_rule_fwd_h = None def _get_fwd_h(): global _chunk_gated_delta_rule_fwd_h if _chunk_gated_delta_rule_fwd_h is None: - from cula.ops.chunk_delta_h_sm100 import chunk_gated_delta_rule_fwd_h + from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h _chunk_gated_delta_rule_fwd_h = chunk_gated_delta_rule_fwd_h return _chunk_gated_delta_rule_fwd_h @@ -347,12 +347,12 @@ def intracard_merge( For split seq [s0, s1, ..., s_{n-1}]: h0_sj = m_{j-1} @ h0_{j-1} + he_{j-1}. Returns (initial_states_merge [num_non_first, H, K, V] fp32, num_non_first). """ - from cula.ops.cp.merge import merge_fwd + from cula.ops.kda.sm100.cp.merge import launch_merge if num_non_first == 0: return None, 0 - initial_states_merge = merge_fwd( + initial_states_merge = launch_merge( hm=hm, seq_starts=merge_seq_starts, seq_counts=merge_seq_counts, @@ -414,10 +414,12 @@ def intracard_fwd_h( chunk_indices: torch.Tensor | None = None, max_splits: int = 32, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Intra-card CP chunk_delta_h forward; drop-in replacement for chunk_gated_delta_rule_fwd_h. + """Intra-card CP chunk_delta_h forward; splits long sequences and runs + pre_scan -> merge -> fwd_h on the sub-sequences. - Splits long sequences, runs pre_scan → merge → fwd_h on sub-sequences. - Falls back to the non-CP path when guards indicate no benefit. + Pure CP executor: raises NotSplittableError when the shape cannot be + meaningfully split. The caller owns the fallback-vs-raise policy (the + pre-split heuristic lives in sm100_intracard_cp_decision, not here). """ assert cu_seqlens is not None, "intracard_fwd_h requires cu_seqlens (varlen mode)" @@ -429,21 +431,6 @@ def intracard_fwd_h( if cu_seqlens_cpu is None: cu_seqlens_cpu = cu_seqlens.cpu() - if not should_use_intracard_cp(cu_seqlens_cpu, num_sms, H, chunk_size): - return _get_fwd_h()( - k=k, - w=w, - u=u, - gk=gk, - initial_state=initial_state, - output_final_state=output_final_state, - chunk_size=chunk_size, - save_new_value=save_new_value, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - _no_cp=True, - ) - cu_list = cu_seqlens_cpu.tolist() num_seqs = len(cu_list) - 1 max_seq_len = max(cu_list[i + 1] - cu_list[i] for i in range(num_seqs)) @@ -473,19 +460,9 @@ def intracard_fwd_h( split_info = False if not split_info: - return _get_fwd_h()( - k=k, - w=w, - u=u, - gk=gk, - initial_state=initial_state, - output_final_state=output_final_state, - chunk_size=chunk_size, - save_new_value=save_new_value, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - _no_cp=True, - ) + from cula.ops.kda.policy import NotSplittableError + + raise NotSplittableError("SM100 intracard CP is not meaningfully splittable for this shape.") N_orig = len(cu_seqlens_cpu) - 1 diff --git a/cula/ops/cp/merge.py b/cula/ops/kda/sm100/cp/merge.py similarity index 83% rename from cula/ops/cp/merge.py rename to cula/ops/kda/sm100/cp/merge.py index 1136b501..75d3fd12 100644 --- a/cula/ops/cp/merge.py +++ b/cula/ops/kda/sm100/cp/merge.py @@ -23,88 +23,10 @@ import cutlass.cute as cute import cutlass.utils as utils import torch -from cutlass._mlir import ir -from cutlass._mlir.dialects import llvm as _llvm from cutlass.cute.nvgpu import cpasync from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor, make_fake_stream -from cutlass.cutlass_dsl import T as _T - - -# --------------------------------------------------------------------------- -# Inline PTX helpers: SM80 warp-level TF32 MMA (mma.sync.m16n8k8.tf32.tf32.f32) -# --------------------------------------------------------------------------- -def _to_ir(v, loc=None, ip=None): - """Convert DSL Numeric to an MLIR Value; pass through if already a Value.""" - if hasattr(v, "ir_value"): - return v.ir_value(loc=loc, ip=ip) - return v - - -@cutlass.dsl_user_op -def _cvt_f32_to_tf32(f, *, loc=None, ip=None): - """Round-to-nearest convert fp32 -> tf32 (stored as i32 bit pattern).""" - f_ir = _to_ir(f, loc=loc, ip=ip) - result = _llvm.inline_asm( - _T.i32(), - [f_ir], - "cvt.rna.tf32.f32 $0, $1;", - "=r,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return cutlass.Int32(result) - - -@cutlass.dsl_user_op -def _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): - """One mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 instruction. - - Inputs: - a0..a3: tf32 bits (Int32) — A fragment of 16x8 tile - b0..b1: tf32 bits (Int32) — B fragment of 8x8 tile - c0..c3: Float32 — accumulator in - Returns: - (d0, d1, d2, d3) Float32 — accumulator out - """ - ins = [ - _to_ir(a0, loc=loc, ip=ip), - _to_ir(a1, loc=loc, ip=ip), - _to_ir(a2, loc=loc, ip=ip), - _to_ir(a3, loc=loc, ip=ip), - _to_ir(b0, loc=loc, ip=ip), - _to_ir(b1, loc=loc, ip=ip), - _to_ir(c0, loc=loc, ip=ip), - _to_ir(c1, loc=loc, ip=ip), - _to_ir(c2, loc=loc, ip=ip), - _to_ir(c3, loc=loc, ip=ip), - ] - struct_ty = ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>") - ret = _llvm.inline_asm( - struct_ty, - ins, - "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " - "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", - "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - d0 = _llvm.extractvalue(_T.f32(), ret, [0], loc=loc, ip=ip) - d1 = _llvm.extractvalue(_T.f32(), ret, [1], loc=loc, ip=ip) - d2 = _llvm.extractvalue(_T.f32(), ret, [2], loc=loc, ip=ip) - d3 = _llvm.extractvalue(_T.f32(), ret, [3], loc=loc, ip=ip) - return ( - cutlass.Float32(d0), - cutlass.Float32(d1), - cutlass.Float32(d2), - cutlass.Float32(d3), - ) +from cula.ops.ptx import cvt_f32_to_tf32, mma_m16n8k8_tf32 # --------------------------------------------------------------------------- # Compile-time constants (thread/vector layout) @@ -116,7 +38,7 @@ def _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=No _VEC = 4 # 128-bit vectorized fp32 cp.async -class ChunkDeltaRuleMerge: +class Merge: """Prefix-scan merge kernel. H/K/V/BV kept as Python ints on ``self`` so layout construction is static. @@ -364,21 +286,21 @@ def kernel( for mi in cutlass.range_constexpr(M_TILES): row_a = warp_id * 32 + mi * 16 + q row_b = row_a + 8 - a_frag[mi, 0] = _cvt_f32_to_tf32(sM[row_a, k_base + rp]) - a_frag[mi, 1] = _cvt_f32_to_tf32(sM[row_b, k_base + rp]) - a_frag[mi, 2] = _cvt_f32_to_tf32(sM[row_a, k_base + rp + 4]) - a_frag[mi, 3] = _cvt_f32_to_tf32(sM[row_b, k_base + rp + 4]) + a_frag[mi, 0] = cvt_f32_to_tf32(sM[row_a, k_base + rp]) + a_frag[mi, 1] = cvt_f32_to_tf32(sM[row_b, k_base + rp]) + a_frag[mi, 2] = cvt_f32_to_tf32(sM[row_a, k_base + rp + 4]) + a_frag[mi, 3] = cvt_f32_to_tf32(sM[row_b, k_base + rp + 4]) # Pre-load + cvt B. For m16n8k8 TF32, B[8x8] per-lane (col-major): # b0: (rp, q) # b1: (rp+4, q) for nj in cutlass.range_constexpr(N_TILES): col_b = nj * 8 + q - b_frag[nj, 0] = _cvt_f32_to_tf32(sH[k_base + rp, col_b]) - b_frag[nj, 1] = _cvt_f32_to_tf32(sH[k_base + rp + 4, col_b]) + b_frag[nj, 0] = cvt_f32_to_tf32(sH[k_base + rp, col_b]) + b_frag[nj, 1] = cvt_f32_to_tf32(sH[k_base + rp + 4, col_b]) # MMAs for mi in cutlass.range_constexpr(M_TILES): for nj in cutlass.range_constexpr(N_TILES): - d0, d1, d2, d3 = _mma_m16n8k8_tf32( + d0, d1, d2, d3 = mma_m16n8k8_tf32( a_frag[mi, 0], a_frag[mi, 1], a_frag[mi, 2], @@ -430,7 +352,7 @@ def kernel( # Compile cache # --------------------------------------------------------------------------- def _compile_merge_variant(H: int, K: int, V: int, has_h0: int): - kernel_obj = ChunkDeltaRuleMerge(H=H, K=K, V=V, BV=_BV_DEFAULT, has_h0=has_h0) + kernel_obj = Merge(H=H, K=K, V=V, BV=_BV_DEFAULT, has_h0=has_h0) sym_s = cute.sym_int() sym_nnf = cute.sym_int() @@ -485,7 +407,7 @@ def _get_compiled_merge(H: int, K: int, V: int, has_h0: int): # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- -def merge_fwd( +def launch_merge( hm: torch.Tensor, seq_starts: list[int], seq_counts: list[int], diff --git a/cula/ops/cp/pre_scan.py b/cula/ops/kda/sm100/cp/pre_scan.py similarity index 100% rename from cula/ops/cp/pre_scan.py rename to cula/ops/kda/sm100/cp/pre_scan.py diff --git a/cula/ops/chunk_delta_h_sm100.py b/cula/ops/kda/sm100/delta_h.py similarity index 99% rename from cula/ops/chunk_delta_h_sm100.py rename to cula/ops/kda/sm100/delta_h.py index c4c84af1..c341bae9 100644 --- a/cula/ops/chunk_delta_h_sm100.py +++ b/cula/ops/kda/sm100/delta_h.py @@ -18,7 +18,6 @@ """ import argparse -import os as _os import cutlass import cutlass.cute as cute @@ -36,21 +35,12 @@ from fla.ops.utils import prepare_chunk_indices, prepare_lens from fla.utils import tensor_cache -from cula.utils import USE_FAST_MATH, assert_blackwell +from cula.ops.kda.policy import sm100_intracard_cp_decision +from cula.utils import USE_FAST_MATH, assert_blackwell, get_device_sm_count COMPILE_OPTIONS = "--enable-tvm-ffi --generate-line-info --ptxas-options '--verbose'" -# Intracard CP auto-dispatch -def _intracard_cp_enabled() -> bool: - """Return whether intracard-CP is currently enabled (runtime check). - - Env var truthiness matches FLA: any value other than "0" enables it. - Default (unset) is "0" → disabled. - """ - return _os.environ.get("CULA_INTRACARD_CP", "0") != "0" - - # in FLA, cumsum returns int64 tensor by default @tensor_cache def prepare_chunk_offsets_i32( @@ -2028,6 +2018,7 @@ def chunk_gated_delta_rule_fwd_h( persistent: bool = True, _no_cp: bool = False, cu_seqlens_cpu: torch.Tensor | None = None, + use_intracard_cp=None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """ ChunkDeltaRuleFwdH forward pass — FLA-compatible API. @@ -2059,19 +2050,23 @@ def chunk_gated_delta_rule_fwd_h( v_new: [B, T, HV, V] bf16 (or None if save_new_value=False) final_state: [N, HV, K, V] fp32 (or None if output_final_state=False) """ - # --- Intracard CP auto-dispatch --- - if _intracard_cp_enabled() and not _no_cp and cu_seqlens is not None and g is None and torch.is_inference_mode_enabled(): - from cula.ops.cp.chunk_delta_h import intracard_fwd_h, should_use_intracard_cp - from cula.utils import get_device_sm_count - - # Materialize cu_seqlens_cpu once here to avoid repeated D2H sync inside intracard_fwd_h. - _cu_seqlens_cpu = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.cpu() - if should_use_intracard_cp( - _cu_seqlens_cpu, - get_device_sm_count(k.device), - k.shape[2], - chunk_size, - ): + # --- Intracard CP dispatch (policy: cula.ops.kda.policy) --- + cp_decision = sm100_intracard_cp_decision( + mode=use_intracard_cp, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + g=g, + num_qk_heads=k.shape[2], + chunk_size=chunk_size, + is_inference=torch.is_inference_mode_enabled(), + sm_count_provider=lambda: get_device_sm_count(k.device), + no_cp=_no_cp, + ) + if cp_decision.enabled: + from cula.ops.kda.policy import NotSplittableError + from cula.ops.kda.sm100.cp.chunk_delta_h import intracard_fwd_h + + try: return intracard_fwd_h( k=k, w=w, @@ -2083,8 +2078,12 @@ def chunk_gated_delta_rule_fwd_h( save_new_value=save_new_value, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - cu_seqlens_cpu=_cu_seqlens_cpu, + cu_seqlens_cpu=cu_seqlens_cpu, ) + except NotSplittableError: + if cp_decision.force: + raise + # "auto" path: shape not splittable -- fall through to the serial body below. B, T, H, K_dim = k.shape HV = u.shape[2] diff --git a/cula/ops/fwd_o_sm100.py b/cula/ops/kda/sm100/fwd_o.py similarity index 100% rename from cula/ops/fwd_o_sm100.py rename to cula/ops/kda/sm100/fwd_o.py diff --git a/cula/ops/lightning/__init__.py b/cula/ops/lightning/__init__.py new file mode 100644 index 00000000..012655a7 --- /dev/null +++ b/cula/ops/lightning/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Lightning Attention backend kernels (prefill + decode) — non-KDA.""" diff --git a/cula/ops/la_decode.py b/cula/ops/lightning/decode.py similarity index 100% rename from cula/ops/la_decode.py rename to cula/ops/lightning/decode.py diff --git a/cula/ops/lightning_attn_sm100.py b/cula/ops/lightning/prefill_sm100.py similarity index 100% rename from cula/ops/lightning_attn_sm100.py rename to cula/ops/lightning/prefill_sm100.py diff --git a/cula/ops/ptx.py b/cula/ops/ptx.py new file mode 100644 index 00000000..595538f0 --- /dev/null +++ b/cula/ops/ptx.py @@ -0,0 +1,152 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared inline PTX and MLIR helpers for CuTeDSL kernels (SM80+).""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.dialects import vector as _vector +from cutlass.cutlass_dsl import T as _T +from cutlass.cutlass_dsl import dsl_user_op + + +def _to_ir(v, loc=None, ip=None): + if hasattr(v, "ir_value"): + return v.ir_value(loc=loc, ip=ip) + return v + + +@cutlass.dsl_user_op +def cvt_f32_to_tf32(f, *, loc=None, ip=None): + f_ir = _to_ir(f, loc=loc, ip=ip) + result = _llvm.inline_asm( + _T.i32(), + [f_ir], + "cvt.rna.tf32.f32 $0, $1;", + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@cutlass.dsl_user_op +def mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + ins = [_to_ir(x, loc=loc, ip=ip) for x in (a0, a1, a2, a3, b0, b1, c0, c1, c2, c3)] + struct_ty = ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>") + ret = _llvm.inline_asm( + struct_ty, + ins, + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = _llvm.extractvalue(_T.f32(), ret, [0], loc=loc, ip=ip) + d1 = _llvm.extractvalue(_T.f32(), ret, [1], loc=loc, ip=ip) + d2 = _llvm.extractvalue(_T.f32(), ret, [2], loc=loc, ip=ip) + d3 = _llvm.extractvalue(_T.f32(), ret, [3], loc=loc, ip=ip) + return ( + cutlass.Float32(d0), + cutlass.Float32(d1), + cutlass.Float32(d2), + cutlass.Float32(d3), + ) + + +# --------------------------------------------------------------------------- +# MLIR vector utilities (architecture-independent) +# --------------------------------------------------------------------------- + + +@cute.jit +def reinterpret_cast(vec, src_type, src_num, tgt_type): + """Zero-cost reinterpret of a vector's element type (single vector.bitcast).""" + tgt_num = src_num * src_type.width // tgt_type.width + + @dsl_user_op + def _do(v, *, loc=None, ip=None): + tgt_vec_ty = ir.VectorType.get([tgt_num], tgt_type.mlir_type) + return _vector.bitcast(tgt_vec_ty, _to_ir(v, loc, ip), loc=loc, ip=ip) + + return _do(vec) + + +@cute.jit +def subvec(vec, offset, size): + """Extract a contiguous sub-vector (vector.extract_strided_slice).""" + + @dsl_user_op + def _do(v, *, loc=None, ip=None): + ir_v = _to_ir(v, loc, ip) + elem_ty = ir.VectorType(ir_v.type).element_type + res_ty = ir.VectorType.get([size], elem_ty) + return _vector.extract_strided_slice( + res_ty, + ir_v, + offsets=[offset], + sizes=[size], + strides=[1], + loc=loc, + ip=ip, + ) + + return _do(vec) + + +_STORE_256B_ASM = "st.global.L1::no_allocate.v8.f32 [$0], {$1, $2, $3, $4, $5, $6, $7, $8};" +_STORE_256B_CONSTRAINTS = "l,r,r,r,r,r,r,r,r" + + +@cute.jit +def store_256b(gmem_ptr, vec): + """Store 256 bits (8 x 32-bit) to global memory, bypassing L1 allocation.""" + + @dsl_user_op + def _do(addr, v, *, loc=None, ip=None): + i32_ty = ir.IntegerType.get_signless(32) + ir_v = _to_ir(v, loc, ip) + elems = [ + _vector.extractelement( + ir_v, + position=_arith.constant(i32_ty, i, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + for i in range(8) + ] + operands = [_to_ir(addr, loc, ip)] + elems + _llvm.inline_asm( + ir.Type.parse("!llvm.void"), + operands, + _STORE_256B_ASM, + _STORE_256B_CONSTRAINTS, + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + _do(gmem_ptr, vec) diff --git a/cula/ops/ptx_umma_ext.py b/cula/ops/ptx_umma_ext.py deleted file mode 100644 index 2e8caeea..00000000 --- a/cula/ops/ptx_umma_ext.py +++ /dev/null @@ -1,961 +0,0 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""CuteDSL UMMA extension wrappers for SM100 (Blackwell) ``tcgen05.mma``. - -CuteDSL's high-level ``cute.gemm()`` / ``make_tiled_mma()`` API does not -expose all ``tcgen05.mma`` instruction variants. This module provides -low-level wrappers for the two categories currently needed: - -1. **Masked MMA** – SS and TS forms with the 128-bit ``disable-output-lane`` - mask operand (``{m0, m1, m2, m3}``). Implemented via the native - ``nvvm.tcgen05_mma`` MLIR op with its ``write_disable_mask`` parameter - (``vector<4xi32>``). - -2. **Weight-stationary (WS) MMA** – ``tcgen05.mma.ws`` SS / TS forms for - both ``kind::tf32`` and ``kind::f16``. Implemented via - ``llvm.inline_asm``. - ----------------------------------------------------------------------- -PTX instruction forms ----------------------------------------------------------------------- -SS (SMEM A, SMEM B): - tcgen05.mma.cta_group::1.kind::tf32 [tmem_c], desc_a, desc_b, - desc_val, {m0,m1,m2,m3}, p; - -TS (TMEM A, SMEM B): - tcgen05.mma.cta_group::1.kind::tf32 [tmem_c], [tmem_a], desc_b, - desc_val, {m0,m1,m2,m3}, p; - -WS_SS (weight-stationary, SMEM A, SMEM B): - tcgen05.mma.ws.cta_group::1.kind::tf32 [tmem_c], desc_a, desc_b, - desc_val, p; - tcgen05.mma.ws.cta_group::1.kind::f16 [tmem_c], desc_a, desc_b, - desc_val, p; - -WS_TS (weight-stationary, TMEM A, SMEM B): - tcgen05.mma.ws.cta_group::1.kind::tf32 [tmem_c], [tmem_a], desc_b, - desc_val, p; - tcgen05.mma.ws.cta_group::1.kind::f16 [tmem_c], [tmem_a], desc_b, - desc_val, p; - ----------------------------------------------------------------------- -Disable-output-lane mask layout (4 × uint32 = 128 bits) ----------------------------------------------------------------------- -Each uint32 covers 32 M-dimension rows (8 rows × 4 elements per group). - 0x00000000 → group is ACTIVE (output written) - 0xFFFFFFFF → group is DISABLED (output suppressed) - -Predefined SS mask constants (SMEM A variants): - SS_NO_MASK = (0, 0, 0, 0) all rows active - SS_MASK0 = (0, 0xFF…, 0, 0xFF…) odd groups disabled - SS_MASK1 = (0xFF…, 0, 0xFF…, 0) even groups disabled - SS_MASK2 = (0xFF…, 0xFF…, 0, 0xFF…) group 2 only active - SS_MASK3 = (0xFF…, 0xFF…, 0xFF…, 0) group 3 only active - -Predefined TS mask constants (TMEM A variants): - TS_NO_MASK = (0, 0, 0, 0) all rows active - TS_MASK0 = (0, 0xFF…, 0xFF…, 0xFF…) group 0 only active - TS_MASK1 = (0xFF…, 0, 0xFF…, 0xFF…) group 1 only active - TS_MASK2 = (0xFF…, 0xFF…, 0, 0xFF…) group 2 only active - TS_MASK3 = (0xFF…, 0xFF…, 0xFF…, 0) group 3 only active - TS_MASK02 = (0, 0xFF…, 0, 0xFF…) groups 0,2 only active - TS_MASK13 = (0xFF…, 0, 0xFF…, 0) groups 1,3 only active - -Public API (all decorated with @cute.jit) ----------------------------------------------------------------------- -Descriptor helpers (call inside @cute.jit): - Tcgen05SmemDescriptor — 64-bit SMEM descriptor object - initialize_tcgen05_descriptor — fill descriptor bitfields - -Low-level primitives (pass mask words explicitly): - tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, - mask0, mask1, mask2, mask3) - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, - mask0, mask1, mask2, mask3) - tcgen05mma_ws_ss_tf32(desc_a, desc_b, tmem_c, desc_val, scale_out) - tcgen05mma_ws_ts_tf32(tmem_a, desc_b, tmem_c, desc_val, scale_out) - tcgen05mma_ws_ss_f16(desc_a, desc_b, tmem_c, desc_val, scale_out) - tcgen05mma_ws_ts_f16(tmem_a, desc_b, tmem_c, desc_val, scale_out) - -Named convenience wrappers (pre-set masks, pass only MMA operands): - tcgen05mma_ss_no_mask / tcgen05mma_ss_mask0 / …mask1 / …mask2 / …mask3 - tcgen05mma_ts_no_mask / tcgen05mma_ts_mask0 / …mask1 / …mask2 / …mask3 - tcgen05mma_ts_mask02 / tcgen05mma_ts_mask13 -""" - -__all__ = [ - # descriptor helpers - "Tcgen05SmemDescriptor", - "initialize_tcgen05_descriptor", - # low-level primitives - "tcgen05mma_ss", - "tcgen05mma_ts", - "tcgen05mma_ws_ss_tf32", - "tcgen05mma_ws_ts_tf32", - "tcgen05mma_ws_ss_f16", - "tcgen05mma_ws_ts_f16", - # SS named wrappers - "tcgen05mma_ss_no_mask", - "tcgen05mma_ss_mask0", - "tcgen05mma_ss_mask1", - "tcgen05mma_ss_mask2", - "tcgen05mma_ss_mask3", - # TS named wrappers - "tcgen05mma_ts_no_mask", - "tcgen05mma_ts_mask0", - "tcgen05mma_ts_mask1", - "tcgen05mma_ts_mask2", - "tcgen05mma_ts_mask3", - "tcgen05mma_ts_mask02", - "tcgen05mma_ts_mask13", - # collector enums (re-exported for convenience) - "CollectorBBuffer", - "CollectorOp", -] - -import cutlass -import cutlass.cute as cute -from cutlass._mlir import ir -from cutlass._mlir.dialects import arith as _arith -from cutlass._mlir.dialects import llvm -from cutlass._mlir.dialects import nvvm as _nvvm -from cutlass.cutlass_dsl import dsl_user_op - -# Re-export collector enums for caller convenience. -CollectorBBuffer = _nvvm.Tcgen05MMACollectorBBuffer -CollectorOp = _nvvm.Tcgen05MMACollectorOp - -# --------------------------------------------------------------------------- -# Mask constants (4 × uint32). 0 = ACTIVE, 0xFFFFFFFF = DISABLED. -# --------------------------------------------------------------------------- -_ALL_ACTIVE = 0x00000000 -_ALL_OFF = 0xFFFFFFFF - -# SS masks (SMEM A, SMEM B) -SS_NO_MASK = (_ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE) -SS_MASK0 = (_ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {0,F,0,F} -SS_MASK1 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE) # {F,0,F,0} -SS_MASK2 = (_ALL_OFF, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {F,F,0,F} -SS_MASK3 = (_ALL_OFF, _ALL_OFF, _ALL_OFF, _ALL_ACTIVE) # {F,F,F,0} - -# TS masks (TMEM A, SMEM B) -TS_NO_MASK = (_ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE) -TS_MASK0 = (_ALL_ACTIVE, _ALL_OFF, _ALL_OFF, _ALL_OFF) # {0,F,F,F} -TS_MASK1 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_OFF) # {F,0,F,F} -TS_MASK2 = (_ALL_OFF, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {F,F,0,F} -TS_MASK3 = (_ALL_OFF, _ALL_OFF, _ALL_OFF, _ALL_ACTIVE) # {F,F,F,0} -TS_MASK02 = (_ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) # {0,F,0,F} -TS_MASK13 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE) # {F,0,F,0} - - -# --------------------------------------------------------------------------- -# Tcgen05SmemDescriptor — 64-bit SMEM descriptor stored as 2×Int32 -# --------------------------------------------------------------------------- - - -class Tcgen05SmemDescriptor: - """64-bit shared-memory descriptor for tcgen05 MMA (Blackwell / SM100). - - The descriptor encodes SMEM base address, leading/stride byte offsets, - swizzle mode, and other fields required by the ``tcgen05.mma`` PTX - instruction to locate a matrix tile in shared memory. - - 64-bit layout (PTX ISA Table 40):: - - Bit 63 Bit 0 - ┌──────────┬────────┬─────┬──────────┬────┬──────────┬──────┬──────────────┐ - │ 63 61 │ 60 53 │ 52 │ 51 49 │ 48 │ 45 32 │31 30 │ 29 16│15 14│ 13 0│ - │layout_typ│ reservd│l_abs│base_offst│ 46 │ SBO │ rsvd │ LBO │rsvd │start_adr│ - │ (3 bit) │ (8 bit)│(1b) │ (3 bit) │=0b001│(14 bit)│(2 b) │(14 bit)│(2b) │(14 bit) │ - └──────────┴────────┴─────┴──────────┴────┴──────────┴──────┴────────┴─────┴─────────┘ - - Field descriptions: - - - **start_address** [bits 0-13]: SMEM base pointer, encoded as - ``smem_ptr >> 4`` (16-byte aligned). The hardware reconstructs the - full address as ``encoded_value << 4``. - - - **LBO** (Leading Byte Offset) [bits 16-29]: distance in bytes between - consecutive elements along the leading dimension, encoded as - ``lbo_bytes >> 4``. When ``lbo_mode=1`` this is an absolute byte - address rather than a relative offset. - - - **SBO** (Stride Byte Offset) [bits 32-45]: distance in bytes between - consecutive elements along the stride dimension, encoded as - ``sbo_bytes >> 4``. - - - **version** [bits 46-48]: fixed constant ``0b001`` (= 1). - - - **base_offset** [bits 49-51]: 3-bit alignment correction when the - SMEM tile does not start at a natural swizzle-pattern boundary - (1024B for 128B swizzle, 512B for 64B, 256B for 32B). - Computed as ``(start_addr >> 7) & 0x7``. Usually 0. - - - **lbo_mode** (leading_abs) [bit 52]: 0 → LBO is a relative byte - offset; 1 → LBO is an absolute byte address. - - - **layout_type** (swizzle_mode) [bits 61-63]: - - 0 = SWIZZLE_NONE - - 1 = SWIZZLE_128B_BASE32B (128-byte pattern, 32-byte atom) - - 2 = SWIZZLE_128B (128-byte pattern) - - 4 = SWIZZLE_64B (64-byte pattern) - - 6 = SWIZZLE_32B (32-byte pattern) - - Storage: two Int32 registers (desc[0] = low 32 bits, desc[1] = high 32 - bits), recast to a single Int64 for the PTX ``l``-constraint operand. - - Usage inside a @cute.jit kernel:: - - desc = Tcgen05SmemDescriptor() - initialize_tcgen05_descriptor(desc, smem_ptr, lbo, sbo, 0, True, swizzle) - """ - - def __init__(self, desc_64: cute.Int64 = None): - # desc[0]: low 32 bits → start_address[0:14] | LBO[16:30] - # desc[1]: high 32 bits → SBO[0:14] | version[14:16] | base_offset[17:20] - # | lbo_mode[20] | layout_type[29:32] - self.desc = cute.make_rmem_tensor((2,), dtype=cutlass.Int32) - # Alias the 2×i32 as 1×i64 for PTX "l" constraint (64-bit operand) - self.desc_i64 = cute.make_tensor(cute.recast_ptr(self.desc.iterator, dtype=cute.Int64), (1,)) - if desc_64 is not None: - self.desc_i64[0] = desc_64 - - def __add__(self, byte_offset): - """Return a new descriptor offset by ``byte_offset`` bytes. - - Only the start_address field (bits 0-13 of desc[0]) is modified. - Since it is stored in 16-byte units, we add ``byte_offset >> 4``. - All other fields (LBO, SBO, swizzle, etc.) are copied unchanged. - """ - res = cute.make_rmem_tensor((2,), dtype=cutlass.Int32) - res_i64 = cute.make_tensor(cute.recast_ptr(res.iterator, dtype=cute.Int64), (1,)) - res[0] = self.desc[0] + (byte_offset >> 4) # adjust start_address - res[1] = self.desc[1] # high word unchanged - return Tcgen05SmemDescriptor(res_i64[0]) - - -# --------------------------------------------------------------------------- -# initialize_tcgen05_descriptor -# --------------------------------------------------------------------------- - - -def initialize_tcgen05_descriptor( - desc, - start_address, - leading_byte_offset, - stride_byte_offset, - base_offset, - leading_abs, - swizzle_mode, -): - """Pack SMEM descriptor bitfields into *desc* (a Tcgen05SmemDescriptor). - - Constructs the 64-bit descriptor in two 32-bit halves (desc[0] and desc[1]). - All address/offset fields must be pre-divided by 16 (``>> 4``) before - passing, because the hardware stores them in 16-byte granularity. - - Low 32 bits — desc[0]:: - - ┌────────────────┬──────┬──────────────────┐ - │ bits 29…16 │15…14 │ bits 13…0 │ - │ LBO (14 bits) │ rsvd │ start_addr >> 4 │ - └────────────────┴──────┴──────────────────┘ - - - [0:14) start_address >> 4 — SMEM tile base pointer in 16B units. - - [14:16) reserved (0). - - [16:30) leading_byte_offset — LBO in 16B units (caller passes >> 4). - - High 32 bits — desc[1]:: - - ┌────────┬────────┬─────┬──────────┬────────┬──────────────────┐ - │ 31…29 │ 28…21 │ 20 │ 19…17 │ 16…14 │ bits 13…0 │ - │ layout │ rsvd │l_abs│base_off │version │ SBO (14 bits) │ - │ (3 bit)│ (8 bit)│(1b) │ (3 bit) │=0b001 │ │ - └────────┴────────┴─────┴──────────┴────────┴──────────────────┘ - - - [0:14) stride_byte_offset — SBO in 16B units (caller passes >> 4). - - [14:16) version = 1 (fixed constant 0b001, only bit 14 set). - - [17:20) base_offset & 0x7 — swizzle alignment correction. - Typically 0. Non-zero when the tile doesn't start at - the natural swizzle boundary (1024B/512B/256B). - - [20:21) lbo_mode — 0 = LBO is relative offset, 1 = absolute address. - - [29:32) layout_type (swizzle_mode & 0x7): - 0 = SWIZZLE_NONE - 1 = SWIZZLE_128B_BASE32B (Swizzle<2,5,2>) - 2 = SWIZZLE_128B (Swizzle<3,4,3>) - 4 = SWIZZLE_64B (Swizzle<2,4,3>) - 6 = SWIZZLE_32B (Swizzle<1,4,3>) - - Args: - desc: Tcgen05SmemDescriptor to fill. - start_address: CuTeDSL Pointer to the SMEM tile start. - leading_byte_offset: Leading-dimension byte offset, already >> 4. - stride_byte_offset: Stride byte offset, already >> 4. - base_offset: Swizzle alignment correction (raw int, bits 17-19). - leading_abs: Bool — True → LBO is absolute address. - swizzle_mode: Swizzle layout_type integer (bits 29-31). - """ - # Encode start_address: take SMEM pointer, shift right by 4 to get 16B units - ptr_val = start_address.toint() >> 4 - - # --- Low 32 bits (desc[0]) --- - # bits [0:14) = start_address >> 4 - # bits [16:30) = leading_byte_offset (already in 16B units) - desc.desc[0] = cutlass.Int32(ptr_val) | cutlass.Int32(cutlass.Int32(leading_byte_offset) << 16) - - # --- High 32 bits (desc[1]) --- - # bits [0:14) = stride_byte_offset (already in 16B units) - # bit [14] = version = 1 (fixed) - # bits [17:20) = base_offset & 0x7 (swizzle alignment correction) - # bit [20] = lbo_mode (0=relative, 1=absolute) - # bits [29:32) = layout_type (swizzle mode) - desc.desc[1] = ( - cutlass.Int32(stride_byte_offset) - | cutlass.Int32(1 << 14) # version = 1 - | cutlass.Int32(cutlass.Int32(base_offset & 0x7) << 17) - | cutlass.Int32(cutlass.Int32(int(leading_abs)) << 20) - | cutlass.Int32(cutlass.Int32(swizzle_mode & 0x7) << 29) - ) - - -# --------------------------------------------------------------------------- -# Internal helper -# --------------------------------------------------------------------------- - - -def _ir(val, loc=None, ip=None): - """Extract raw MLIR IR value from a CuTeDSL wrapper.""" - return val.ir_value(loc=loc, ip=ip) if hasattr(val, "ir_value") else val - - -# =========================================================================== -# Low-level primitives -# =========================================================================== - -# --------------------------------------------------------------------------- -# tcgen05mma_ss — SMEM A, SMEM B (non-warp-specialised) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ss( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, - mask0: int, - mask1: int, - mask2: int, - mask3: int, -): - """Issue ``tcgen05.mma.cta_group::1.kind::tf32`` with SMEM operands. - - ``mask{0-3}`` are the four uint32 words of the 128-bit - ``disable-output-lane`` mask (0=active, 0xFFFFFFFF=disabled). - - Caller must ensure single-thread execution (e.g. via ``elect_one``); - no internal ``elect.sync`` is performed. - - Args: - desc_a: 64-bit SMEM descriptor for matrix A. - desc_b: 64-bit SMEM descriptor for matrix B. - tmem_c: TMEM base address (uint32) for accumulators C/D. - desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). - scale_out: 1 → accumulate into C, 0 → overwrite C (clear accumulators). - mask0-3: Four uint32 words of the disable-output-lane mask. - """ - - @dsl_user_op - def _do(c_val, da_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - i32_ty = ir.IntegerType.get_signless(32) - i1_ty = ir.IntegerType.get_signless(1) - vec4i32_ty = ir.VectorType.get([4], i32_ty) - - c_ir = _ir(c_val, loc, ip) - d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) - da_ir = _ir(da_val, loc, ip) # i64 SMEM descriptor - db_ir = _ir(db_val, loc, ip) # i64 SMEM descriptor - dv_ir = _ir(dv_val, loc, ip) - sc_ir = _ir(sc_val, loc, ip) - enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - - m0_ir = _ir(m0_val, loc, ip) - m1_ir = _ir(m1_val, loc, ip) - m2_ir = _ir(m2_val, loc, ip) - m3_ir = _ir(m3_val, loc, ip) - - undef = llvm.mlir_undef(vec4i32_ty, loc=loc, ip=ip) - idx0 = _arith.constant(i32_ty, 0, loc=loc, ip=ip) - idx1 = _arith.constant(i32_ty, 1, loc=loc, ip=ip) - idx2 = _arith.constant(i32_ty, 2, loc=loc, ip=ip) - idx3 = _arith.constant(i32_ty, 3, loc=loc, ip=ip) - v = llvm.InsertElementOp(undef, m0_ir, idx0, loc=loc, ip=ip) - v = llvm.InsertElementOp(v, m1_ir, idx1, loc=loc, ip=ip) - v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) - mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) - - _nvvm.tcgen05_mma( - mma_kind=_nvvm.Tcgen05MMAKind.TF32, - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, - d=d_ptr, - a=da_ir, - b=db_ir, - idesc=dv_ir, - enable_input_d=enable_d, - write_disable_mask=mask, - loc=loc, - ip=ip, - ) - - _do( - cutlass.Int32(tmem_c), - desc_a.desc_i64[0], - desc_b.desc_i64[0], - cutlass.Int32(desc_val), - cutlass.Int32(scale_out), - cutlass.Int32(mask0), - cutlass.Int32(mask1), - cutlass.Int32(mask2), - cutlass.Int32(mask3), - ) - - -# --------------------------------------------------------------------------- -# tcgen05mma_ts — TMEM A, SMEM B (non-warp-specialised) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ts( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, - mask0: int, - mask1: int, - mask2: int, - mask3: int, -): - """Issue ``tcgen05.mma.cta_group::1.kind::tf32`` with TMEM A operand. - - Matrix A is read from TMEM via indirect addressing ``[tmem_a]``. - Matrix B is read from SMEM via descriptor. - Caller must ensure single-thread execution (e.g. via ``elect_one``). - - Args: - tmem_a: TMEM base address (uint32) for matrix A. - desc_b: 64-bit SMEM descriptor for matrix B. - tmem_c: TMEM base address (uint32) for accumulators C/D. - desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). - scale_out: 1 → accumulate into C, 0 → overwrite C. - mask0-3: Four uint32 words of the disable-output-lane mask. - """ - - @dsl_user_op - def _do(c_val, a_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - i32_ty = ir.IntegerType.get_signless(32) - i1_ty = ir.IntegerType.get_signless(1) - vec4i32_ty = ir.VectorType.get([4], i32_ty) - - c_ir = _ir(c_val, loc, ip) - a_ir = _ir(a_val, loc, ip) - d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) - a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) - b_ir = _ir(db_val, loc, ip) - dv_ir = _ir(dv_val, loc, ip) - sc_ir = _ir(sc_val, loc, ip) - enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - - m0_ir = _ir(m0_val, loc, ip) - m1_ir = _ir(m1_val, loc, ip) - m2_ir = _ir(m2_val, loc, ip) - m3_ir = _ir(m3_val, loc, ip) - - undef = llvm.mlir_undef(vec4i32_ty, loc=loc, ip=ip) - idx0 = _arith.constant(i32_ty, 0, loc=loc, ip=ip) - idx1 = _arith.constant(i32_ty, 1, loc=loc, ip=ip) - idx2 = _arith.constant(i32_ty, 2, loc=loc, ip=ip) - idx3 = _arith.constant(i32_ty, 3, loc=loc, ip=ip) - v = llvm.InsertElementOp(undef, m0_ir, idx0, loc=loc, ip=ip) - v = llvm.InsertElementOp(v, m1_ir, idx1, loc=loc, ip=ip) - v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) - mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) - - _nvvm.tcgen05_mma( - mma_kind=_nvvm.Tcgen05MMAKind.TF32, - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, - d=d_ptr, - a=a_ptr, - b=b_ir, - idesc=dv_ir, - enable_input_d=enable_d, - write_disable_mask=mask, - loc=loc, - ip=ip, - ) - - _do( - cutlass.Int32(tmem_c), - cutlass.Int32(tmem_a), - desc_b.desc_i64[0], - cutlass.Int32(desc_val), - cutlass.Int32(scale_out), - cutlass.Int32(mask0), - cutlass.Int32(mask1), - cutlass.Int32(mask2), - cutlass.Int32(mask3), - ) - - -# --------------------------------------------------------------------------- -# tcgen05mma_ws_ss_tf32 — weight-stationary, SMEM A, SMEM B, kind::tf32 -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ws_ss_tf32( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, - collector_b_buffer=None, - collector_op=None, -): - """Issue ``tcgen05.mma.ws.cta_group::1.kind::tf32`` (weight-stationary form). - - This variant does NOT take a ``disable-output-lane`` mask; the - optional ``zero-column-mask-desc`` operand is omitted. - - Args: - desc_a: 64-bit SMEM descriptor for matrix A. - desc_b: 64-bit SMEM descriptor for matrix B. - tmem_c: TMEM base address (uint32) for accumulators C/D. - desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). - scale_out: 1 → accumulate, 0 → overwrite. - collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). - Defaults to None (hardware default: ``b0::discard``). - collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). - Defaults to None (hardware default: discard). - """ - - @dsl_user_op - def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - i1_ty = ir.IntegerType.get_signless(1) - - c_ir = _ir(c_val, loc, ip) - d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) - da_ir = _ir(da_val, loc, ip) - db_ir = _ir(db_val, loc, ip) - dv_ir = _ir(dv_val, loc, ip) - sc_ir = _ir(sc_val, loc, ip) - enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - - _nvvm.tcgen05_mma_ws( - mma_kind=_nvvm.Tcgen05MMAKind.TF32, - d=d_ptr, - a=da_ir, - b=db_ir, - idesc=dv_ir, - enable_input_d=enable_d, - collector_b_buffer=collector_b_buffer, - collector_op=collector_op, - loc=loc, - ip=ip, - ) - - _do( - cutlass.Int32(tmem_c), - desc_a.desc_i64[0], - desc_b.desc_i64[0], - cutlass.Int32(desc_val), - cutlass.Int32(scale_out), - ) - - -# --------------------------------------------------------------------------- -# tcgen05mma_ws_ss_f16 — weight-stationary, SMEM A, SMEM B, kind::f16 -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ws_ss_f16( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, - collector_b_buffer=None, - collector_op=None, -): - """Issue ``tcgen05.mma.ws.cta_group::1.kind::f16`` (weight-stationary form). - - Same as the tf32 variant but uses ``.kind::f16`` for half-precision - input types (f16 / bf16). K dimension is 16 instead of 8. - - This variant does NOT take a ``disable-output-lane`` mask; the - optional ``zero-column-mask-desc`` operand is omitted. - - Args: - desc_a: 64-bit SMEM descriptor for matrix A. - desc_b: 64-bit SMEM descriptor for matrix B. - tmem_c: TMEM base address (uint32) for accumulators C/D. - desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). - scale_out: 1 → accumulate, 0 → overwrite. - collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). - Defaults to None (hardware default: ``b0::discard``). - collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). - Defaults to None (hardware default: discard). - """ - - @dsl_user_op - def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - i1_ty = ir.IntegerType.get_signless(1) - - c_ir = _ir(c_val, loc, ip) - d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) - da_ir = _ir(da_val, loc, ip) - db_ir = _ir(db_val, loc, ip) - dv_ir = _ir(dv_val, loc, ip) - sc_ir = _ir(sc_val, loc, ip) - enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - - _nvvm.tcgen05_mma_ws( - mma_kind=_nvvm.Tcgen05MMAKind.F16, - d=d_ptr, - a=da_ir, - b=db_ir, - idesc=dv_ir, - enable_input_d=enable_d, - collector_b_buffer=collector_b_buffer, - collector_op=collector_op, - loc=loc, - ip=ip, - ) - - _do( - cutlass.Int32(tmem_c), - desc_a.desc_i64[0], - desc_b.desc_i64[0], - cutlass.Int32(desc_val), - cutlass.Int32(scale_out), - ) - - -# --------------------------------------------------------------------------- -# tcgen05mma_ws_ts_tf32 — weight-stationary, TMEM A, SMEM B, kind::tf32 -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ws_ts_tf32( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, - collector_b_buffer=None, - collector_op=None, -): - """Issue ``tcgen05.mma.ws.cta_group::1.kind::tf32`` with TMEM A (weight-stationary). - - Matrix A is read from TMEM via indirect addressing ``[tmem_a]``. - Matrix B is read from SMEM via descriptor. - This variant does NOT take a ``disable-output-lane`` mask; the - optional ``zero-column-mask-desc`` operand is omitted. - - Args: - tmem_a: TMEM base address (uint32) for matrix A. - desc_b: 64-bit SMEM descriptor for matrix B. - tmem_c: TMEM base address (uint32) for accumulators C/D. - desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). - scale_out: 1 → accumulate, 0 → overwrite. - collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). - Defaults to None (hardware default: ``b0::discard``). - collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). - Defaults to None (hardware default: discard). - """ - - @dsl_user_op - def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - i1_ty = ir.IntegerType.get_signless(1) - - c_ir = _ir(c_val, loc, ip) - d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) - a_ir = _ir(a_val, loc, ip) - a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) - db_ir = _ir(db_val, loc, ip) - dv_ir = _ir(dv_val, loc, ip) - sc_ir = _ir(sc_val, loc, ip) - enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - - _nvvm.tcgen05_mma_ws( - mma_kind=_nvvm.Tcgen05MMAKind.TF32, - d=d_ptr, - a=a_ptr, - b=db_ir, - idesc=dv_ir, - enable_input_d=enable_d, - collector_b_buffer=collector_b_buffer, - collector_op=collector_op, - loc=loc, - ip=ip, - ) - - _do( - cutlass.Int32(tmem_c), - cutlass.Int32(tmem_a), - desc_b.desc_i64[0], - cutlass.Int32(desc_val), - cutlass.Int32(scale_out), - ) - - -# --------------------------------------------------------------------------- -# tcgen05mma_ws_ts_f16 — weight-stationary, TMEM A, SMEM B, kind::f16 -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ws_ts_f16( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, - collector_b_buffer=None, - collector_op=None, -): - """Issue ``tcgen05.mma.ws.cta_group::1.kind::f16`` with TMEM A (weight-stationary). - - Same as the tf32 variant but uses ``.kind::f16`` for half-precision - input types (f16 / bf16). K dimension is 16 instead of 8. - - Matrix A is read from TMEM via indirect addressing ``[tmem_a]``. - Matrix B is read from SMEM via descriptor. - This variant does NOT take a ``disable-output-lane`` mask; the - optional ``zero-column-mask-desc`` operand is omitted. - - Args: - tmem_a: TMEM base address (uint32) for matrix A. - desc_b: 64-bit SMEM descriptor for matrix B. - tmem_c: TMEM base address (uint32) for accumulators C/D. - desc_val: High 32 bits of the UMMA instruction descriptor (idescE>>32). - scale_out: 1 → accumulate, 0 → overwrite. - collector_b_buffer: Optional ``CollectorBBuffer`` enum (B0–B3). - Defaults to None (hardware default: ``b0::discard``). - collector_op: Optional ``CollectorOp`` enum (FILL/USE/LASTUSE/DISCARD). - Defaults to None (hardware default: discard). - """ - - @dsl_user_op - def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): - ptr6_ty = llvm.PointerType.get(address_space=6) - i1_ty = ir.IntegerType.get_signless(1) - - c_ir = _ir(c_val, loc, ip) - d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) - a_ir = _ir(a_val, loc, ip) - a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) - db_ir = _ir(db_val, loc, ip) - dv_ir = _ir(dv_val, loc, ip) - sc_ir = _ir(sc_val, loc, ip) - enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - - _nvvm.tcgen05_mma_ws( - mma_kind=_nvvm.Tcgen05MMAKind.F16, - d=d_ptr, - a=a_ptr, - b=db_ir, - idesc=dv_ir, - enable_input_d=enable_d, - collector_b_buffer=collector_b_buffer, - collector_op=collector_op, - loc=loc, - ip=ip, - ) - - _do( - cutlass.Int32(tmem_c), - cutlass.Int32(tmem_a), - desc_b.desc_i64[0], - cutlass.Int32(desc_val), - cutlass.Int32(scale_out), - ) - - -# =========================================================================== -# Named convenience wrappers -# =========================================================================== -# These call the low-level primitives with pre-set mask constants so callers -# do not need to repeat the literal values. Signature: same as the base -# function but without the mask0-3 args. - -# --------------------------------------------------------------------------- -# SS named wrappers (SMEM A) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ss_no_mask( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """SS MMA with no output-lane disable (all rows active).""" - tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_NO_MASK[0], SS_NO_MASK[1], SS_NO_MASK[2], SS_NO_MASK[3]) - - -@cute.jit -def tcgen05mma_ss_mask0( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """SS MMA: mask={0, 0xF…, 0, 0xF…} — groups 0,2 active (1,3 disabled).""" - tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK0[0], SS_MASK0[1], SS_MASK0[2], SS_MASK0[3]) - - -@cute.jit -def tcgen05mma_ss_mask1( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """SS MMA: mask={0xF…, 0, 0xF…, 0} — groups 1,3 active (0,2 disabled).""" - tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK1[0], SS_MASK1[1], SS_MASK1[2], SS_MASK1[3]) - - -@cute.jit -def tcgen05mma_ss_mask2( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """SS MMA: mask={0xF…, 0xF…, 0, 0xF…} — group 2 only active.""" - tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK2[0], SS_MASK2[1], SS_MASK2[2], SS_MASK2[3]) - - -@cute.jit -def tcgen05mma_ss_mask3( - desc_a: Tcgen05SmemDescriptor, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """SS MMA: mask={0xF…, 0xF…, 0xF…, 0} — group 3 only active.""" - tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK3[0], SS_MASK3[1], SS_MASK3[2], SS_MASK3[3]) - - -# --------------------------------------------------------------------------- -# TS named wrappers (TMEM A) -# --------------------------------------------------------------------------- - - -@cute.jit -def tcgen05mma_ts_no_mask( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA with no output-lane disable (all rows active).""" - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_NO_MASK[0], TS_NO_MASK[1], TS_NO_MASK[2], TS_NO_MASK[3]) - - -@cute.jit -def tcgen05mma_ts_mask0( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA: mask={0, 0xF…, 0xF…, 0xF…} — group 0 only active.""" - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK0[0], TS_MASK0[1], TS_MASK0[2], TS_MASK0[3]) - - -@cute.jit -def tcgen05mma_ts_mask1( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA: mask={0xF…, 0, 0xF…, 0xF…} — group 1 only active.""" - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK1[0], TS_MASK1[1], TS_MASK1[2], TS_MASK1[3]) - - -@cute.jit -def tcgen05mma_ts_mask2( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA: mask={0xF…, 0xF…, 0, 0xF…} — group 2 only active.""" - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK2[0], TS_MASK2[1], TS_MASK2[2], TS_MASK2[3]) - - -@cute.jit -def tcgen05mma_ts_mask3( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA: mask={0xF…, 0xF…, 0xF…, 0} — group 3 only active.""" - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK3[0], TS_MASK3[1], TS_MASK3[2], TS_MASK3[3]) - - -@cute.jit -def tcgen05mma_ts_mask02( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA: mask={0, 0xF…, 0, 0xF…} — groups 0,2 active (1,3 disabled). - - Used in the KDA intra-chunk backward kernel for the QK/KG phase where - only even row-groups of the M tile contribute to the triangular region. - """ - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK02[0], TS_MASK02[1], TS_MASK02[2], TS_MASK02[3]) - - -@cute.jit -def tcgen05mma_ts_mask13( - tmem_a: int, - desc_b: Tcgen05SmemDescriptor, - tmem_c: int, - desc_val: int, - scale_out: int, -): - """TS MMA: mask={0xF…, 0, 0xF…, 0} — groups 1,3 active (0,2 disabled). - - Used in the KDA intra-chunk backward kernel for the QK/KG phase where - only odd row-groups of the M tile contribute to the triangular region. - """ - tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK13[0], TS_MASK13[1], TS_MASK13[2], TS_MASK13[3]) diff --git a/cula/ops/sm100/__init__.py b/cula/ops/sm100/__init__.py new file mode 100644 index 00000000..c90887f8 --- /dev/null +++ b/cula/ops/sm100/__init__.py @@ -0,0 +1,2 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 diff --git a/cula/ops/sm100/ptx.py b/cula/ops/sm100/ptx.py new file mode 100644 index 00000000..1e95b412 --- /dev/null +++ b/cula/ops/sm100/ptx.py @@ -0,0 +1,788 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SM100 (Blackwell) Tensor Memory intrinsics and UMMA extensions for CuTeDSL.""" + +__all__ = [ + # TMEM load/store/copy + "tcgen05_ld_32x32b", + "tcgen05_st_32x32b", + "tcgen05_cp_128x256b", + "tcgen05_cp_128x128b", + "tcgen05_fence_before", + "tcgen05_fence_after", + "umma_arrive", + "umma_arrive_noelect", + # descriptor helpers + "Tcgen05SmemDescriptor", + "initialize_tcgen05_descriptor", + # low-level MMA primitives + "tcgen05mma_ss", + "tcgen05mma_ts", + "tcgen05mma_ws_ss_tf32", + "tcgen05mma_ws_ts_tf32", + "tcgen05mma_ws_ss_f16", + "tcgen05mma_ws_ts_f16", + # SS named wrappers + "tcgen05mma_ss_no_mask", + "tcgen05mma_ss_mask0", + "tcgen05mma_ss_mask1", + "tcgen05mma_ss_mask2", + "tcgen05mma_ss_mask3", + # TS named wrappers + "tcgen05mma_ts_no_mask", + "tcgen05mma_ts_mask0", + "tcgen05mma_ts_mask1", + "tcgen05mma_ts_mask2", + "tcgen05mma_ts_mask3", + "tcgen05mma_ts_mask02", + "tcgen05mma_ts_mask13", + # collector enums (re-exported for convenience) + "CollectorBBuffer", + "CollectorOp", +] + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import nvvm as _nvvm +from cutlass.cute.arch import elect_one +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.typing import Int32 +from cutlass.cutlass_dsl import dsl_user_op + +CollectorBBuffer = _nvvm.Tcgen05MMACollectorBBuffer +CollectorOp = _nvvm.Tcgen05MMACollectorOp + + +def _to_ir(val, loc=None, ip=None): + return val.ir_value(loc=loc, ip=ip) if hasattr(val, "ir_value") else val + + +# =========================================================================== +# Tcgen05SmemDescriptor — 64-bit SMEM descriptor stored as 2×Int32 +# =========================================================================== + + +class Tcgen05SmemDescriptor: + """64-bit shared-memory descriptor for tcgen05 MMA (Blackwell / SM100). + + The descriptor encodes SMEM base address, leading/stride byte offsets, + swizzle mode, and other fields required by the ``tcgen05.mma`` PTX + instruction to locate a matrix tile in shared memory. + + 64-bit layout (PTX ISA Table 40):: + + Bit 63 Bit 0 + ┌──────────┬────────┬─────┬──────────┬────┬──────────┬──────┬──────────────┐ + │ 63 61 │ 60 53 │ 52 │ 51 49 │ 48 │ 45 32 │31 30 │ 29 16│15 14│ 13 0│ + │layout_typ│ reservd│l_abs│base_offst│ 46 │ SBO │ rsvd │ LBO │rsvd │start_adr│ + │ (3 bit) │ (8 bit)│(1b) │ (3 bit) │=0b001│(14 bit)│(2 b) │(14 bit)│(2b) │(14 bit) │ + └──────────┴────────┴─────┴──────────┴────┴──────────┴──────┴────────┴─────┴─────────┘ + + Storage: two Int32 registers (desc[0] = low 32 bits, desc[1] = high 32 + bits), recast to a single Int64 for the PTX ``l``-constraint operand. + + Usage inside a @cute.jit kernel:: + + desc = Tcgen05SmemDescriptor() + initialize_tcgen05_descriptor(desc, smem_ptr, lbo, sbo, 0, True, swizzle) + """ + + def __init__(self, desc_64: cute.Int64 = None): + self.desc = cute.make_rmem_tensor((2,), dtype=cutlass.Int32) + self.desc_i64 = cute.make_tensor(cute.recast_ptr(self.desc.iterator, dtype=cute.Int64), (1,)) + if desc_64 is not None: + self.desc_i64[0] = desc_64 + + def __add__(self, byte_offset): + """Return a new descriptor offset by ``byte_offset`` bytes.""" + res = cute.make_rmem_tensor((2,), dtype=cutlass.Int32) + res_i64 = cute.make_tensor(cute.recast_ptr(res.iterator, dtype=cute.Int64), (1,)) + res[0] = self.desc[0] + (byte_offset >> 4) + res[1] = self.desc[1] + return Tcgen05SmemDescriptor(res_i64[0]) + + +def initialize_tcgen05_descriptor( + desc, + start_address, + leading_byte_offset, + stride_byte_offset, + base_offset, + leading_abs, + swizzle_mode, +): + """Pack SMEM descriptor bitfields into *desc* (a Tcgen05SmemDescriptor). + + All address/offset fields must be pre-divided by 16 (``>> 4``) before + passing, because the hardware stores them in 16-byte granularity. + + Args: + desc: Tcgen05SmemDescriptor to fill. + start_address: CuTeDSL Pointer to the SMEM tile start. + leading_byte_offset: Leading-dimension byte offset, already >> 4. + stride_byte_offset: Stride byte offset, already >> 4. + base_offset: Swizzle alignment correction (raw int, bits 17-19). + leading_abs: Bool — True → LBO is absolute address. + swizzle_mode: Swizzle layout_type integer (bits 29-31). + """ + ptr_val = start_address.toint() >> 4 + + desc.desc[0] = cutlass.Int32(ptr_val) | cutlass.Int32(cutlass.Int32(leading_byte_offset) << 16) + + desc.desc[1] = ( + cutlass.Int32(stride_byte_offset) + | cutlass.Int32(1 << 14) + | cutlass.Int32(cutlass.Int32(base_offset & 0x7) << 17) + | cutlass.Int32(cutlass.Int32(int(leading_abs)) << 20) + | cutlass.Int32(cutlass.Int32(swizzle_mode & 0x7) << 29) + ) + + +# =========================================================================== +# TMEM load / store / copy (tcgen05.ld / tcgen05.st / tcgen05.cp) +# =========================================================================== + + +@cute.jit +def tcgen05_ld_32x32b(num: int, taddr: int): + """Load *num* × 32-bit values from TMEM → an opaque ``vector``.""" + + @dsl_user_op + def _do(addr_val, *, loc=None, ip=None): + i32_ty = ir.IntegerType.get_signless(32) + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + vec_i32_ty = ir.VectorType.get([num], i32_ty) + return _nvvm.tcgen05_ld( + res=vec_i32_ty, + shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, + num=num, + tmem_addr=tmem_ptr, + loc=loc, + ip=ip, + ) + + return _do(Int32(taddr)) + + +@cute.jit +def tcgen05_st_32x32b(num: int, taddr: int, vec): + """Store *num* × 32-bit values from an opaque vector → TMEM.""" + + @dsl_user_op + def _do(addr_val, vec_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + _nvvm.tcgen05_st( + shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, + num=num, + tmem_addr=tmem_ptr, + r=_to_ir(vec_val, loc, ip), + loc=loc, + ip=ip, + ) + + _do(Int32(taddr), vec) + + +@cute.jit +def tcgen05_cp_128x256b(taddr: int, smem_desc: Tcgen05SmemDescriptor): + """Async copy SMEM → TMEM with shape ``128x256b`` (``cta_group::1``).""" + + @dsl_user_op + def _do(addr_val, desc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + _nvvm.tcgen05_cp( + shape=_nvvm.Tcgen05CpShape.SHAPE_128x256b, + taddr=tmem_ptr, + smem_desc=_to_ir(desc_val, loc, ip), + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + loc=loc, + ip=ip, + ) + + _do(Int32(taddr), smem_desc.desc_i64[0]) + + +@cute.jit +def tcgen05_cp_128x128b(taddr: int, smem_desc: Tcgen05SmemDescriptor): + """Async copy SMEM → TMEM with shape ``128x128b`` (``cta_group::1``).""" + + @dsl_user_op + def _do(addr_val, desc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) + _nvvm.tcgen05_cp( + shape=_nvvm.Tcgen05CpShape.SHAPE_128x128b, + taddr=tmem_ptr, + smem_desc=_to_ir(desc_val, loc, ip), + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + loc=loc, + ip=ip, + ) + + _do(Int32(taddr), smem_desc.desc_i64[0]) + + +@cute.jit +def tcgen05_fence_before(): + """tcgen05.fence::before_thread_sync — non-blocking ordering fence.""" + _nvvm.tcgen05_fence(kind=_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) + + +@cute.jit +def tcgen05_fence_after(): + """tcgen05.fence::after_thread_sync — non-blocking ordering fence.""" + _nvvm.tcgen05_fence(kind=_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC) + + +@cute.jit +def umma_arrive(mbar_ptr: cute.Pointer): + """tcgen05.commit.cta_group::1.mbarrier::arrive::one — signal MMA done.""" + with elect_one(): + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + + +@cute.jit +def umma_arrive_noelect(mbar_ptr: cute.Pointer): + """tcgen05.commit.cta_group::1.mbarrier::arrive::one — signal MMA done.""" + tcgen05.commit(mbar_ptr, cta_group=tcgen05.CtaGroup.ONE) + + +# =========================================================================== +# Disable-output-lane mask constants (4 × uint32) +# =========================================================================== + +_ALL_ACTIVE = 0x00000000 +_ALL_OFF = 0xFFFFFFFF + +# SS masks (SMEM A, SMEM B) +SS_NO_MASK = (_ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE) +SS_MASK0 = (_ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) +SS_MASK1 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE) +SS_MASK2 = (_ALL_OFF, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) +SS_MASK3 = (_ALL_OFF, _ALL_OFF, _ALL_OFF, _ALL_ACTIVE) + +# TS masks (TMEM A, SMEM B) +TS_NO_MASK = (_ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE, _ALL_ACTIVE) +TS_MASK0 = (_ALL_ACTIVE, _ALL_OFF, _ALL_OFF, _ALL_OFF) +TS_MASK1 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_OFF) +TS_MASK2 = (_ALL_OFF, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) +TS_MASK3 = (_ALL_OFF, _ALL_OFF, _ALL_OFF, _ALL_ACTIVE) +TS_MASK02 = (_ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE, _ALL_OFF) +TS_MASK13 = (_ALL_OFF, _ALL_ACTIVE, _ALL_OFF, _ALL_ACTIVE) + + +# =========================================================================== +# Low-level MMA primitives +# =========================================================================== + + +@cute.jit +def tcgen05mma_ss( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + mask0: int, + mask1: int, + mask2: int, + mask3: int, +): + """Issue ``tcgen05.mma.cta_group::1.kind::tf32`` with SMEM operands.""" + + @dsl_user_op + def _do(c_val, da_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i32_ty = ir.IntegerType.get_signless(32) + i1_ty = ir.IntegerType.get_signless(1) + vec4i32_ty = ir.VectorType.get([4], i32_ty) + + c_ir = _to_ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + da_ir = _to_ir(da_val, loc, ip) + db_ir = _to_ir(db_val, loc, ip) + dv_ir = _to_ir(dv_val, loc, ip) + sc_ir = _to_ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + m0_ir = _to_ir(m0_val, loc, ip) + m1_ir = _to_ir(m1_val, loc, ip) + m2_ir = _to_ir(m2_val, loc, ip) + m3_ir = _to_ir(m3_val, loc, ip) + + undef = llvm.mlir_undef(vec4i32_ty, loc=loc, ip=ip) + idx0 = _arith.constant(i32_ty, 0, loc=loc, ip=ip) + idx1 = _arith.constant(i32_ty, 1, loc=loc, ip=ip) + idx2 = _arith.constant(i32_ty, 2, loc=loc, ip=ip) + idx3 = _arith.constant(i32_ty, 3, loc=loc, ip=ip) + v = llvm.InsertElementOp(undef, m0_ir, idx0, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m1_ir, idx1, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) + mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) + + _nvvm.tcgen05_mma( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + d=d_ptr, + a=da_ir, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + write_disable_mask=mask, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + desc_a.desc_i64[0], + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + cutlass.Int32(mask0), + cutlass.Int32(mask1), + cutlass.Int32(mask2), + cutlass.Int32(mask3), + ) + + +@cute.jit +def tcgen05mma_ts( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + mask0: int, + mask1: int, + mask2: int, + mask3: int, +): + """Issue ``tcgen05.mma.cta_group::1.kind::tf32`` with TMEM A operand.""" + + @dsl_user_op + def _do(c_val, a_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i32_ty = ir.IntegerType.get_signless(32) + i1_ty = ir.IntegerType.get_signless(1) + vec4i32_ty = ir.VectorType.get([4], i32_ty) + + c_ir = _to_ir(c_val, loc, ip) + a_ir = _to_ir(a_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) + b_ir = _to_ir(db_val, loc, ip) + dv_ir = _to_ir(dv_val, loc, ip) + sc_ir = _to_ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + m0_ir = _to_ir(m0_val, loc, ip) + m1_ir = _to_ir(m1_val, loc, ip) + m2_ir = _to_ir(m2_val, loc, ip) + m3_ir = _to_ir(m3_val, loc, ip) + + undef = llvm.mlir_undef(vec4i32_ty, loc=loc, ip=ip) + idx0 = _arith.constant(i32_ty, 0, loc=loc, ip=ip) + idx1 = _arith.constant(i32_ty, 1, loc=loc, ip=ip) + idx2 = _arith.constant(i32_ty, 2, loc=loc, ip=ip) + idx3 = _arith.constant(i32_ty, 3, loc=loc, ip=ip) + v = llvm.InsertElementOp(undef, m0_ir, idx0, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m1_ir, idx1, loc=loc, ip=ip) + v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) + mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) + + _nvvm.tcgen05_mma( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + d=d_ptr, + a=a_ptr, + b=b_ir, + idesc=dv_ir, + enable_input_d=enable_d, + write_disable_mask=mask, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + cutlass.Int32(tmem_a), + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + cutlass.Int32(mask0), + cutlass.Int32(mask1), + cutlass.Int32(mask2), + cutlass.Int32(mask3), + ) + + +# --------------------------------------------------------------------------- +# Weight-stationary variants +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ws_ss_tf32( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::tf32`` (weight-stationary, SS).""" + + @dsl_user_op + def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _to_ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + da_ir = _to_ir(da_val, loc, ip) + db_ir = _to_ir(db_val, loc, ip) + dv_ir = _to_ir(dv_val, loc, ip) + sc_ir = _to_ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + d=d_ptr, + a=da_ir, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + desc_a.desc_i64[0], + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +@cute.jit +def tcgen05mma_ws_ss_f16( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::f16`` (weight-stationary, SS).""" + + @dsl_user_op + def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _to_ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + da_ir = _to_ir(da_val, loc, ip) + db_ir = _to_ir(db_val, loc, ip) + dv_ir = _to_ir(dv_val, loc, ip) + sc_ir = _to_ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.F16, + d=d_ptr, + a=da_ir, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + desc_a.desc_i64[0], + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +@cute.jit +def tcgen05mma_ws_ts_tf32( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::tf32`` with TMEM A (weight-stationary).""" + + @dsl_user_op + def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _to_ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + a_ir = _to_ir(a_val, loc, ip) + a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) + db_ir = _to_ir(db_val, loc, ip) + dv_ir = _to_ir(dv_val, loc, ip) + sc_ir = _to_ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.TF32, + d=d_ptr, + a=a_ptr, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + cutlass.Int32(tmem_a), + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +@cute.jit +def tcgen05mma_ws_ts_f16( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, + collector_b_buffer=None, + collector_op=None, +): + """Issue ``tcgen05.mma.ws.cta_group::1.kind::f16`` with TMEM A (weight-stationary).""" + + @dsl_user_op + def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): + ptr6_ty = llvm.PointerType.get(address_space=6) + i1_ty = ir.IntegerType.get_signless(1) + + c_ir = _to_ir(c_val, loc, ip) + d_ptr = llvm.inttoptr(ptr6_ty, c_ir, loc=loc, ip=ip) + a_ir = _to_ir(a_val, loc, ip) + a_ptr = llvm.inttoptr(ptr6_ty, a_ir, loc=loc, ip=ip) + db_ir = _to_ir(db_val, loc, ip) + dv_ir = _to_ir(dv_val, loc, ip) + sc_ir = _to_ir(sc_val, loc, ip) + enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) + + _nvvm.tcgen05_mma_ws( + mma_kind=_nvvm.Tcgen05MMAKind.F16, + d=d_ptr, + a=a_ptr, + b=db_ir, + idesc=dv_ir, + enable_input_d=enable_d, + collector_b_buffer=collector_b_buffer, + collector_op=collector_op, + loc=loc, + ip=ip, + ) + + _do( + cutlass.Int32(tmem_c), + cutlass.Int32(tmem_a), + desc_b.desc_i64[0], + cutlass.Int32(desc_val), + cutlass.Int32(scale_out), + ) + + +# =========================================================================== +# Named convenience wrappers (pre-set mask constants) +# =========================================================================== + +# --------------------------------------------------------------------------- +# SS named wrappers (SMEM A) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ss_no_mask( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA with no output-lane disable (all rows active).""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_NO_MASK[0], SS_NO_MASK[1], SS_NO_MASK[2], SS_NO_MASK[3]) + + +@cute.jit +def tcgen05mma_ss_mask0( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0, 0xF…, 0, 0xF…} — groups 0,2 active (1,3 disabled).""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK0[0], SS_MASK0[1], SS_MASK0[2], SS_MASK0[3]) + + +@cute.jit +def tcgen05mma_ss_mask1( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0xF…, 0, 0xF…, 0} — groups 1,3 active (0,2 disabled).""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK1[0], SS_MASK1[1], SS_MASK1[2], SS_MASK1[3]) + + +@cute.jit +def tcgen05mma_ss_mask2( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0xF…, 0xF…, 0, 0xF…} — group 2 only active.""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK2[0], SS_MASK2[1], SS_MASK2[2], SS_MASK2[3]) + + +@cute.jit +def tcgen05mma_ss_mask3( + desc_a: Tcgen05SmemDescriptor, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """SS MMA: mask={0xF…, 0xF…, 0xF…, 0} — group 3 only active.""" + tcgen05mma_ss(desc_a, desc_b, tmem_c, desc_val, scale_out, SS_MASK3[0], SS_MASK3[1], SS_MASK3[2], SS_MASK3[3]) + + +# --------------------------------------------------------------------------- +# TS named wrappers (TMEM A) +# --------------------------------------------------------------------------- + + +@cute.jit +def tcgen05mma_ts_no_mask( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA with no output-lane disable (all rows active).""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_NO_MASK[0], TS_NO_MASK[1], TS_NO_MASK[2], TS_NO_MASK[3]) + + +@cute.jit +def tcgen05mma_ts_mask0( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0, 0xF…, 0xF…, 0xF…} — group 0 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK0[0], TS_MASK0[1], TS_MASK0[2], TS_MASK0[3]) + + +@cute.jit +def tcgen05mma_ts_mask1( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0, 0xF…, 0xF…} — group 1 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK1[0], TS_MASK1[1], TS_MASK1[2], TS_MASK1[3]) + + +@cute.jit +def tcgen05mma_ts_mask2( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0xF…, 0, 0xF…} — group 2 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK2[0], TS_MASK2[1], TS_MASK2[2], TS_MASK2[3]) + + +@cute.jit +def tcgen05mma_ts_mask3( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0xF…, 0xF…, 0} — group 3 only active.""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK3[0], TS_MASK3[1], TS_MASK3[2], TS_MASK3[3]) + + +@cute.jit +def tcgen05mma_ts_mask02( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0, 0xF…, 0, 0xF…} — groups 0,2 active (1,3 disabled).""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK02[0], TS_MASK02[1], TS_MASK02[2], TS_MASK02[3]) + + +@cute.jit +def tcgen05mma_ts_mask13( + tmem_a: int, + desc_b: Tcgen05SmemDescriptor, + tmem_c: int, + desc_val: int, + scale_out: int, +): + """TS MMA: mask={0xF…, 0, 0xF…, 0} — groups 1,3 active (0,2 disabled).""" + tcgen05mma_ts(tmem_a, desc_b, tmem_c, desc_val, scale_out, TS_MASK13[0], TS_MASK13[1], TS_MASK13[2], TS_MASK13[3]) diff --git a/cula/utils.py b/cula/utils.py index 8b8e0ab1..99d19cd4 100644 --- a/cula/utils.py +++ b/cula/utils.py @@ -83,8 +83,8 @@ def assert_hopper(device: torch.device | str | int | None = None) -> None: def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callable: """Return the appropriate ``kda_prefill`` implementation for *device*. - - sm100/sm103 (Blackwell) → cula.kda.blackwell_fused_fwd.flash_kda_prefill - - sm90 (Hopper) → cula.kda.kda_prefill_hopper + - sm100/sm103 (Blackwell) → cula.ops.kda.experimental.sm100_fused (WIP) + - sm90 (Hopper) → cula.kda.hopper_fused_fwd Args: device: CUDA device to query. Defaults to the currently active device. @@ -94,13 +94,13 @@ def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callabl """ major, minor = get_device_sm_version(device) if major == 10 and minor in (0, 3): - from cula.kda import kda_prefill_blackwell + from cula.ops.kda.experimental.sm100_fused.wrapper import flash_kda_prefill - return kda_prefill_blackwell + return flash_kda_prefill elif major == 9 and minor == 0: - from cula.kda import kda_prefill_hopper + from cula.kda.hopper_fused_fwd import cula_kda_prefill - return kda_prefill_hopper + return cula_kda_prefill else: raise RuntimeError( f"Unsupported CUDA compute capability sm_{major}{minor}. " @@ -109,29 +109,16 @@ def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callabl def get_pre_scan(device: torch.device | str | int | None = None) -> Callable: - """Return the appropriate ``chunk_delta_rule_pre_scan`` implementation for *device*. - - - sm100/sm103 (Blackwell) → cula.ops.cp.pre_scan (CuTeDSL SM100 kernel) - - sm90 (Hopper) → cula.ops.cp.pre_scan_sm90 (to be implemented) - - Args: - device: CUDA device to query. Defaults to the currently active device. - - Raises: - RuntimeError: If the device architecture is not supported. - """ + """Return the intracard-CP pre_scan implementation for *device*.""" major, minor = get_device_sm_version(device) if major == 10 and minor in (0, 3): - from cula.ops.cp.pre_scan import chunk_delta_rule_pre_scan + from cula.ops.kda.sm100.cp.pre_scan import chunk_delta_rule_pre_scan return chunk_delta_rule_pre_scan - elif major == 9 and minor == 0: - raise NotImplementedError("The Hopper (SM90) implementation of pre_scan is not yet available.") - else: - raise RuntimeError( - f"Unsupported CUDA compute capability sm_{major}{minor}. " - f"Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." - ) + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Intracard CP pre_scan is currently available only on SM100/SM103." + ) @cute.jit diff --git a/docs/chunk_delta_h_pipeline.md b/docs/chunk_delta_h_pipeline.md index 7cd06451..0d97baf6 100644 --- a/docs/chunk_delta_h_pipeline.md +++ b/docs/chunk_delta_h_pipeline.md @@ -1,6 +1,6 @@ # ChunkDeltaRuleFwdH Pipeline -> 文件: `cula/ops/chunk_delta_h.py` +> 文件: `cula/ops/kda/sm100/delta_h.py` > 类名: `ChunkDeltaRuleFwdH` ## 计算公式 @@ -24,7 +24,7 @@ $$h_{new} = 2^{gk} \cdot h + \text{update}$$ | 6 | Store | TMA S2G 写 h_out 和 v_new | | 7 | Empty | 占位 | -**总线程**: 256 (8 warps) +**总线程**: 256 (8 warps) **寄存器分配**: CUDA=232, Others=40, occ=1 (仅支持) ## MMA 操作 diff --git a/pyproject.toml b/pyproject.toml index fe93e562..f628d3a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,8 +80,8 @@ force-sort-within-sections = false "__init__.py" = ["F401"] "cula/_version.py" = ["UP007"] # TODO: fix undefined names (exp_g, chunk_kda_bwd_dqkwg) — WIP code -"cula/ops/kda_fully_fused_sm100_wip.py" = ["F821"] -"cula/kda/blackwell_fused_fwd.py" = ["F821"] +"cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py" = ["F821"] +"cula/ops/kda/experimental/sm100_fused/wrapper.py" = ["F821"] [tool.setuptools_scm] write_to = "cula/_version.py" diff --git a/tests/test_chunk_delta_h.py b/tests/test_chunk_delta_h.py index 01cdb157..75223a15 100644 --- a/tests/test_chunk_delta_h.py +++ b/tests/test_chunk_delta_h.py @@ -21,8 +21,10 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import importlib.util +pytestmark = pytest.mark.sm100_only + _spec = importlib.util.spec_from_file_location( - "chunk_delta_h_sm100", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "chunk_delta_h_sm100.py") + "chunk_delta_h", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "kda", "sm100", "delta_h.py") ) _mod = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_mod) diff --git a/tests/test_compare_with_fla.py b/tests/test_compare_with_fla.py index 3468c9c4..f74ddea2 100644 --- a/tests/test_compare_with_fla.py +++ b/tests/test_compare_with_fla.py @@ -33,7 +33,7 @@ from cutlass.cute.runtime import from_dlpack # noqa: E402 # Our implementation -from cula.ops.chunk_delta_h_sm100 import ChunkDeltaRuleFwdH # noqa: E402 +from cula.ops.kda.sm100.delta_h import ChunkDeltaRuleFwdH # noqa: E402 def fla_reference_chunk_fwd_h( diff --git a/tests/test_fwd_o.py b/tests/test_fwd_o.py index 51b30003..9ef7b213 100644 --- a/tests/test_fwd_o.py +++ b/tests/test_fwd_o.py @@ -37,7 +37,7 @@ import importlib.util _spec = importlib.util.spec_from_file_location( - "fwd_o_sm100", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "fwd_o_sm100.py") + "fwd_o", os.path.join(os.path.dirname(__file__), "..", "cula", "ops", "kda", "sm100", "fwd_o.py") ) _mod = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_mod) @@ -49,6 +49,8 @@ from fla.ops.gla.chunk import chunk_gla_fwd_o_gk as triton_chunk_gla_fwd_o_gk # noqa: E402 +pytestmark = pytest.mark.sm100_only + # ── Constants ── K, V, BT = 128, 128, 64 DTYPE = torch.bfloat16 diff --git a/tests/test_intracard_cp.py b/tests/test_intracard_cp.py index d478190a..85d077a6 100644 --- a/tests/test_intracard_cp.py +++ b/tests/test_intracard_cp.py @@ -32,13 +32,13 @@ from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h as fla_fwd_h # noqa: E402 from fla.utils import assert_close # noqa: E402 (RMSE-relative + atol short-circuit + NaN check) -from cula.ops.chunk_delta_h_sm100 import chunk_gated_delta_rule_fwd_h # noqa: E402 -from cula.ops.cp.chunk_delta_h import ( # noqa: E402 +from cula.ops.kda.sm100.cp.chunk_delta_h import ( # noqa: E402 compute_subseq_len, intracard_fwd_h, prepare_subseq_cu_seqlens, should_use_intracard_cp, ) +from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h # noqa: E402 from cula.utils import get_device_sm_count # noqa: E402 # Constants & tolerances — aligned with existing cuLA tests (see below). @@ -128,19 +128,40 @@ def run_cula_cp(k, w, u, gk, h0, cu, **kw): def run_intracard_direct(k, w, u, gk, h0, cu, *, output_final_state=True, save_new_value=True): - """Direct CP call — skips the auto-dispatch heuristic.""" - return intracard_fwd_h( - k=k, - w=w, - u=u, - gk=gk, - initial_state=h0, - output_final_state=output_final_state, - chunk_size=BT, - save_new_value=save_new_value, - cu_seqlens=cu, - cu_seqlens_cpu=cu.cpu(), - ) + """Direct CP call — skips the auto-dispatch heuristic. + + intracard_fwd_h is a pure executor that raises NotSplittableError when the + post-split occupancy guard rejects; mirror the production caller's graceful + fallback to the serial path so configs that don't engage CP still return. + """ + from cula.ops.kda.policy import NotSplittableError + + try: + return intracard_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=h0, + output_final_state=output_final_state, + chunk_size=BT, + save_new_value=save_new_value, + cu_seqlens=cu, + cu_seqlens_cpu=cu.cpu(), + ) + except NotSplittableError: + return chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=h0, + output_final_state=output_final_state, + chunk_size=BT, + save_new_value=save_new_value, + cu_seqlens=cu, + _no_cp=True, + ) def run_fla(k, w, u, gk, h0, cu, **kw): @@ -234,6 +255,19 @@ def assert_cp_splits(cu, H, total_T): assert split_info, "config must exercise the split path" +def test_forced_cp_not_splittable_raises(): + """use_intracard_cp=True on an unsplittable shape must raise NotSplittableError.""" + from cula.ops.kda.policy import NotSplittableError + + # A single one-chunk sequence cannot be meaningfully split. + cu = torch.tensor([0, BT], dtype=torch.int32, device=DEVICE) + k = torch.randn(1, BT, 1, K, device=DEVICE, dtype=torch.bfloat16) + w = torch.randn(1, BT, 1, K, device=DEVICE, dtype=torch.bfloat16) + u = torch.randn(1, BT, 1, V, device=DEVICE, dtype=torch.bfloat16) + with torch.inference_mode(), pytest.raises(NotSplittableError): + chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, cu_seqlens=cu, use_intracard_cp=True) + + # ====================== Dispatch path: CP vs no-CP ====================== # Verifies chunk_gated_delta_rule_fwd_h routes to CP under env+inference_mode, # and matches the same-kernel no-CP baseline. diff --git a/tests/test_la_decode.py b/tests/test_la_decode.py index 5b57ac58..b8708336 100644 --- a/tests/test_la_decode.py +++ b/tests/test_la_decode.py @@ -30,7 +30,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.ops.la_decode import linear_attention_decode +from cula.ops.lightning.decode import linear_attention_decode try: from fla.ops.common.fused_recurrent import fused_recurrent_fwd @@ -231,9 +231,10 @@ def test_vs_fla(B): # End-to-End Prefill -> Decode Test # --------------------------------------------------------------------------- +@pytest.mark.sm100_only def test_prefill_decode_e2e(): """Verify prefill output state passes directly into decode without transpose.""" - from cula.ops.lightning_attn_sm100 import lightning_attn_fwd + from cula.ops.lightning.prefill_sm100 import lightning_attn_fwd B, S, H, D = 2, 64, 8, 128 scale = D**-0.5 diff --git a/tests/test_la_decode_pool.py b/tests/test_la_decode_pool.py index c9a75796..d3205328 100644 --- a/tests/test_la_decode_pool.py +++ b/tests/test_la_decode_pool.py @@ -28,7 +28,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.ops.la_decode import linear_attention_decode +from cula.ops.lightning.decode import linear_attention_decode def torch_la_decode_ref(q, k, v, state, decay_scales, scale): diff --git a/tests/test_lightning_attn.py b/tests/test_lightning_attn.py index 5e52f86c..46c59b99 100644 --- a/tests/test_lightning_attn.py +++ b/tests/test_lightning_attn.py @@ -35,7 +35,7 @@ warnings.filterwarnings("ignore", category=DeprecationWarning) sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen # noqa: E402 +from cula.ops.lightning.prefill_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen # noqa: E402 try: from fla.ops.simple_gla import chunk_simple_gla diff --git a/tests/test_ptx_umma_masked.py b/tests/test_ptx_umma_masked.py index e7bc3196..3d6859b3 100644 --- a/tests/test_ptx_umma_masked.py +++ b/tests/test_ptx_umma_masked.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """ -Standalone CuteDSL test for ptx_umma_masked.py inline PTX MMA wrappers. +Standalone CuteDSL test for SM100 masked MMA inline PTX wrappers. Tests: 1. tcgen05mma_ss_no_mask -- M=64, N=64, K=8, TF32, all rows active → matches torch.mm @@ -16,7 +16,7 @@ All descriptor values are computed via make_umma_smem_desc / smem_descriptor_to_int (proven correct in test_umma_ptx_jit.py). Wrapped in Tcgen05SmemDescriptor for API -compatibility with ptx_umma_masked.py convenience wrappers. +compatibility with cula.ops.sm100.ptx convenience wrappers. """ import pathlib @@ -30,6 +30,7 @@ import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils +import pytest import torch from cutlass.cute.arch import ( elect_one, @@ -48,13 +49,15 @@ from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Float32, Int32, Int64, TFloat32 -from cula.ops.ptx_umma_ext import ( +from cula.ops.sm100.ptx import ( Tcgen05SmemDescriptor, tcgen05mma_ss_mask0, tcgen05mma_ss_mask1, tcgen05mma_ss_no_mask, ) +pytestmark = pytest.mark.sm100_only + M_DIM, N_DIM, K_DIM = 64, 64, 8 TMEM_COLS = 64 diff --git a/tests/test_ptx_umma_ws.py b/tests/test_ptx_umma_ws.py index da211a9a..210028b4 100644 --- a/tests/test_ptx_umma_ws.py +++ b/tests/test_ptx_umma_ws.py @@ -16,7 +16,7 @@ - tmem region 0: accumulator for both phases - tmem region 1: holds A data for TS phase (populated via R2T store) -SMEM layout follows the same conventions as test_ptx_umma_masked.py. +SMEM layout follows the same conventions as test_ptx_umma_masked. """ import pathlib @@ -30,6 +30,7 @@ import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils +import pytest import torch from cutlass.cute.arch import ( elect_one, @@ -46,19 +47,18 @@ from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import BFloat16, Float32, Int32, Int64, TFloat32 -from cula.ops.intrinsics_sm100 import ( - store_256b, - subvec, - tcgen05_ld_32x32b, -) -from cula.ops.ptx_umma_ext import ( +from cula.ops.ptx import store_256b, subvec +from cula.ops.sm100.ptx import ( CollectorBBuffer, CollectorOp, Tcgen05SmemDescriptor, + tcgen05_ld_32x32b, tcgen05mma_ws_ss_f16, tcgen05mma_ws_ss_tf32, ) +pytestmark = pytest.mark.sm100_only + M_DIM, N_DIM = 64, 64 # TODO: support arbitrary K K_DIM_TF32 = 8 # kind::tf32 → K>=8, tile size From ca6f1745fc9a8bbcaec273ec9bdac9fea0f697c6 Mon Sep 17 00:00:00 2001 From: Haisha Zhao <33570593+Hyaloid@users.noreply.github.com> Date: Wed, 8 Jul 2026 00:00:06 +0800 Subject: [PATCH 28/34] feat: intracard cp for sm90 (#86) * feat: intracard cp for sm90 * tune threshold & compare with fla-intracard-cp --------- Co-authored-by: Chaofan Yu <103550325+icavan@users.noreply.github.com> --- benchmarks/bench_intracard_cp.py | 2 +- benchmarks/bench_intracard_cp_sm90.py | 448 +++++++++++++++++ csrc/api/kda_sm90.cu | 49 +- csrc/api/pybind.cu | 99 ++++ csrc/kda/sm90/collective/mainloop_kda_fwd.hpp | 28 +- csrc/kda/sm90/kda_fwd_sm90.cu | 25 +- csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu | 122 ++--- csrc/kda/sm90/kernel/kernel_kda_fwd.hpp | 5 + csrc/kda/sm90/prefill_kernel.hpp | 5 +- csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh | 8 +- cula/kda/__init__.py | 2 + cula/kda/auto_route.py | 112 +++++ cula/kda/cp_context.py | 452 ++++++++++++++++++ cula/kda/cp_h_boundary.py | 190 ++++++++ cula/kda/gate_l2norm_fused.py | 194 ++++++++ cula/kda/hopper_fused_fwd_opt.py | 390 +++++++++++++++ cula/kda/l2norm_qk_fused.py | 119 +++++ cula/kda/wy_intra.py | 354 ++++++++++++++ cula/kda/wy_recompute.py | 137 ++++++ tests/test_intracard_cp_sm90.py | 439 +++++++++++++++++ 20 files changed, 3075 insertions(+), 105 deletions(-) create mode 100644 benchmarks/bench_intracard_cp_sm90.py create mode 100644 csrc/api/pybind.cu create mode 100644 cula/kda/auto_route.py create mode 100644 cula/kda/cp_context.py create mode 100644 cula/kda/cp_h_boundary.py create mode 100644 cula/kda/gate_l2norm_fused.py create mode 100644 cula/kda/hopper_fused_fwd_opt.py create mode 100644 cula/kda/l2norm_qk_fused.py create mode 100644 cula/kda/wy_intra.py create mode 100644 cula/kda/wy_recompute.py create mode 100644 tests/test_intracard_cp_sm90.py diff --git a/benchmarks/bench_intracard_cp.py b/benchmarks/bench_intracard_cp.py index 4eba5a48..41e9de3a 100644 --- a/benchmarks/bench_intracard_cp.py +++ b/benchmarks/bench_intracard_cp.py @@ -59,7 +59,7 @@ # ============================================================ BT, D = 64, 128 H_VALUES = [4, 8] -WARMUP = 10 +WARMUP = 25 N_ITERS = 100 NCU_MODE = False SANITIZER_MODE = False diff --git a/benchmarks/bench_intracard_cp_sm90.py b/benchmarks/bench_intracard_cp_sm90.py new file mode 100644 index 00000000..c8f85773 --- /dev/null +++ b/benchmarks/bench_intracard_cp_sm90.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""bench_intracard_cp_sm90.py — CP-on vs CP-off (vs FLA baseline) for SM90 KDA prefill. + +Mirrors benchmarks/bench_intracard_cp.py (SM100 version) but for the Hopper +(SM90) path: + + CP_on : cula.kda.kda_prefill_hopper_auto + CP_off : cula.kda.kda_prefill_hopper + FLA : fla.ops.kda.chunk_kda (Triton baseline) + +Reports per-config `pred` (would CP fire?) and `n_sub` (CP-chunk count). When +`pred=N` we still measure CP_on to confirm the bypass adds no regression. The +`CP_on/FLA` column shows the speedup of cuLA's optimized (CP-on) kernel over +the FLA Triton baseline. + +Usage: + python benchmarks/bench_intracard_cp_sm90.py [--ncu] [--sanitizer] +""" + +import argparse +import os +import pathlib +import sys + +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) +os.environ.setdefault("FLA_INTRACARD_CP", "1") + +from fla.ops.common.intracard_cp import compute_subseq_len, prepare_subseq_cu_seqlens +from fla.ops.kda import chunk_kda as fla_chunk_kda + +from benchmarks.utils import ( + SEED, + exclusive_cumsum, + prepare_safe_gate_inputs, + set_seed, + time_cuda_fn, +) +from cula.kda import kda_prefill_hopper, kda_prefill_hopper_auto +from cula.kda.auto_route import _should_use_opt +from cula.kda.cp_context import _calc_cp_seqs, is_dominant_long_seq +from cula.kda.hopper_fused_fwd_opt import FUSED_GATE_L2NORM_VARLEN_AVG_SEQ, _fused_gate_l2norm_threshold +from cula.utils import get_device_sm_count + +# ============================================================ +# Constants +# ============================================================ +BT, D = 64, 128 +H_VALUES = [4, 8] +WARMUP = 10 +N_ITERS = 10 +NCU_MODE = False +SANITIZER_MODE = False + +# (tag, seq_lens) — varlen configs, run with cu_seqlens=cumsum(seq_lens) +CONFIGS = [ + # small varlen — exercises fused gate+l2norm path (packed_T*H <= 65536) + ("4x256", [256] * 4), + ("8x256", [256] * 8), + ("16x256", [256] * 16), + ("4x1K", [1024] * 4), + ("8x1K", [1024] * 8), + ("4x2K", [2048] * 4), + ("1K+512+256+128", [1024, 512, 256, 128]), + ("2K+1K+512+256", [2048, 1024, 512, 256]), + ("1K+1+63+65+129", [1024, 1, 63, 65, 129]), + # single seq + ("T=4K", [4096]), + ("T=8K", [8192]), + ("T=32K", [32768]), + ("T=64K", [65536]), + ("T=128K", [131072]), + # equal-length batches (~32K total) + ("8x4K", [4096] * 8), + ("4x8K", [8192] * 4), + ("2x16K", [16384] * 2), + # asymmetric multi-seq + ("16K+16K", [16384, 16384]), + ("24K+8K", [24576, 8192]), + ("28K+4K", [28672, 4096]), + ("32K+256+256", [32768, 256, 256]), + ("40K+1K+8K", [40960, 1024, 8192]), + ("64K+512+256+128", [65536, 512, 256, 128]), + ("128K+1K", [131072, 1024]), + ("128K+2x1K", [131072, 1024, 1024]), + ("128K+5x1K", [131072] + [1024] * 5), + ("128K+10x1K", [131072] + [1024] * 10), +] + + +# ============================================================ +# Helpers +# ============================================================ +def _bench_warmup_iters(): + warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + return warmup, n_iters + + +def run_call(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, enable_cp, return_state=False): + fn = kda_prefill_hopper_auto if enable_cp else kda_prefill_hopper + out = fn( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + initial_state=None, + output_final_state=return_state, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return out + + +def run_fla_call(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, return_state=False): + # FLA's chunk_kda fuses A_log + dt_bias internally when use_gate_in_kernel=True; pass them via kwargs. + # Wrap in inference_mode so FLA's IntraCardCPBackend gate (intracard.py) activates. + with torch.inference_mode(): + return fla_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=None, + output_final_state=return_state, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=False, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + A_log=A_log, + dt_bias=dt_bias, + ) + + +def accuracy(ref, got): + if ref is None or got is None: + return float("nan"), float("nan") + diff = (ref.float() - got.float()).abs() + return diff.max().item(), diff.mean().item() + + +def predict_cp(seq_lens, H, num_sms, device): + cu = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + raw_batch = len(seq_lens) + packed_seq = sum(seq_lens) + + if raw_batch > 1: + cp_wf = (raw_batch * H <= 16 and packed_seq >= 8192) or ( + packed_seq >= 8192 and H <= 16 and is_dominant_long_seq(seq_lens, H) + ) + else: + cp_wf = (H <= 8 and packed_seq >= 4096) or (H <= 16 and packed_seq >= 4096) or (H <= 32 and packed_seq >= 16384) + if not cp_wf: + return False, 0 + + use_cp, cp_cu, *_ = _calc_cp_seqs(cu, BT, H, num_sms, raw_cu_seqlens_cpu=cu.cpu()) + if not use_cp: + return False, 0 + n_sub = int(cp_cu.numel() - 1) + if n_sub == raw_batch: # no-op split + return False, 0 + return True, n_sub + + +def predict_fla_cp(seq_lens, H, num_sms): + """Mirror fla.ops.common.intracard_cp.intracard_fwd_h gating to predict + whether FLA's intracard CP fires and how many sub-sequences result. + HV (num_v_heads) maps to H here.""" + cu = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int64) + seq_lens_t = torch.diff(cu) + max_seq_len = int(seq_lens_t.max().item()) + subseq_len = compute_subseq_len(max_seq_len, num_sms, H, BT) + if (seq_lens_t < 2 * subseq_len).all(): + return False, 0 + _, split_info, total_subseqs = prepare_subseq_cu_seqlens(cu, subseq_len, BT) + if not split_info: + return False, 0 + return True, total_subseqs + + +def predict_fused_all_pre(q, v, cu_seqlens_for_opt, *, cu_seqlens_is_none, use_gate_in_kernel, use_qk_l2norm_in_kernel): + if not _should_use_opt(q, cu_seqlens_for_opt): + return False + num_qk_heads = q.shape[-2] + num_v_heads = v.shape[-2] + if cu_seqlens_is_none: + avg_seq_ok = True + else: + N = cu_seqlens_for_opt.numel() - 1 + packed_T = q.shape[1] + avg_seq_ok = N <= 1 or packed_T <= N * FUSED_GATE_L2NORM_VARLEN_AVG_SEQ + return ( + use_gate_in_kernel + and use_qk_l2norm_in_kernel + and (q.numel() // q.shape[-1]) <= _fused_gate_l2norm_threshold(cu_seqlens_is_none) + and num_qk_heads == num_v_heads + and avg_seq_ok + ) + + +# ============================================================ +# Benchmark +# ============================================================ +SEP = " " + "─" * 180 +ROW_HEADER = ( + f" {'config':<24s} {'T':>7s} {'pred':>4s} {'sub':>4s} {'fla_cp':>6s} {'fla_sub':>7s} {'fused_pre':>5s}" + f" │ {'o max/mean':>17s} {'ht max/mean':>17s}" + f" │ {'FLA(ms)':>9s} {'CP_off(ms)':>10s} {'CP_on(ms)':>10s} {'CP_on/off':>8s} {'CP_on/FLA':>9s}" +) + + +def _format_row(r): + pred_s = "Y" if r["pred"] else "N" + fla_pred_s = "Y" if r["fla_pred"] else "N" + fused_s = "Y" if r["fused_all_pre"] else "N" + return ( + f" {r['tag']:<24s} {r['total_T']:>7d} {pred_s} {r['n_sub']:>4d} {fla_pred_s} {r['fla_n_sub']:>4d} {fused_s}" + f" │ {r['o_max']:>7.1e}/{r['o_mean']:>7.1e} {r['ht_max']:>7.1e}/{r['ht_mean']:>7.1e}" + f" │ {r['ms_fla']:>9.4f} {r['ms_off']:>10.4f} {r['ms_on']:>10.4f}" + f" {r['speedup']:>7.2f}x {r['speedup_vs_fla']:>8.2f}x" + ) + + +def bench_cp(h_values, configs): + print("\n" + "=" * 110) + print(" BENCHMARK REPORT: Intracard CP (SM90)") + print(" CP-on (kda_prefill_hopper_auto) vs CP-off (kda_prefill_hopper) vs FLA (chunk_kda)") + print(f" D={D} dtype=bf16 safe_gate=True") + wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") + print(f" Warmup={wu} Iters={ni}{mode_tag}") + print("=" * 110) + + device = torch.device("cuda") + num_sms = get_device_sm_count(device) + results = [] + + for H in h_values: + print(f"\n [H={H}]", flush=True) + print(SEP, flush=True) + print(ROW_HEADER, flush=True) + print(SEP, flush=True) + + for tag, seq_lens in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + total_T = sum(seq_lens) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens, seed=SEED) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, lower_bound = inputs["scale"], inputs["lower_bound"] + + pred, n_sub = predict_cp(seq_lens, H, num_sms, device) + fla_pred, fla_n_sub = predict_fla_cp(seq_lens, H, num_sms) + fused_all_pre = predict_fused_all_pre( + q, + v, + cu_seqlens, + cu_seqlens_is_none=False, + use_gate_in_kernel=True, + use_qk_l2norm_in_kernel=True, + ) + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + + try: + o_off, ht_off = run_call(**common, enable_cp=False, return_state=True) + o_on, ht_on = run_call(**common, enable_cp=True, return_state=True) + o_max, o_mean = accuracy(o_off, o_on) + ht_max, ht_mean = accuracy(ht_off, ht_on) + del o_off, ht_off, o_on, ht_on + + ms_off = time_cuda_fn(lambda: run_call(**common, enable_cp=False), *_bench_warmup_iters()) + ms_on = time_cuda_fn(lambda: run_call(**common, enable_cp=True), *_bench_warmup_iters()) + speedup = ms_off / ms_on if ms_on > 0 else float("inf") + try: + ms_fla = time_cuda_fn(lambda: run_fla_call(**common), *_bench_warmup_iters()) + speedup_vs_fla = ms_fla / ms_on if ms_on > 0 else float("inf") + except Exception: + ms_fla = float("nan") + speedup_vs_fla = float("nan") + except torch.cuda.OutOfMemoryError: + ms_off = ms_on = speedup = float("nan") + ms_fla = speedup_vs_fla = float("nan") + o_max = o_mean = ht_max = ht_mean = float("nan") + + row = { + "tag": tag, + "H": H, + "total_T": total_T, + "pred": pred, + "n_sub": n_sub, + "fla_pred": fla_pred, + "fla_n_sub": fla_n_sub, + "fused_all_pre": fused_all_pre, + "ms_off": ms_off, + "ms_on": ms_on, + "ms_fla": ms_fla, + "speedup": speedup, + "speedup_vs_fla": speedup_vs_fla, + "o_max": o_max, + "o_mean": o_mean, + "ht_max": ht_max, + "ht_mean": ht_mean, + } + results.append(row) + print(_format_row(row), flush=True) + + del q, k, v, g, beta, A_log, dt_bias, inputs + torch.cuda.empty_cache() + + print(SEP, flush=True) + + return results + + +# ============================================================ +# Report (summary only — per-row output is streamed inside bench_cp) +# ============================================================ +def print_report(results, h_values): + sep = "=" * 110 + triggered = [r for r in results if r["pred"]] + bypassed = [r for r in results if not r["pred"]] + + print() + print(sep) + print(" Summary") + print(sep) + + if triggered: + speedups = [r["speedup"] for r in triggered if r["speedup"] == r["speedup"]] # NaN filter + if speedups: + geo = 1.0 + for s in speedups: + geo *= s + geo = geo ** (1 / len(speedups)) + print( + f" CP triggered ({len(triggered)} configs): " + f"geo-mean={geo:.2f}x best={max(speedups):.2f}x worst={min(speedups):.2f}x" + ) + + if bypassed: + ratios = [r["ms_on"] / r["ms_off"] for r in bypassed if r["ms_off"] == r["ms_off"] and r["ms_off"] > 0] + if ratios: + print( + f" CP bypassed ({len(bypassed)} configs): " + f"mean overhead={sum(ratios) / len(ratios):.3f}x max={max(ratios):.3f}x " + f"(1.00 = no regression)" + ) + + # cuLA (CP-on) vs FLA speedups + fla_speedups = [r["speedup_vs_fla"] for r in results if r["speedup_vs_fla"] == r["speedup_vs_fla"]] + if fla_speedups: + geo = 1.0 + for s in fla_speedups: + geo *= s + geo = geo ** (1 / len(fla_speedups)) + print( + f" cuLA (CP-on) vs FLA ({len(fla_speedups)} configs): " + f"geo-mean={geo:.2f}x best={max(fla_speedups):.2f}x worst={min(fla_speedups):.2f}x" + ) + tri_fla = [r["speedup_vs_fla"] for r in triggered if r["speedup_vs_fla"] == r["speedup_vs_fla"]] + if tri_fla: + geo_t = 1.0 + for s in tri_fla: + geo_t *= s + geo_t = geo_t ** (1 / len(tri_fla)) + print( + f" └─ CP-triggered subset ({len(tri_fla)} configs): " + f"geo-mean={geo_t:.2f}x best={max(tri_fla):.2f}x worst={min(tri_fla):.2f}x" + ) + + o_maxes = [r["o_max"] for r in results if r["o_max"] == r["o_max"]] + ht_maxes = [r["ht_max"] for r in results if r["ht_max"] == r["ht_max"]] + if o_maxes: + print( + f" Accuracy (CP-on vs CP-off): " + f"o max={max(o_maxes):.2e} avg={sum(o_maxes) / len(o_maxes):.2e} " + f"ht max={max(ht_maxes):.2e} avg={sum(ht_maxes) / len(ht_maxes):.2e}" + ) + + print(sep) + + +# ============================================================ +# Main +# ============================================================ +def main(): + parser = argparse.ArgumentParser(description="bench_intracard_cp_sm90: CP-on vs CP-off") + parser.add_argument("--ncu", action="store_true", help="NCU profiling mode: warmup=1, iters=1") + parser.add_argument("--sanitizer", action="store_true", help="Sanitizer mode: warmup=1, iters=1") + args = parser.parse_args() + + global NCU_MODE, SANITIZER_MODE + if args.ncu: + NCU_MODE = True + print("[NCU mode] warmup=1, iters=1") + if args.sanitizer: + SANITIZER_MODE = True + print("[Sanitizer mode] warmup=1, iters=1") + + results = bench_cp(H_VALUES, CONFIGS) + print_report(results, H_VALUES) + return results + + +if __name__ == "__main__": + main() diff --git a/csrc/api/kda_sm90.cu b/csrc/api/kda_sm90.cu index d80df7cc..d9a4de06 100644 --- a/csrc/api/kda_sm90.cu +++ b/csrc/api/kda_sm90.cu @@ -35,7 +35,9 @@ kda_fwd_prefill( torch::Tensor workspace_buffer, float scale, bool output_final_state, - bool safe_gate) { + bool safe_gate, + OptionalTensor cp_seq_map_, + OptionalTensor raw_cu_seqlens_) { // Q, K: [packed_seq, num_qk_heads, D] // V/O/g: [packed_seq, num_v_heads, D] (GVA: num_v_heads is a positive integer multiple of num_qk_heads) auto packed_seq = q.size(0); @@ -44,6 +46,31 @@ kda_fwd_prefill( auto head_size = q.size(2); auto num_seqs = cu_seqlens.size(0) - 1; + // Intra-card CP plumbing. + int32_t const* cp_seq_map_ptr = nullptr; + int32_t const* raw_cu_seqlens_ptr = nullptr; + int32_t raw_num_seqs = static_cast(num_seqs); + if (cp_seq_map_.has_value()) { + TORCH_CHECK(raw_cu_seqlens_.has_value(), "raw_cu_seqlens must be provided alongside cp_seq_map"); + auto const& cp_seq_map = cp_seq_map_.value(); + auto const& raw_cu_seqlens = raw_cu_seqlens_.value(); + TORCH_CHECK(cp_seq_map.device() == q.device(), "cp_seq_map must be on the same device as q"); + TORCH_CHECK(raw_cu_seqlens.device() == q.device(), "raw_cu_seqlens must be on the same device as q"); + TORCH_CHECK(cp_seq_map.dtype() == torch::kInt32, "cp_seq_map must be int32"); + TORCH_CHECK(raw_cu_seqlens.dtype() == torch::kInt32, "raw_cu_seqlens must be int32"); + TORCH_CHECK(cp_seq_map.is_contiguous(), "cp_seq_map must be contiguous"); + TORCH_CHECK(raw_cu_seqlens.is_contiguous(), "raw_cu_seqlens must be contiguous"); + TORCH_CHECK( + cp_seq_map.size(0) == num_seqs, + "cp_seq_map.size(0) must equal cu_seqlens.size(0)-1, got ", + cp_seq_map.size(0), + " vs ", + num_seqs); + cp_seq_map_ptr = cp_seq_map.data_ptr(); + raw_cu_seqlens_ptr = raw_cu_seqlens.data_ptr(); + raw_num_seqs = static_cast(raw_cu_seqlens.size(0) - 1); + } + // GVA contract on the C++ side. Order matters: check positivity *before* the modulo to // avoid % 0 / division-by-zero UB in case the Python layer passed a degenerate shape. TORCH_CHECK(num_qk_heads > 0, "KDA requires num_qk_heads > 0, got ", num_qk_heads); @@ -64,15 +91,15 @@ kda_fwd_prefill( {packed_seq, num_v_heads, head_size}, torch::TensorOptions().dtype(q.dtype()).device(q.device())); - // output_final_state controls the API side effect. If it is false, ignore - // even an explicitly provided output_state_ buffer so the kernel skips the - // final-state store. + // Allocate output state if not provided. In CP mode the state is keyed + // by raw_num_seqs (one slot per original sequence), not the inflated + // CP-chunk count. OptionalTensor output_state = std::nullopt; if (output_final_state) { output_state = output_state_.has_value() ? output_state_.value() : torch::zeros( - {num_seqs, num_v_heads, head_size, head_size}, + {raw_num_seqs, num_v_heads, head_size, head_size}, torch::TensorOptions().dtype(torch::kFloat32).device(q.device())); } @@ -123,7 +150,7 @@ kda_fwd_prefill( auto& input_state = input_state_.value(); TORCH_CHECK(input_state.dtype() == torch::kFloat32, "input_state must be float32"); TORCH_CHECK(input_state.is_contiguous(), "input_state must be contiguous"); - // Defense in depth: also enforce shape on the C++ side (Python layer should already check). + // In CP mode the leading dim is the CP-chunk count (== num_seqs), otherwise it is the raw num_seqs. TORCH_CHECK( input_state.dim() == 4 && input_state.size(0) == num_seqs && input_state.size(1) == num_v_heads && input_state.size(2) == head_size && input_state.size(3) == head_size, @@ -164,7 +191,10 @@ kda_fwd_prefill( static_cast(packed_seq), scale, safe_gate, - static_cast(sm_count)); + static_cast(sm_count), + cp_seq_map_ptr, + raw_cu_seqlens_ptr, + raw_num_seqs); } else { float const* beta_ptr = beta_.has_value() ? beta_.value().data_ptr() : nullptr; kda::sm90::launch_kda_fwd_prefill_kernel( @@ -186,7 +216,10 @@ kda_fwd_prefill( static_cast(packed_seq), scale, safe_gate, - static_cast(sm_count)); + static_cast(sm_count), + cp_seq_map_ptr, + raw_cu_seqlens_ptr, + raw_num_seqs); } return {output, output_state}; diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu new file mode 100644 index 00000000..5a0f6299 --- /dev/null +++ b/csrc/api/pybind.cu @@ -0,0 +1,99 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) +void +ChunkKDAFwdIntra( + at::Tensor q, + at::Tensor k, + at::Tensor g, + at::Tensor beta, + at::Tensor cu_seqlens, + at::Tensor chunk_indices, + at::Tensor Aqk_out, + at::Tensor Akk_out, + at::Tensor tile_counter, + float scale, + int chunk_size, + bool use_tf32_inverse, + bool unified_gref); +void +ChunkKDAFwdRecompWU( + at::Tensor k, + at::Tensor v, + at::Tensor beta, + at::Tensor A, + at::Tensor g, + at::Tensor cu_seqlens, + at::Tensor chunk_indices, + at::Tensor w_out, + at::Tensor u_out, + at::Tensor kg_out, + int chunk_size, + std::optional q, + std::optional qg_out); +#endif + +#if defined(CULA_SM90A_ENABLED) +std::tuple> +kda_fwd_prefill( + std::optional output_, + std::optional output_state_, + torch::Tensor const& q, + torch::Tensor const& k, + torch::Tensor const& v, + std::optional input_state_, + std::optional alpha_, + std::optional beta_, + torch::Tensor const& cu_seqlens, + torch::Tensor workspace_buffer, + float scale, + bool output_final_state, + bool safe_gate, + std::optional cp_seq_map_, + std::optional raw_cu_seqlens_); +#endif + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "cuLA"; +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) + m.def("chunk_kda_fwd_intra_cuda", &ChunkKDAFwdIntra); + m.def("recompute_w_u_cuda", &ChunkKDAFwdRecompWU); +#endif +#if defined(CULA_SM90A_ENABLED) + m.def( + "kda_fwd_prefill", + &kda_fwd_prefill, + pybind11::arg("output_"), + pybind11::arg("output_state_"), + pybind11::arg("q"), + pybind11::arg("k"), + pybind11::arg("v"), + pybind11::arg("input_state_"), + pybind11::arg("alpha_"), + pybind11::arg("beta_"), + pybind11::arg("cu_seqlens"), + pybind11::arg("workspace_buffer"), + pybind11::arg("scale"), + pybind11::arg("output_final_state"), + pybind11::arg("safe_gate"), + pybind11::arg("cp_seq_map_") = std::nullopt, + pybind11::arg("raw_cu_seqlens_") = std::nullopt); +#endif +} diff --git a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp index 301dbfd4..66563e5a 100644 --- a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp +++ b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp @@ -930,14 +930,26 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { v_head_idx); return; } + // Intra-card CP + int32_t out_seq_idx = seq_idx; + int32_t out_num_seqs = problem_size.num_seqs; + if (problem_size.cp_seq_map != nullptr) { + out_seq_idx = problem_size.cp_seq_map[seq_idx]; + out_num_seqs = problem_size.raw_num_seqs; + int32_t this_end = problem_size.cu_seqlens[seq_idx + 1]; + int32_t raw_end = problem_size.raw_cu_seqlens[out_seq_idx + 1]; + if (this_end != raw_end) { + return; + } + } DPRINTF0_WG("[%d,%d,%d,%d]>> save tKVrKV -> tKVgKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); // GVA: state is stored per V/O head. int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); auto gKV = make_tensor( make_gmem_ptr(params.ptr_output_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, problem_size.num_seqs)))( - _, _, state_head_idx, seq_idx); // (KDim, VDim), K-contiguous + make_layout(make_shape(Int{}, Int{}, num_state_heads, out_num_seqs)))( + _, _, state_head_idx, out_seq_idx); // (KDim, VDim), K-contiguous auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); @@ -1371,10 +1383,18 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { if constexpr (!kInitStateFromInput) { clear(tKVrKV); - compute_loop_body(0, /*is_first_block_=*/cute::true_type{}, /*is_final_block_=*/cute::false_type{}); + if (num_blocks == 1) { + compute_loop_body(0, /*is_first_block_=*/cute::true_type{}, /*is_final_block_=*/cute::true_type{}); + } else { + compute_loop_body(0, /*is_first_block_=*/cute::true_type{}, /*is_final_block_=*/cute::false_type{}); + } } else { kv_load(tKVrKV); // GMEM -> Register, only once at the beginning - compute_loop_body(0, /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::false_type{}); + if (num_blocks == 1) { + compute_loop_body(0, /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::true_type{}); + } else { + compute_loop_body(0, /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::false_type{}); + } } CUTE_NO_UNROLL for (int blk = 1; blk < num_blocks - 1; ++blk) { diff --git a/csrc/kda/sm90/kda_fwd_sm90.cu b/csrc/kda/sm90/kda_fwd_sm90.cu index d668db9b..ed855db6 100644 --- a/csrc/kda/sm90/kda_fwd_sm90.cu +++ b/csrc/kda/sm90/kda_fwd_sm90.cu @@ -53,7 +53,10 @@ launch_kda_fwd_prefill_kernel_gbai( int32_t head_size, int64_t total_seqlen, float scale, - int32_t sm_count); + int32_t sm_count, + int32_t const* cp_seq_map, + int32_t const* raw_cu_seqlens, + int32_t raw_num_seqs); template < typename ArchTag, // TODO: hide this @@ -81,7 +84,10 @@ launch_kda_fwd_prefill_kernel( int64_t total_seqlen, float scale, bool safe_gate, - int32_t sm_count = 0) { + int32_t sm_count, + int32_t const* cp_seq_map, + int32_t const* raw_cu_seqlens, + int32_t raw_num_seqs) { bool needs_beta = beta != nullptr; bool needs_alpha = alpha != nullptr; bool init_state = input_state != nullptr; @@ -105,7 +111,10 @@ launch_kda_fwd_prefill_kernel( head_size, \ total_seqlen, \ scale, \ - sm_count); + sm_count, \ + cp_seq_map, \ + raw_cu_seqlens, \ + raw_num_seqs); if (init_state) { if (needs_beta && needs_alpha && safe_gate) { LAUNCH(true, true, true, true); @@ -146,7 +155,10 @@ launch_kda_fwd_prefill_kernel( int64_t total_seqlen, float scale, bool safe_gate, - int32_t sm_count); + int32_t sm_count, + int32_t const* cp_seq_map, + int32_t const* raw_cu_seqlens, + int32_t raw_num_seqs); // TBeta=bf16 template void @@ -169,6 +181,9 @@ launch_kda_fwd_prefill_kernel( int64_t total_seqlen, float scale, bool safe_gate, - int32_t sm_count); + int32_t sm_count, + int32_t const* cp_seq_map, + int32_t const* raw_cu_seqlens, + int32_t raw_num_seqs); } // namespace kda::sm90 diff --git a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu index 309cefa0..0da2986e 100644 --- a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu +++ b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu @@ -23,92 +23,44 @@ namespace kda::sm90 { using namespace cute; using bf16 = cute::bfloat16_t; -// SafeGate=true, InitState=false -template void -launch_kda_fwd_prefill_kernel_gbai( - cudaStream_t, - bf16*, - float*, - bf16 const*, - bf16 const*, - bf16 const*, - float const*, - float const*, - float const*, - int32_t const*, - uint8_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int64_t, - float, - int32_t); +#define INSTANTIATE_GBAI(NeedsBeta, NeedsAlpha, InitState, SafeGate, TBeta) \ + template void launch_kda_fwd_prefill_kernel_gbai< \ + NeedsBeta, \ + NeedsAlpha, \ + InitState, \ + SafeGate, \ + cutlass::arch::Sm90, \ + bf16, \ + bf16, \ + float, \ + TBeta>( \ + cudaStream_t, \ + bf16*, \ + float*, \ + bf16 const*, \ + bf16 const*, \ + bf16 const*, \ + float const*, \ + float const*, \ + TBeta const*, \ + int32_t const*, \ + uint8_t*, \ + int32_t, \ + int32_t, \ + int32_t, \ + int32_t, \ + int64_t, \ + float, \ + int32_t, \ + int32_t const*, \ + int32_t const*, \ + int32_t) -// SafeGate=true, InitState=true -template void -launch_kda_fwd_prefill_kernel_gbai( - cudaStream_t, - bf16*, - float*, - bf16 const*, - bf16 const*, - bf16 const*, - float const*, - float const*, - float const*, - int32_t const*, - uint8_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int64_t, - float, - int32_t); +INSTANTIATE_GBAI(true, true, false, true, float); +INSTANTIATE_GBAI(true, true, true, true, float); +INSTANTIATE_GBAI(true, true, false, true, bf16); +INSTANTIATE_GBAI(true, true, true, true, bf16); -// SafeGate=true, InitState=false, BetaBF16 -template void -launch_kda_fwd_prefill_kernel_gbai( - cudaStream_t, - bf16*, - float*, - bf16 const*, - bf16 const*, - bf16 const*, - float const*, - float const*, - bf16 const*, - int32_t const*, - uint8_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int64_t, - float, - int32_t); - -// SafeGate=true, InitState=true, BetaBF16 -template void -launch_kda_fwd_prefill_kernel_gbai( - cudaStream_t, - bf16*, - float*, - bf16 const*, - bf16 const*, - bf16 const*, - float const*, - float const*, - bf16 const*, - int32_t const*, - uint8_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int64_t, - float, - int32_t); +#undef INSTANTIATE_GBAI } // namespace kda::sm90 diff --git a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp index 4f8ba027..ac597f7b 100644 --- a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp +++ b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp @@ -140,6 +140,11 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { int32_t num_qk_heads; int32_t num_v_heads; int32_t head_size; // d + + // For intra-card CP + int32_t const* cp_seq_map = nullptr; + int32_t const* raw_cu_seqlens = nullptr; + int32_t raw_num_seqs = 0; }; using ProblemShape = VarlenProblemShape; diff --git a/csrc/kda/sm90/prefill_kernel.hpp b/csrc/kda/sm90/prefill_kernel.hpp index d56fafae..6e54c3a3 100644 --- a/csrc/kda/sm90/prefill_kernel.hpp +++ b/csrc/kda/sm90/prefill_kernel.hpp @@ -46,6 +46,9 @@ launch_kda_fwd_prefill_kernel( int64_t total_seqlen, float scale, bool safe_gate, - int32_t sm_count = 0); + int32_t sm_count = 0, + int32_t const* cp_seq_map = nullptr, + int32_t const* raw_cu_seqlens = nullptr, + int32_t raw_num_seqs = 0); } // namespace kda::sm90 diff --git a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh index 72f13a6f..c53f2ae3 100644 --- a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh +++ b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh @@ -58,7 +58,10 @@ launch_kda_fwd_prefill_kernel_gbai( int32_t head_size, int64_t total_seqlen, float scale, - int32_t sm_count) { + int32_t sm_count, + int32_t const* cp_seq_map = nullptr, + int32_t const* raw_cu_seqlens = nullptr, + int32_t raw_num_seqs = 0) { #if defined(CULA_SM90A_ENABLED) constexpr bool HopperSupported = true; #else @@ -123,6 +126,9 @@ launch_kda_fwd_prefill_kernel_gbai( .num_qk_heads = num_qk_heads, .num_v_heads = num_v_heads, .head_size = head_size, + .cp_seq_map = cp_seq_map, + .raw_cu_seqlens = raw_cu_seqlens, + .raw_num_seqs = raw_num_seqs, }, .mainloop = { diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index 8baa41e2..f0d21e62 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -19,6 +19,8 @@ "kda_decode", "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", + "kda_prefill_hopper_opt", + "kda_prefill_hopper_auto", ] _LAZY = { diff --git a/cula/kda/auto_route.py b/cula/kda/auto_route.py new file mode 100644 index 00000000..54b90284 --- /dev/null +++ b/cula/kda/auto_route.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import torch + +from cula.kda.cp_context import is_dominant_long_seq +from cula.kda.hopper_fused_fwd import cula_kda_prefill as _basic +from cula.kda.hopper_fused_fwd_opt import FUSED_GATE_L2NORM_TH_VARLEN +from cula.kda.hopper_fused_fwd_opt import cula_kda_prefill_opt as _opt + + +def _should_use_opt(q: torch.Tensor, cu_seqlens: torch.Tensor | None) -> bool: + """Pick opt vs basic based on H100 measurements.""" + B = q.shape[0] + T = q.shape[1] + H = q.shape[2] + + if cu_seqlens is not None: + N = cu_seqlens.numel() - 1 + if N > 1: + packed_T = q.shape[1] + if packed_T * H <= FUSED_GATE_L2NORM_TH_VARLEN: + return True + if N * H <= 16 and T >= 8192: + return True + if H <= 16 and T >= 32768 + N - 1: + cu_list = cu_seqlens.tolist() + seqlens = [cu_list[i + 1] - cu_list[i] for i in range(N)] + if is_dominant_long_seq(seqlens, H): + return True + return False + # N == 1 falls through to the single-sequence logic below. + + # Fused gate+l2norm reliably wins at very small T*H even with B>1. + if T * H <= 6000: + return True + + if B == 1: + # T=1024 H=8/16 gets a small win from fused l2norm_qk (T*H<10000). + if H <= 16 and T <= 1024: + return True + if H <= 8: + return T >= 4096 # CP kicks in + elif H <= 16: + return T >= 4096 + elif H <= 32: + return T >= 16384 + else: # H >= 64 + return False # base ties or wins + + if B == 2: + if H == 8: + return T >= 4096 + return False # B=2 H>=16 mostly ties + + # B >= 4 : B*H >= 32 already saturates a sizable fraction of SMs, CP + # buys little; basic and opt tie. Default to basic (cheaper wrapper). + return False + + +def cula_kda_prefill_auto( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = True, + use_qk_l2norm_in_kernel: bool = True, + use_gate_in_kernel: bool = True, + safe_gate: bool = True, + lower_bound: float | None = -5.0, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, + **kwargs, +): + if _should_use_opt(q, cu_seqlens): + return _opt( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + auto_cp=True, + **kwargs, + ) + return _basic( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + **kwargs, + ) diff --git a/cula/kda/cp_context.py b/cula/kda/cp_context.py new file mode 100644 index 00000000..976870ed --- /dev/null +++ b/cula/kda/cp_context.py @@ -0,0 +1,452 @@ +from __future__ import annotations + +import functools +import math +import weakref + +import torch +from fla.utils import tensor_cache + +from cula.utils import get_device_sm_count + + +@tensor_cache +def _create_cu_seqlens(batch_size: int, num_tokens: int, device_idx: int, dtype: torch.dtype) -> torch.Tensor: + return torch.arange(batch_size + 1, dtype=dtype, device=f"cuda:{device_idx}") * num_tokens + + +@functools.lru_cache(maxsize=32) +def _create_full_cu_seqlens_2(T: int, device_idx: int, dtype: torch.dtype) -> torch.Tensor: + return torch.tensor([0, T], dtype=dtype, device=f"cuda:{device_idx}") + + +_CP_SEQS_CACHE: dict = {} +_SLOT_MAP_CACHE: dict = {} + + +DOMINANT_LONG_SEQ_MIN_LEN = 32768 +DOMINANT_LONG_SEQ_MAX_H = 16 + + +def is_dominant_long_seq( + seqlens: list[int], + H: int, + min_long_len: int = DOMINANT_LONG_SEQ_MIN_LEN, + max_H: int = DOMINANT_LONG_SEQ_MAX_H, +) -> bool: + if not seqlens or max_H < H: + return False + longest = max(seqlens) + if longest < min_long_len: + return False + return 2 * longest >= sum(seqlens) + + +def _seqlens_from_cu(cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor | None = None) -> list[int]: + """Return per-seq lengths from a cu_seqlens tensor (CPU sync if no CPU copy).""" + src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + cu_list = src.tolist() + return [cu_list[i + 1] - cu_list[i] for i in range(len(cu_list) - 1)] + + +def _get_slot_map(cp_cu_seqlens: torch.Tensor, T: int, chunk_size: int) -> torch.Tensor: + key = (id(cp_cu_seqlens), T, chunk_size) + hit = _SLOT_MAP_CACHE.get(key) + if hit is not None: + weak_ref, cached = hit + if weak_ref() is cp_cu_seqlens: + return cached + # id was reused after the original tensor was collected — recompute. + num_chunks = (T + chunk_size - 1) // chunk_size + cp_starts = (cp_cu_seqlens[:-1] // chunk_size).to(torch.int32) + slot_map = torch.full((num_chunks,), -1, dtype=torch.int32, device=cp_cu_seqlens.device) + slots = torch.arange(cp_starts.numel(), dtype=torch.int32, device=cp_cu_seqlens.device) + slot_map[cp_starts.long()] = slots + _SLOT_MAP_CACHE[key] = (weakref.ref(cp_cu_seqlens), slot_map) + return slot_map + + +def _calc_cp_seqs_cached(raw_cu_seqlens, chunk_size, num_v_heads, sm_count, raw_cu_seqlens_cpu=None): + key = (id(raw_cu_seqlens), chunk_size, num_v_heads, sm_count) + hit = _CP_SEQS_CACHE.get(key) + if hit is not None: + weak_ref, cached = hit + if weak_ref() is raw_cu_seqlens: + return cached + + val = _calc_cp_seqs(raw_cu_seqlens, chunk_size, num_v_heads, sm_count, raw_cu_seqlens_cpu=raw_cu_seqlens_cpu) + _CP_SEQS_CACHE[key] = (weakref.ref(raw_cu_seqlens), val) + return val + + +def _calc_cp_seqs( + raw_cu_seqlens: torch.Tensor, + chunk_size: int, + num_v_heads: int, + sm_count: int, + raw_cu_seqlens_cpu: torch.Tensor | None = None, +) -> tuple[bool, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Decide whether intra-card CP pays off and, if so, build the split tables.""" + device = raw_cu_seqlens.device + seqlen_dtype = raw_cu_seqlens.dtype + + raw_cu_seqlens_list = (raw_cu_seqlens_cpu if raw_cu_seqlens_cpu is not None else raw_cu_seqlens).tolist() + raw_batch_size = len(raw_cu_seqlens_list) - 1 + seqlens = [raw_cu_seqlens_list[i + 1] - raw_cu_seqlens_list[i] for i in range(raw_batch_size)] + num_chunks = [(s + chunk_size - 1) // chunk_size for s in seqlens] + if max(num_chunks) <= 0: + return False, None, None, None, None + + H = num_v_heads + V_BLOCKS = 1 # bump to 2 once main kernel supports V-blocking + target_cp_batch = max(1, sm_count // (H * V_BLOCKS)) + total_chunks = sum(num_chunks) + # mlc * cp_batch * chunk_size ≈ T per raw seq → mlc = total_chunks / cp_batch. + target_mlc = max(1, total_chunks // (max(1, target_cp_batch))) + # Snap to nearest power of 2; clamp to ≥ 4 to keep multi-stage pipelining alive. + max_local_chunks = 2 ** round(math.log2(max(target_mlc, 1.0))) + max_local_chunks = max(max_local_chunks, 4) + max_local_tokens = max_local_chunks * chunk_size + + cp_cu_seqlens: list[int] = [] + ht_mask: list[bool] = [] + seq_map_c2r: list[int] = [] + seq_map_r2c: list[int] = [0] + + for i, c in enumerate(num_chunks): + s = raw_cu_seqlens_list[i] + e = raw_cu_seqlens_list[i + 1] + if c > max_local_chunks: + cut = s + while True: + cp_cu_seqlens.append(cut) + ht_mask.append(False) + seq_map_c2r.append(i) + remaining = e - cut + if remaining <= max_local_tokens + chunk_size: + break + cut += max_local_tokens + ht_mask[-1] = True + else: + cp_cu_seqlens.append(s) + ht_mask.append(True) + seq_map_c2r.append(i) + seq_map_r2c.append(len(cp_cu_seqlens)) + cp_cu_seqlens.append(raw_cu_seqlens_list[-1]) + + Be = total_chunks / max(num_chunks) + use_cp = (Be * H <= 40) or (Be * H <= 56 and max(num_chunks) >= 128) + # Additional cuLA-specific guard: never bother if there is only one CP-chunk + # (no split happened). + if len(cp_cu_seqlens) - 1 == raw_batch_size: + use_cp = False + + if use_cp and raw_batch_size == 1: + T_max = max(seqlens) + if H <= 8 or H <= 16: + if T_max < 4096: + use_cp = False + elif H <= 32: + if T_max < 16384: + use_cp = False + else: # H >= 64 + use_cp = False + elif use_cp and raw_batch_size > 1: + T_packed = sum(seqlens) + native_grid = raw_batch_size * H + dominant_exception = is_dominant_long_seq(seqlens, H) + + unaligned = any(s % chunk_size != 0 for s in seqlens[:-1]) or (raw_cu_seqlens_list[-1] % chunk_size != 0) + if unaligned: + use_cp = False + elif native_grid > 16 and not dominant_exception: + # Native grid already big enough that CP's lift is marginal. + use_cp = False + elif T_packed * H <= 32768: + use_cp = False + elif H <= 8: + if T_packed < 8192: + use_cp = False + elif H <= 16: + if T_packed < 4096: + use_cp = False + else: # H >= 32 with native_grid <= 16 → only B=1 fits this, handled above + use_cp = False + + if not use_cp: + return False, None, None, None, None + + cp_cu_seqlens_t = torch.tensor(cp_cu_seqlens, dtype=seqlen_dtype, device=device) + seq_map_c2r_t = torch.tensor(seq_map_c2r, dtype=seqlen_dtype, device=device) + seq_map_r2c_t = torch.tensor(seq_map_r2c, dtype=seqlen_dtype, device=device) + ht_mask_t = torch.tensor(ht_mask, dtype=torch.bool, device=device) + return True, cp_cu_seqlens_t, seq_map_r2c_t, seq_map_c2r_t, ht_mask_t + + +def _build_raw_seq_idx( + cp_cu_seqlens: torch.Tensor, seq_map_c2r: torch.Tensor, T: int, chunk_size: int +) -> tuple[torch.Tensor, list[int]]: + """Per-chunk raw seq id: raw_seq_idx[i_t] = raw_seq containing chunk i_t.""" + NT = (T + chunk_size - 1) // chunk_size + out = torch.empty(NT, dtype=torch.int32, device=cp_cu_seqlens.device) + # cp_cu_seqlens[:-1] // chunk_size gives the chunk-index start of each CP-chunk. + cp_starts = (cp_cu_seqlens[:-1] // chunk_size).to(torch.int64) + cp_ends = ((cp_cu_seqlens[1:] + chunk_size - 1) // chunk_size).to(torch.int64) + + c2r_cpu = seq_map_c2r.tolist() + starts = cp_starts.tolist() + ends = cp_ends.tolist() + out_cpu = [0] * NT + for i in range(len(starts)): + out[starts[i] : ends[i]] = c2r_cpu[i] + for j in range(starts[i], ends[i]): + out_cpu[j] = c2r_cpu[i] + return out, out_cpu + + +_RAW_SEQ_IDX_CACHE: dict = {} + + +def _get_raw_seq_idx( + cp_cu_seqlens: torch.Tensor, seq_map_c2r: torch.Tensor, T: int, chunk_size: int +) -> tuple[torch.Tensor, list[int]]: + """weakref-guarded cache for the per-chunk raw_seq_idx tensor.""" + key = (id(cp_cu_seqlens), T, chunk_size) + hit = _RAW_SEQ_IDX_CACHE.get(key) + if hit is not None: + weak_ref, cached = hit + if weak_ref() is cp_cu_seqlens: + return cached + val = _build_raw_seq_idx(cp_cu_seqlens, seq_map_c2r, T, chunk_size) + _RAW_SEQ_IDX_CACHE[key] = (weakref.ref(cp_cu_seqlens), val) + return val + + +def _compute_cp_h0_via_fla_h( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor, + beta: torch.Tensor, + scale: float, + raw_h0: torch.Tensor | None, + cp_cu_seqlens: torch.Tensor, + seq_map_c2r: torch.Tensor, + chunk_size: int, +) -> torch.Tensor: + """Compute cp_h0 directly via FLA's per-chunk h tensor — exact, no mt needed.""" + from cula.kda.cp_h_boundary import kda_cp_h0_boundary + from cula.kda.wy_intra import kda_intra_native + + cp_batch = cp_cu_seqlens.size(0) - 1 + raw_batch = int(seq_map_c2r.max().item()) + 1 if seq_map_c2r.numel() > 0 else 1 + T = k.size(1) + + if T > CP_PREPROCESS_TILE_TOKENS: + return _compute_cp_h0_via_fla_h_tiled( + k=k, + v=v, + g_cumsum=g_cumsum, + beta=beta, + raw_h0=raw_h0, + cp_cu_seqlens=cp_cu_seqlens, + seq_map_c2r=seq_map_c2r, + chunk_size=chunk_size, + ) + + w, u, _, kg = kda_intra_native( + k=k, + v=v, + gk=g_cumsum, + beta=beta, + chunk_size=chunk_size, + ) + + slot_map = _get_slot_map(cp_cu_seqlens, T, chunk_size) + + if raw_batch > 1: + if raw_h0 is None: + H_v = v.size(2) + K = k.size(3) + V = v.size(3) + raw_h0 = torch.zeros(raw_batch, H_v, V, K, dtype=torch.float32, device=k.device) + h0_chunk0 = raw_h0[0:1] # state at chunk 0 (always raw seq 0) + raw_seq_idx, _ = _get_raw_seq_idx(cp_cu_seqlens, seq_map_c2r, T, chunk_size) + else: + h0_chunk0 = raw_h0 + raw_seq_idx = None + + cp_h0 = kda_cp_h0_boundary( + kg=kg, + w=w, + u=u, + g_cumsum=g_cumsum, + h0=h0_chunk0, + slot_map=slot_map, + num_cp=cp_batch, + chunk_size=chunk_size, + raw_h0_dense=raw_h0 if raw_batch > 1 else None, + raw_seq_idx=raw_seq_idx, + ) + del w, u, kg + return cp_h0 + + +CP_PREPROCESS_TILE_TOKENS = 16384 + + +def _compute_cp_h0_via_fla_h_tiled( + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor, + beta: torch.Tensor, + raw_h0: torch.Tensor | None, + cp_cu_seqlens: torch.Tensor, + seq_map_c2r: torch.Tensor, + chunk_size: int, +) -> torch.Tensor: + from cula.kda.cp_h_boundary import kda_cp_h0_boundary + from cula.kda.wy_intra import kda_intra_native + + assert k.size(0) == 1, "tiled CP preprocess assumes packed [1, T, H, K] layout" + cp_batch = cp_cu_seqlens.size(0) - 1 + raw_batch = int(seq_map_c2r.max().item()) + 1 if seq_map_c2r.numel() > 0 else 1 + T = k.size(1) + H = k.size(2) + K = k.size(3) + V = v.size(3) + + if raw_h0 is None: + raw_h0 = torch.zeros(raw_batch, H, V, K, dtype=torch.float32, device=k.device) + device = k.device + + assert CP_PREPROCESS_TILE_TOKENS % chunk_size == 0 + tile_tokens = CP_PREPROCESS_TILE_TOKENS + + global_slot_map = _get_slot_map(cp_cu_seqlens, T, chunk_size) + # Per-chunk raw-seq map (length NT_total). For raw_batch==1 this is all + # zeros and we don't even pass it to the kernel. + if raw_batch > 1: + global_raw_seq_idx, global_raw_seq_idx_cpu = _get_raw_seq_idx(cp_cu_seqlens, seq_map_c2r, T, chunk_size) + else: + global_raw_seq_idx = None + global_raw_seq_idx_cpu = None + + cp_h0 = torch.empty(cp_batch, H, V, K, dtype=torch.float32, device=device) + + exit_state = torch.empty(H, V, K, dtype=torch.float32, device=device) + + BT = chunk_size + BC = 16 + max_tile_T = min(tile_tokens, T) + buf_w = torch.empty(1, max_tile_T, H, K, dtype=k.dtype, device=device) + buf_u = torch.empty(1, max_tile_T, H, V, dtype=v.dtype, device=device) + buf_kg = torch.empty(1, max_tile_T, H, K, dtype=k.dtype, device=device) + buf_Akkd = torch.empty(1, max_tile_T, H, BC, dtype=torch.float32, device=device) + buf_Akk = torch.empty(1, max_tile_T, H, BT, dtype=k.dtype, device=device) + + n_tiles = (T + tile_tokens - 1) // tile_tokens + for tile_idx in range(n_tiles): + s = tile_idx * tile_tokens + e = min(s + tile_tokens, T) + is_last_tile = tile_idx == n_tiles - 1 + s_chunk = s // chunk_size + e_chunk = (e + chunk_size - 1) // chunk_size + + k_t = k[:, s:e] + v_t = v[:, s:e] + g_t = g_cumsum[:, s:e] + beta_t = beta[:, s:e] + + w_t, u_t, _, kg_t = kda_intra_native( + k=k_t, + v=v_t, + gk=g_t, + beta=beta_t, + chunk_size=chunk_size, + out_w=buf_w, + out_u=buf_u, + out_kg=buf_kg, + out_Akkd=buf_Akkd, + out_Akk=buf_Akk, + ) + + slot_map_t = global_slot_map[s_chunk:e_chunk] + + if raw_batch > 1: + first_raw_in_tile = global_raw_seq_idx_cpu[s_chunk] + if tile_idx == 0: + h0_in = raw_h0[first_raw_in_tile : first_raw_in_tile + 1] + else: + prev_tile_last_raw = global_raw_seq_idx_cpu[s_chunk - 1] + if first_raw_in_tile == prev_tile_last_raw: + h0_in = exit_state + else: + h0_in = raw_h0[first_raw_in_tile : first_raw_in_tile + 1] + raw_seq_idx_t = global_raw_seq_idx[s_chunk:e_chunk] + else: + h0_in = raw_h0 if tile_idx == 0 else exit_state + raw_seq_idx_t = None + + kda_cp_h0_boundary( + kg=kg_t, + w=w_t, + u=u_t, + g_cumsum=g_t, + h0=h0_in, + slot_map=slot_map_t, + num_cp=cp_batch, + chunk_size=chunk_size, + cp_h0_out=cp_h0, + # Skip writing exit_state on the last tile — no consumer. + exit_state=None if is_last_tile else exit_state, + raw_h0_dense=raw_h0 if raw_batch > 1 else None, + raw_seq_idx=raw_seq_idx_t, + ) + + return cp_h0 + + +def intra_card_cp_preprocess( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + raw_h0: torch.Tensor | None, + raw_cu_seqlens: torch.Tensor | None, + chunk_size: int = 64, + raw_cu_seqlens_cpu: torch.Tensor | None = None, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + assert k.dim() == 4 and k.size(0) == 1, "expected packed [1, T, H, K]" + num_v_heads = v.size(2) + sm_count = get_device_sm_count(k.device) + + if raw_cu_seqlens is None: + raw_cu_seqlens = _create_cu_seqlens(1, k.size(1), k.device.index, torch.int32) + + use_cp, cp_cu_seqlens, seq_map_r2c, seq_map_c2r, ht_mask = _calc_cp_seqs_cached( + raw_cu_seqlens, + chunk_size, + num_v_heads, + sm_count, + raw_cu_seqlens_cpu=raw_cu_seqlens_cpu, + ) + if not use_cp: + return None, None, None, None + + cp_h0 = _compute_cp_h0_via_fla_h( + q=q, + k=k, + v=v, + g_cumsum=g, + beta=beta, + scale=scale, + raw_h0=raw_h0, + cp_cu_seqlens=cp_cu_seqlens, + seq_map_c2r=seq_map_c2r, + chunk_size=chunk_size, + ) + + return cp_h0, cp_cu_seqlens, seq_map_c2r, raw_cu_seqlens diff --git a/cula/kda/cp_h_boundary.py b/cula/kda/cp_h_boundary.py new file mode 100644 index 00000000..b584bee9 --- /dev/null +++ b/cula/kda/cp_h_boundary.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BV in [16, 32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT"], +) +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_EXIT_STATE": lambda args: args["exit_state"] is not None, + "MULTI_RAW": lambda args: args["raw_seq_idx"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def _kda_h_boundary_kernel( + kg, + w, + u, + gk, + h0, # fp32 [H, V, K] OR [1, H, V, K] — state at chunk 0 + cp_h0_out, # fp32 [num_cp, H, V, K], cuLA layout + slot_map, # int32 [NT], slot_map[i_t] = output slot or -1 + exit_state, # fp32 [H, V, K] or None — end-of-tile state for next tile + raw_h0_dense, # fp32 [raw_batch, H, V, K] — per-raw-seq h0 for cross-seq resets + raw_seq_idx, # int32 [NT] — raw_seq_idx[i_t] = which raw seq chunk i_t belongs to + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_EXIT_STATE: tl.constexpr, + MULTI_RAW: tl.constexpr, +): + # Grid: (V_tiles, H). One CTA per (V-tile, head) walks all T chunks serially. + i_v, i_h = tl.program_id(0), tl.program_id(1) + NT = tl.cdiv(T, BT) + + # State tiles: 2× [BV, 64] holds the (V_tile=BV, K=128) state split into K=64 halves. + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + + # Per-batch (B=1) offset into per-head buffers. + kg_ptr = kg + i_h * K + w_ptr = w + i_h * K + u_ptr = u + i_h * V + gk_ptr = gk + i_h * K + h0_ptr = h0 + i_h * V * K if USE_INITIAL_STATE else h0 # h0 [H, V, K] or [1, H, V, K] + + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0_ptr, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_h0_2 = tl.make_block_ptr(h0_ptr, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + + if MULTI_RAW: + prev_raw = tl.load(raw_seq_idx + 0) + else: + prev_raw = 0 # unused + + for i_t in range(NT): + if MULTI_RAW: + if i_t > 0: + cur_raw = tl.load(raw_seq_idx + i_t) + if cur_raw != prev_raw: + rh_base = raw_h0_dense + cur_raw.to(tl.int64) * (H * V * K) + i_h * (V * K) + p_rh_1 = tl.make_block_ptr(rh_base, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_rh_2 = tl.make_block_ptr(rh_base, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + b_h1 = tl.load(p_rh_1, boundary_check=(0, 1)).to(tl.float32) + b_h2 = tl.load(p_rh_2, boundary_check=(0, 1)).to(tl.float32) + prev_raw = cur_raw + + # ----- Conditional boundary store ----- + # slot = slot_map[i_t]; if slot >= 0, write h to cp_h0_out[slot, i_h, ...] + slot = tl.load(slot_map + i_t) + is_boundary = slot >= 0 + # Use slot=0 as safe target when not boundary; mask blocks the store. + safe_slot = tl.maximum(slot, 0).to(tl.int64) + out_base = cp_h0_out + safe_slot * (H * V * K) + i_h * (V * K) + p_out_1 = tl.make_block_ptr(out_base, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_out_2 = tl.make_block_ptr(out_base, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + if is_boundary: + tl.store(p_out_1, b_h1, boundary_check=(0, 1)) + tl.store(p_out_2, b_h2, boundary_check=(0, 1)) + + # ----- Compute v_new = u - w @ h ----- + # w [BT, K] split into 2× [BT, 64] for the matmuls + p_w1 = tl.make_block_ptr(w_ptr, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_v_acc = tl.dot(b_w1, tl.trans(b_h1).to(b_w1.dtype)) + p_w2 = tl.make_block_ptr(w_ptr, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_v_acc += tl.dot(b_w2, tl.trans(b_h2).to(b_w2.dtype)) + + # u [BT, V_tile] + p_u = tl.make_block_ptr(u_ptr, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_u, boundary_check=(0, 1)) - b_v_acc + + # ----- Apply gate decay: h *= exp2(gk_last) ----- + last_idx = tl.minimum((i_t + 1) * BT, T) - 1 + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk_ptr + last_idx * (H * K) + o_k1, mask=(o_k1 < K), other=0.0).to(tl.float32) + b_h1 *= tl.math.exp2(b_gk_last1)[None, :] + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load(gk_ptr + last_idx * (H * K) + o_k2, mask=(o_k2 < K), other=0.0).to(tl.float32) + b_h2 *= tl.math.exp2(b_gk_last2)[None, :] + + # ----- State update: h += kg^T @ v_new (transpose_state_layout=True form) ----- + b_v_bf = b_v.to(kg.dtype.element_ty) + p_kg1 = tl.make_block_ptr(kg_ptr, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1)) + b_kg1 = tl.load(p_kg1, boundary_check=(0, 1)) + b_h1 += tl.trans(tl.dot(b_kg1, b_v_bf)) + p_kg2 = tl.make_block_ptr(kg_ptr, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1)) + b_kg2 = tl.load(p_kg2, boundary_check=(0, 1)) + b_h2 += tl.trans(tl.dot(b_kg2, b_v_bf)) + + # ----- After the loop: optionally store the end-of-tile state for the next tile ----- + if STORE_EXIT_STATE: + es_ptr = exit_state + i_h * V * K + p_es_1 = tl.make_block_ptr(es_ptr, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_es_2 = tl.make_block_ptr(es_ptr, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + tl.store(p_es_1, b_h1, boundary_check=(0, 1)) + tl.store(p_es_2, b_h2, boundary_check=(0, 1)) + + +def kda_cp_h0_boundary( + kg: torch.Tensor, # bf16 [1, T, H, K] + w: torch.Tensor, # bf16 [1, T, H, K] + u: torch.Tensor, # bf16 [1, T, H, V] + g_cumsum: torch.Tensor, # fp32 [1, T, H, K] + h0: torch.Tensor | None, # fp32 [H, V, K] (or [1, H, V, K]) — state at this tile's chunk 0 + slot_map: torch.Tensor, # int32 [NT], slot_map[i_t] = output slot or -1 + num_cp: int, + chunk_size: int = 64, + cp_h0_out: torch.Tensor | None = None, # pre-allocated [num_cp, H, V, K] fp32; allocated if None + exit_state: torch.Tensor | None = None, # if not None, kernel writes the end-of-T state to [H, V, K] fp32 + raw_h0_dense: torch.Tensor | None = None, # fp32 [raw_batch, H, V, K] for cross-raw-seq resets + raw_seq_idx: torch.Tensor | None = None, # int32 [NT] mapping chunk → raw seq idx (None for single-raw) +) -> torch.Tensor: + assert kg.is_contiguous() and w.is_contiguous() and u.is_contiguous() + assert kg.size(0) == 1 and w.size(0) == 1 and u.size(0) == 1 and g_cumsum.size(0) == 1 + assert slot_map.dtype == torch.int32 and slot_map.is_contiguous() + if raw_seq_idx is not None: + assert raw_h0_dense is not None, "raw_h0_dense required when raw_seq_idx is set" + assert raw_seq_idx.dtype == torch.int32 and raw_seq_idx.is_contiguous() + + T = kg.size(1) + H = kg.size(2) + K = kg.size(3) + V = u.size(3) + assert K == 128 and V == 128, f"Phase 1 kernel hard-codes K=V=128, got K={K} V={V}" + + if cp_h0_out is None: + cp_h0_out = torch.empty(num_cp, H, V, K, dtype=torch.float32, device=kg.device) + + BT = chunk_size + + # Grid depends on BV (autotune-selected); use meta-grid. + def grid(meta): + return (V // meta["BV"], H) + + _kda_h_boundary_kernel[grid]( + kg, + w, + u, + g_cumsum, + h0, + cp_h0_out, + slot_map, + exit_state, + raw_h0_dense, + raw_seq_idx, + T, + H=H, + K=K, + V=V, + BT=BT, + ) + return cp_h0_out diff --git a/cula/kda/gate_l2norm_fused.py b/cula/kda/gate_l2norm_fused.py new file mode 100644 index 00000000..aa93f69f --- /dev/null +++ b/cula/kda/gate_l2norm_fused.py @@ -0,0 +1,194 @@ +import torch +import triton +import triton.language as tl +from fla.ops.utils.constant import RCP_LN2 as _RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.utils.softplus import softplus + +# Triton requires module-level constants used inside @jit kernels to be +# wrapped in tl.constexpr. +RCP_LN2 = tl.constexpr(_RCP_LN2) + + +@triton.jit +def _gate_l2norm_fused_kernel( + # Pointers + g_ptr, + A_log_ptr, + dt_bias_ptr, # gate inputs + q_ptr, + k_ptr, # qk inputs + g_out_ptr, # gate output (fp32 cumsum) + yq_ptr, + yk_ptr, # qk outputs (bf16) + rstd_q_ptr, + rstd_k_ptr, # qk rstd outputs (fp32) + cu_seqlens_ptr, + chunk_indices_ptr, + # Scalars + lower_bound, + eps_l2, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t = tl.program_id(0) + i_bh = tl.program_id(1) + i_h = i_bh % H + + if IS_VARLEN: + i_n = tl.load(chunk_indices_ptr + i_t * 2).to(tl.int32) + i_t_local = tl.load(chunk_indices_ptr + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens_ptr + i_n).to(tl.int32) + eos = tl.load(cu_seqlens_ptr + i_n + 1).to(tl.int32) + T_seq = eos - bos + bt_base = bos + i_t_local * BT + valid_t = (i_t_local * BT + tl.arange(0, BT)) < T_seq + else: + bt_base = (i_bh // H) * T + i_t * BT + valid_t = (i_t * BT + tl.arange(0, BT)) < T + + rows = bt_base + tl.arange(0, BT) + + cols = tl.arange(0, BD) # (BD,) + valid_d = cols < D + + offs = rows[:, None] * (H * D) + i_h * D + cols[None, :] + mask = valid_t[:, None] & valid_d[None, :] + + # ==================================================================== + # GATE: cumsum( transform(g + bias) ) * RCP_LN2 + # ==================================================================== + b_g = tl.load(g_ptr + offs, mask=mask, other=0.0).to(tl.float32) + + if HAS_BIAS: + b_bias = tl.load(dt_bias_ptr + i_h * D + cols, mask=valid_d, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + b_A = tl.load(A_log_ptr + i_h).to(tl.float32) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(tl.exp(b_A) * b_g) + else: + b_gate = -tl.exp(b_A) * softplus(b_g) + + b_gate_cs = tl.cumsum(b_gate, axis=0) * RCP_LN2 + # zero out the tokens beyond T so we don't pollute g_out tail rows + b_gate_cs = tl.where(mask, b_gate_cs, 0.0) + tl.store(g_out_ptr + offs, b_gate_cs, mask=mask) + + # ==================================================================== + # L2-NORM Q — per-row normalisation along D, rstd written to (B*T*H,) + # ==================================================================== + b_q = tl.load(q_ptr + offs, mask=mask, other=0.0).to(tl.float32) + b_q_sq = tl.sum(b_q * b_q, axis=1) # (BT,) + b_rstd_q = 1.0 / tl.sqrt(b_q_sq + eps_l2) # (BT,) + b_yq = b_q * b_rstd_q[:, None] + tl.store(yq_ptr + offs, b_yq.to(yq_ptr.dtype.element_ty), mask=mask) + + # rstd is shape (B, T, H,) contiguous in (b*T*H + t*H + h) order + rstd_offs = rows * H + i_h # (BT,) + tl.store(rstd_q_ptr + rstd_offs, b_rstd_q, mask=valid_t) + + # ==================================================================== + # L2-NORM K + # ==================================================================== + b_k = tl.load(k_ptr + offs, mask=mask, other=0.0).to(tl.float32) + b_k_sq = tl.sum(b_k * b_k, axis=1) + b_rstd_k = 1.0 / tl.sqrt(b_k_sq + eps_l2) + b_yk = b_k * b_rstd_k[:, None] + tl.store(yk_ptr + offs, b_yk.to(yk_ptr.dtype.element_ty), mask=mask) + tl.store(rstd_k_ptr + rstd_offs, b_rstd_k, mask=valid_t) + + +def gate_l2norm_fused_fwd( + g: torch.Tensor, # (B, T, H, D) bf16 + q: torch.Tensor, # (B, T, H, D) bf16 + k: torch.Tensor, # (B, T, H, D) bf16 + A_log: torch.Tensor, # (H,) fp32 + dt_bias: torch.Tensor | None, # (H*D,) fp32 or None + lower_bound: float | None, # if not None and safe_gate -> use lb*sigmoid path + chunk_size: int = 64, + eps_l2: float = 1e-6, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, +): + """One-launch fused preprocessing. + + Returns: + g_out: (B, T, H, D) fp32 -- gate cumsum * RCP_LN2 + y_q: (B, T, H, D) bf16 -- l2-normalised q + y_k: (B, T, H, D) bf16 -- l2-normalised k + rstd_q, rstd_k: (B, T, H) fp32 -- 1/sqrt(sum^2 + eps) + """ + assert g.shape == q.shape == k.shape, f"shapes must match: g{g.shape} q{q.shape} k{k.shape}" + assert g.is_contiguous() and q.is_contiguous() and k.is_contiguous(), "all inputs must be contiguous" + B, T, H, D = g.shape + assert chunk_size == 64, "only chunk_size=64 supported (matches SM90 main kernel)" + + is_varlen = cu_seqlens is not None + if is_varlen: + assert B == 1, "varlen path expects packed B=1 layout" + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + + if dt_bias is None: + dt_bias_arg = A_log # any valid fp32 pointer; HAS_BIAS=False suppresses use + has_bias = False + else: + dt_bias_arg = dt_bias + has_bias = True + + use_lower_bound = lower_bound is not None + if not use_lower_bound: + lower_bound = 0.0 # unused, but Triton needs a value + + g_out = torch.empty_like(g, dtype=torch.float32) + y_q = torch.empty_like(q) + y_k = torch.empty_like(k) + rstd_q = torch.empty((B, T, H), dtype=torch.float32, device=q.device) + rstd_k = torch.empty((B, T, H), dtype=torch.float32, device=q.device) + + BD = triton.next_power_of_2(D) + if is_varlen: + NT = chunk_indices.shape[0] + grid = (NT, H) + cu_seqlens_arg = cu_seqlens + chunk_indices_arg = chunk_indices + else: + NT = triton.cdiv(T, chunk_size) # ceil — partial last chunk handled by mask + grid = (NT, B * H) + cu_seqlens_arg = A_log + chunk_indices_arg = A_log + num_warps = 1 if BD <= 128 else (2 if BD <= 256 else 4) + + _gate_l2norm_fused_kernel[grid]( + g_ptr=g, + A_log_ptr=A_log, + dt_bias_ptr=dt_bias_arg, + q_ptr=q, + k_ptr=k, + g_out_ptr=g_out, + yq_ptr=y_q, + yk_ptr=y_k, + rstd_q_ptr=rstd_q, + rstd_k_ptr=rstd_k, + cu_seqlens_ptr=cu_seqlens_arg, + chunk_indices_ptr=chunk_indices_arg, + lower_bound=lower_bound, + eps_l2=eps_l2, + T=T, + H=H, + D=D, + BT=chunk_size, + BD=BD, + HAS_BIAS=has_bias, + USE_LOWER_BOUND=use_lower_bound, + IS_VARLEN=is_varlen, + num_warps=num_warps, + ) + return g_out, y_q, y_k, rstd_q, rstd_k diff --git a/cula/kda/hopper_fused_fwd_opt.py b/cula/kda/hopper_fused_fwd_opt.py new file mode 100644 index 00000000..49dc6b71 --- /dev/null +++ b/cula/kda/hopper_fused_fwd_opt.py @@ -0,0 +1,390 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Optimized Hopper KDA prefill: fused gate+l2norm preprocessing + intra-card CP.""" + +import torch +from einops import rearrange +from fla.modules.l2norm import l2norm_fwd +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + +import cula.cudac as cula_cuda +from cula.kda.cp_context import intra_card_cp_preprocess, is_dominant_long_seq +from cula.kda.gate_l2norm_fused import gate_l2norm_fused_fwd +from cula.kda.l2norm_qk_fused import l2norm_fwd_qk +from cula.utils import _get_cache_buf, assert_hopper, get_device_sm_count, prepare_uniform_cu_seqlens + +FUSED_L2NORM_QK_TH_MAX = 10000 + +FUSED_GATE_L2NORM_TH_FIXED = 16384 + +FUSED_GATE_L2NORM_TH_VARLEN = 65536 + +FUSED_GATE_L2NORM_VARLEN_AVG_SEQ = 256 + + +def _fused_gate_l2norm_threshold(cu_seqlens_is_none): + return FUSED_GATE_L2NORM_TH_FIXED if cu_seqlens_is_none else FUSED_GATE_L2NORM_TH_VARLEN + + +FUSED_GATE_L2NORM_TH_MAX = FUSED_GATE_L2NORM_TH_VARLEN + + +def _inference_forward( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + safe_gate, + lower_bound, + cu_seqlens, + chunk_indices, + auto_cp, + cu_seqlens_cpu=None, +): + chunk_size = 64 + batch_size, seq_len, num_qk_heads, head_dim = q.shape + num_v_heads = v.shape[-2] + + cu_seqlens_is_none = cu_seqlens is None + if cu_seqlens_is_none: + cu_seqlens = prepare_uniform_cu_seqlens(batch_size, seq_len, q.device, torch.int32) + if batch_size != 1: + q, k, v, g, beta = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (q, k, v, g, beta)) + + if cu_seqlens_is_none: + avg_seq_ok = True + else: + N = cu_seqlens.numel() - 1 + packed_T = q.shape[1] + avg_seq_ok = N <= 1 or packed_T <= N * FUSED_GATE_L2NORM_VARLEN_AVG_SEQ + + fused_all_pre = ( + use_gate_in_kernel + and use_qk_l2norm_in_kernel + and (q.numel() // q.shape[-1]) <= _fused_gate_l2norm_threshold(cu_seqlens_is_none) + and num_qk_heads == num_v_heads + and avg_seq_ok + ) + + if fused_all_pre: + if chunk_indices is None and not cu_seqlens_is_none: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) + g_out, yq, yk, _, _ = gate_l2norm_fused_fwd( + g=g, + q=q, + k=k, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound if safe_gate else None, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens if not cu_seqlens_is_none else None, + chunk_indices=chunk_indices if not cu_seqlens_is_none else None, + ) + g, q, k = g_out, yq, yk + else: + if chunk_indices is None and not cu_seqlens_is_none: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) + if use_gate_in_kernel: + g = kda_gate_chunk_cumsum( + g=g, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + else: + g = chunk_local_cumsum( + g=g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + if use_qk_l2norm_in_kernel: + D = q.shape[-1] + n_rows = q.numel() // D + if n_rows <= FUSED_L2NORM_QK_TH_MAX: + q_flat = q.view(-1, D) + k_flat = k.view(-1, D) + yq_flat, yk_flat, _, _ = l2norm_fwd_qk(q_flat, k_flat) + q = yq_flat.view_as(q) + k = yk_flat.view_as(k) + else: + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + + packed_seq = batch_size * seq_len + q = q.reshape(packed_seq, num_qk_heads, head_dim).contiguous() + k = k.reshape(packed_seq, num_qk_heads, head_dim).contiguous() + v = v.reshape(packed_seq, num_v_heads, head_dim).contiguous() + g = g.reshape(packed_seq, num_v_heads, head_dim).contiguous() + beta = beta.reshape(packed_seq, num_v_heads).contiguous() + + cp_seq_map = None + raw_cu_seqlens_for_cp = None + + def _dominant_long_seq_gate() -> bool: + if cu_seqlens is None: + return False + N_ = cu_seqlens.numel() - 1 + if N_ <= 1 or N_ * num_v_heads <= 16: + return False # falls under the main branch + if num_v_heads > 16 or packed_seq < 8192: + return False + cu_src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + cu_list = cu_src.tolist() + seqlens = [cu_list[i + 1] - cu_list[i] for i in range(N_)] + return is_dominant_long_seq(seqlens, num_v_heads) + + cp_would_fire = auto_cp and ( + # Multi-seq varlen: CP only when grid is starved AND total T amortizes. + ( + cu_seqlens is not None + and (cu_seqlens.numel() - 1) > 1 + and (cu_seqlens.numel() - 1) * num_v_heads <= 16 + and packed_seq >= 8192 + ) + # Multi-seq dominant-seq exception. + or _dominant_long_seq_gate() + or + # Single sequence: per-H T thresholds matching _calc_cp_seqs. + ( + (cu_seqlens is None or cu_seqlens.numel() - 1 == 1) + and ( + (num_v_heads <= 8 and packed_seq >= 4096) + or (num_v_heads <= 16 and packed_seq >= 4096) + or (num_v_heads <= 32 and packed_seq >= 16384) + ) + ) + ) + if cp_would_fire: + q4 = q.view(1, packed_seq, num_qk_heads, head_dim) + k4 = k.view(1, packed_seq, num_qk_heads, head_dim) + v4 = v.view(1, packed_seq, num_v_heads, head_dim) + g4 = g.view(1, packed_seq, num_v_heads, head_dim) + beta4 = beta.view(1, packed_seq, num_v_heads) + cp_h0, cp_cu_seqlens, cp_seq_map, raw_cu_seqlens_for_cp = intra_card_cp_preprocess( + q=q4, + k=k4, + v=v4, + g=g4, + beta=beta4, + scale=scale, + raw_h0=initial_state, + raw_cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + raw_cu_seqlens_cpu=cu_seqlens_cpu, + ) + del q4, k4, v4, g4, beta4 + if cp_seq_map is not None: + cu_seqlens = cp_cu_seqlens + initial_state = cp_h0 + + sm_count = get_device_sm_count(q.device) + workspace_buffer = _get_cache_buf("hopper_kda_fwd_workspace", sm_count * 128, q.device) + + o, final_state = cula_cuda.kda_fwd_prefill( + None, + None, + q, + k, + v, + initial_state, + g, + beta, + cu_seqlens, + workspace_buffer, + scale, + output_final_state, + safe_gate, + cp_seq_map_=cp_seq_map, + raw_cu_seqlens_=raw_cu_seqlens_for_cp, + ) + o = rearrange(o, "(b t) h d -> b t h d", b=batch_size) + return o.to(q.dtype), final_state + + +class HopperChunkKDAFunctionOpt(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, + auto_cp: bool = True, + cu_seqlens_cpu: torch.IntTensor | None = None, + ): + return _inference_forward( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + safe_gate, + lower_bound, + cu_seqlens, + chunk_indices, + auto_cp, + cu_seqlens_cpu=cu_seqlens_cpu, + ) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + raise NotImplementedError("Backward pass is not implemented yet.") + + +@torch.compiler.disable +def cula_kda_prefill_opt( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, + auto_cp: bool = True, + cu_seqlens_cpu: torch.IntTensor | None = None, + **kwargs, +): + assert_hopper() + assert safe_gate, "Only support safe_gate=True." + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError(f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`.") + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if initial_state is not None: + assert initial_state.dtype == torch.float32, "initial_state must be in float32." + + A_log, dt_bias = None, None + if use_gate_in_kernel: + assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." + A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") + if safe_gate: + if lower_bound is None: + raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") + if not (-5 <= lower_bound < 0): + raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") + + assert q.shape == k.shape, "q and k must have the same shape." + assert q.shape[:2] == v.shape[:2] == g.shape[:2], "q, k, v, g must share batch and sequence dimensions." + batch_size, seq_len, num_qk_heads, head_dim = q.shape + num_v_heads = v.shape[-2] + assert num_qk_heads > 0 and num_v_heads > 0 + assert num_v_heads % num_qk_heads == 0 + assert g.shape == (batch_size, seq_len, num_v_heads, head_dim) + assert v.shape == (batch_size, seq_len, num_v_heads, head_dim) + assert beta.shape == (batch_size, seq_len, num_v_heads) + assert q.dtype == k.dtype == v.dtype == torch.bfloat16, "q, k, v must be in bfloat16." + assert beta.dtype == torch.bfloat16 or beta.dtype == torch.float32, "beta must be in bfloat16 or float32." + assert q.shape[-1] == k.shape[-1] == v.shape[-1] == 128, "Currently we only support head dim of 128 for KDA" + if scale is None: + scale = k.shape[-1] ** -0.5 + + needs_grad = torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v, g, beta) if t is not None) + if not needs_grad: + o, final_state = _inference_forward( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + safe_gate, + lower_bound, + cu_seqlens, + chunk_indices, + auto_cp, + cu_seqlens_cpu=cu_seqlens_cpu, + ) + return o, (final_state if output_final_state else None) + + o, final_state = HopperChunkKDAFunctionOpt.apply( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + safe_gate, + lower_bound, + cu_seqlens, + chunk_indices, + auto_cp, + cu_seqlens_cpu, + ) + + return o, (final_state if output_final_state else None) diff --git a/cula/kda/l2norm_qk_fused.py b/cula/kda/l2norm_qk_fused.py new file mode 100644 index 00000000..234d57bf --- /dev/null +++ b/cula/kda/l2norm_qk_fused.py @@ -0,0 +1,119 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Fused l2-norm for paired (q, k) tensors — one Triton kernel handles both. + +cuLA's baseline `cula_kda_prefill` calls `l2norm_fwd(q)` and `l2norm_fwd(k)` +as two separate Triton kernel launches. Each launch costs ~50-80 μs of CPU +overhead (Python wrapper + torch.empty + CUDA driver dispatch), even though +the actual GPU work is tiny for D=128. + +The two operations are mathematically identical and operate on disjoint +inputs/outputs. By writing one kernel whose grid covers both q and k +(distinguishing them via `tl.program_id(1)`), we cut the Python-driver +overhead in half — saving one launch (~50 μs) per fwd at any T. + +Combined with skipping the gate-stream optimization (which didn't pay off +due to torch.cuda.stream context overhead — see hopper_fused_fwd_opt.py), +this is the cleanest small-T speedup we found. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _l2norm_fwd_qk_kernel( + q_ptr, + k_ptr, # input pointers (T*H, D) each + yq_ptr, + yk_ptr, # output pointers (T*H, D) each + rstd_q_ptr, + rstd_k_ptr, # output rstd (T*H,) each + eps, + D, + BD: tl.constexpr, +): + i_row = tl.program_id(0) # 0..T*H-1 + i_qk = tl.program_id(1) # 0=q, 1=k + + cols = tl.arange(0, BD) + mask = cols < D + + # is_q is uniform across all threads in the block (driven by + # tl.program_id(1)), so the if/else compiles to a single conditional + # branch — no warp divergence, and only one tensor is actually loaded. + is_q = i_qk == 0 + base_off = i_row * D + if is_q: + b_x = tl.load(q_ptr + base_off + cols, mask=mask, other=0.0).to(tl.float32) + else: + b_x = tl.load(k_ptr + base_off + cols, mask=mask, other=0.0).to(tl.float32) + + b_rstd = 1.0 / tl.sqrt(tl.sum(b_x * b_x) + eps) + b_y = b_x * b_rstd + + # Symmetric stores — same uniform-branch logic as the load above. + if is_q: + tl.store(yq_ptr + base_off + cols, b_y, mask=mask) + tl.store(rstd_q_ptr + i_row, b_rstd) + else: + tl.store(yk_ptr + base_off + cols, b_y, mask=mask) + tl.store(rstd_k_ptr + i_row, b_rstd) + + +def l2norm_fwd_qk( + q: torch.Tensor, + k: torch.Tensor, + eps: float = 1e-6, +): + """L2-normalize q and k along the last dim, in a single fused kernel. + + Args: + q, k: shape (..., D). Last dim is normalised; preceding dims are + treated as a flat "row" index. q and k must have identical shape. + eps: numerical safety eps. + + Returns: + (y_q, y_k, rstd_q, rstd_k) + y_q, y_k: normalized outputs, same shape as q, k. + rstd_q, rstd_k: 1/sqrt(sum(x^2)+eps), shape q.shape[:-1]. + """ + assert q.shape == k.shape, f"q.shape {q.shape} != k.shape {k.shape}" + assert q.dtype == k.dtype + assert q.device == k.device + assert q.is_contiguous() and k.is_contiguous(), "q, k must be contiguous" + + D = q.shape[-1] + T = q.numel() // D + + y_q = torch.empty_like(q) + y_k = torch.empty_like(k) + rstd_q = torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device) + rstd_k = torch.empty(k.shape[:-1], dtype=torch.float32, device=k.device) + + BD = triton.next_power_of_2(D) + if D > BD or 65536 // q.element_size() < BD: + raise RuntimeError(f"D={D} too large for fused l2norm_fwd_qk") + + # Grid: (T*H, 2) — program_id(1) picks q (0) or k (1). + # Heuristic on num_warps: small D (e.g. 128) needs only 1 warp; up to 4. + num_warps = 1 if BD <= 256 else (2 if BD <= 1024 else 4) + _l2norm_fwd_qk_kernel[(T, 2)]( + q_ptr=q, + k_ptr=k, + yq_ptr=y_q, + yk_ptr=y_k, + rstd_q_ptr=rstd_q, + rstd_k_ptr=rstd_k, + eps=eps, + D=D, + BD=BD, + num_warps=num_warps, + ) + return y_q, y_k, rstd_q, rstd_k diff --git a/cula/kda/wy_intra.py b/cula/kda/wy_intra.py new file mode 100644 index 00000000..58e741c6 --- /dev/null +++ b/cula/kda/wy_intra.py @@ -0,0 +1,354 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from cula.kda.wy_recompute import kda_recompute_w_u + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4] for num_stages in [2, 3, 4] + ], + key=["H", "K", "BT", "BC"], +) +@triton.jit(do_not_specialize=["T"]) +def _kda_intra_sub_chunk_kernel( + k, + g, + beta, + Akkd, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + bos = i_b * T + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + # Per-head pointer offsets + k_base = k + (bos * H + i_h) * K + g_base = g + (bos * H + i_h) * K + beta_base = beta + bos * H + i_h + Akkd_base = Akkd + (bos * H + i_h) * BC + + # Load k, g (BC × BK), beta (BC) for this sub-chunk + p_k = tl.make_block_ptr(k_base, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g_base, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta_base, (T,), (H,), (i_ti,), (BC,), (0,)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + o_gn = i_ti + tl.minimum(BC // 2, T - i_ti - 1) + o_k = tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(g + (bos * H + i_h) * K + o_gn * (H * K) + o_k, mask=m_k, other=0.0).to(tl.float32) + + b_gm = (b_g - b_gn[None, :]).to(tl.float32) + b_gq = tl.where(m_c[:, None], tl.math.exp2(b_gm), 0.0) + b_gk = tl.where(m_c[:, None], tl.math.exp2(-b_gm), 0.0) + + b_kgt = tl.trans(b_k * b_gk) + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Akkd = tl.make_block_ptr(Akkd_base, (T, BC), (H * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Akkd, b_Akk.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + tl.debug_barrier() + + b_Ai = -b_Akk + for i in range(2, tl.minimum(BC, T - i_ti)): + b_a = -tl.load(Akkd_base + (i_ti + i) * (H * BC) + o_i) + b_a = tl.where(o_i < i, b_a, 0.0) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + + b_Ai += m_I + + tl.store(p_Akkd, b_Ai.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + + +_SOLVE_DOT_PRECISION = tl.constexpr("tf32") + + +@triton.autotune( + configs=[triton.Config({"BK": BK}, num_warps=num_warps) for BK in [32, 64] for num_warps in [1, 2, 4]], + key=["H", "K", "BC"], +) +@triton.jit(do_not_specialize=["T"]) +def _kda_intra_inter_solve_kernel( + k, + g, + beta, + Akkd, + Akk, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + bos = i_b * T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + k_base = k + (bos * H + i_h) * K + g_base = g + (bos * H + i_h) * K + Akk_base = Akk + (bos * H + i_h) * BT + Akkd_base = Akkd + (bos * H + i_h) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr(k_base, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_g0 = tl.make_block_ptr(g_base, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + # sub-chunk 1 (vs sub-chunk 0) + if i_tc1 < T: + p_k1 = tl.make_block_ptr(k_base, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_g1 = tl.make_block_ptr(g_base, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + b_gn1 = tl.load(g + (bos * H + i_h) * K + i_tc1 * (H * K) + o_k, mask=m_k, other=0.0).to(tl.float32) + b_gqn = tl.where(m_tc1[:, None], tl.math.exp2(b_g1 - b_gn1[None, :]), 0.0) + b_kgt = tl.trans(b_k0 * tl.math.exp2(b_gn1[None, :] - b_g0)) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + # sub-chunk 2 (vs 0 and 1) + if i_tc2 < T: + p_k2 = tl.make_block_ptr(k_base, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g_base, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + b_gn2 = tl.load(g + (bos * H + i_h) * K + i_tc2 * (H * K) + o_k, mask=m_k, other=0.0).to(tl.float32) + b_gqn2 = tl.where(m_tc2[:, None], tl.math.exp2(b_g2 - b_gn2[None, :]), 0.0) + b_kg2 = b_k2 * b_gqn2 + b_kgt0 = tl.trans(b_k0 * tl.math.exp2(b_gn2[None, :] - b_g0)) + b_Akk20 += tl.dot(b_kg2, b_kgt0) + b_kgt1 = tl.trans(b_k1 * tl.math.exp2(b_gn2[None, :] - b_g1)) + b_Akk21 += tl.dot(b_kg2, b_kgt1) + + # sub-chunk 3 (vs 0, 1, 2) + if i_tc3 < T: + p_k3 = tl.make_block_ptr(k_base, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g_base, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + b_gn3 = tl.load(g + (bos * H + i_h) * K + i_tc3 * (H * K) + o_k, mask=m_k, other=0.0).to(tl.float32) + b_gqn3 = tl.where(m_tc3[:, None], tl.math.exp2(b_g3 - b_gn3[None, :]), 0.0) + b_kg3 = b_k3 * b_gqn3 + b_kgt0 = tl.trans(b_k0 * tl.math.exp2(b_gn3[None, :] - b_g0)) + b_Akk30 += tl.dot(b_kg3, b_kgt0) + b_kgt1 = tl.trans(b_k1 * tl.math.exp2(b_gn3[None, :] - b_g1)) + b_Akk31 += tl.dot(b_kg3, b_kgt1) + b_kgt2 = tl.trans(b_k2 * tl.math.exp2(b_gn3[None, :] - b_g2)) + b_Akk32 += tl.dot(b_kg3, b_kgt2) + + beta_base = beta + bos * H + i_h + if i_tc1 < T: + p_b1 = tl.make_block_ptr(beta_base, (T,), (H,), (i_tc1,), (BC,), (0,)) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_b2 = tl.make_block_ptr(beta_base, (T,), (H,), (i_tc2,), (BC,), (0,)) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_b3 = tl.make_block_ptr(beta_base, (T,), (H,), (i_tc3,), (BC,), (0,)) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + # Load 4 inverted diagonal blocks (from sub_chunk kernel) + p_Akk00 = tl.make_block_ptr(Akkd_base, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd_base, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akkd_base, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akkd_base, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=_SOLVE_DOT_PRECISION), + b_Ai00, + input_precision=_SOLVE_DOT_PRECISION, + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=_SOLVE_DOT_PRECISION), + b_Ai11, + input_precision=_SOLVE_DOT_PRECISION, + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=_SOLVE_DOT_PRECISION), + b_Ai22, + input_precision=_SOLVE_DOT_PRECISION, + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=_SOLVE_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=_SOLVE_DOT_PRECISION), + input_precision=_SOLVE_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=_SOLVE_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=_SOLVE_DOT_PRECISION), + input_precision=_SOLVE_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=_SOLVE_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=_SOLVE_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=_SOLVE_DOT_PRECISION), + input_precision=_SOLVE_DOT_PRECISION, + ) + + # Store 10 blocks to the full BT×BT Akk buffer. + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + tl.store(p, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + tl.store(p, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + tl.store(p, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + tl.store(p, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + tl.store(p, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)) + tl.store(p, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + tl.store(p, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + tl.store(p, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) + tl.store(p, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p = tl.make_block_ptr(Akk_base, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)) + tl.store(p, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def kda_intra_native( + k: torch.Tensor, # bf16 [1, T, H, K] + v: torch.Tensor, # bf16 [1, T, H, V] + gk: torch.Tensor, # fp32 [1, T, H, K] + beta: torch.Tensor, # bf16 [1, T, H] + chunk_size: int = 64, + q: torch.Tensor | None = None, # bf16 [1, T, H, K], for qg (optional) + need_qg: bool = False, + out_w: torch.Tensor | None = None, # bf16 [1, T, H, K] + out_u: torch.Tensor | None = None, # bf16 [1, T, H, V] + out_kg: torch.Tensor | None = None, # bf16 [1, T, H, K] + out_Akkd: torch.Tensor | None = None, # fp32 [1, T, H, BC] + out_Akk: torch.Tensor | None = None, # bf16 [1, T, H, BT] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]: + B, T, H, K = k.shape + V = v.shape[-1] + assert B == 1, "intra expects packed [1, T, H, K] input" + assert K == 128 and V == 128, f"specialized for K=V=128, got K={K} V={V}" + BT = chunk_size + BC = 16 + NT = (T + BT - 1) // BT + NC = BT // BC + + if out_Akkd is not None: + Akkd = out_Akkd[:B, :T] + else: + Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) + if out_Akk is not None: + Akk = out_Akk[:B, :T] + Akk.zero_() + else: + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + + # Step 1: per-sub-chunk diagonal Akk inversion + BK_sub = triton.next_power_of_2(K) # =128 for K=128 + grid_sub = (NT, NC, B * H) + _kda_intra_sub_chunk_kernel[grid_sub]( + k=k, + g=gk, + beta=beta, + Akkd=Akkd, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK_sub, + ) + + # Step 2: per-chunk off-diagonal + assemble full Akk_inv + grid_inter = (NT, B * H) + _kda_intra_inter_solve_kernel[grid_inter]( + k=k, + g=gk, + beta=beta, + Akkd=Akkd, + Akk=Akk, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + ) + if out_Akkd is None: + del Akkd + + # Step 3: recompute w, u, kg (and optionally qg) from Akk + w, u, qg, kg = kda_recompute_w_u( + k=k, + v=v, + beta=beta, + A=Akk, + q=q if need_qg else None, + gk=gk, + chunk_size=chunk_size, + out_w=out_w, + out_u=out_u, + out_kg=out_kg, + ) + # Akk is similarly dead after recompute's kernel is queued. + if out_Akk is None: + del Akk + return w, u, qg, kg diff --git a/cula/kda/wy_recompute.py b/cula/kda/wy_recompute.py new file mode 100644 index 00000000..edfb9693 --- /dev/null +++ b/cula/kda/wy_recompute.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT"], +) +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def _kda_recompute_wuk_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, +): + """K = V = 128, BT = 64 specialized. BK = K, BV = V (no inner loop).""" + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + bos = i_b * T + + # Per-head pointer offsets + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + # ----- u = A @ (β · v) ----- + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb) + p_u = tl.make_block_ptr(u + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + # ----- Load k, gk, compute β·exp2(gk)·k and kg ----- + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_gk = tl.make_block_ptr(gk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + b_exp_gk = tl.math.exp2(b_gk) + + # w = A @ (β · exp2(gk) · k) + b_kb = b_k * b_b[:, None] * b_exp_gk + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + # qg = q · exp2(gk) (optional) + if STORE_QG: + p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * b_exp_gk + p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + + # kg = β · exp2(g_chunk_end - gk) · k (optional, needed by cp_h0 path) + if STORE_KG: + last_idx = tl.minimum(i_t * BT + BT, T) - 1 + o_k = tl.arange(0, K) + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=o_k < K, other=0.0).to(tl.float32) + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_kg = b_k * tl.where(m_t[:, None], tl.math.exp2(b_gn[None, :] - b_gk), 0.0) + p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + +def kda_recompute_w_u( + k: torch.Tensor, # bf16 [B, T, H, K] + v: torch.Tensor, # bf16 [B, T, H, V] + beta: torch.Tensor, # bf16 [B, T, H] + A: torch.Tensor, # bf16 [B, T, H, BT] + q: torch.Tensor | None, # bf16 [B, T, H, K], or None + gk: torch.Tensor, # fp32 [B, T, H, K] + chunk_size: int = 64, + out_w: torch.Tensor | None = None, + out_u: torch.Tensor | None = None, + out_kg: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + B, T, H, K = k.shape + V = v.shape[-1] + BT = chunk_size + assert K == 128 and V == 128, f"specialized for K=V=128, got K={K} V={V}" + assert A.shape[-1] == BT, f"expected A.shape[-1]={BT}, got {A.shape[-1]}" + + w = out_w[:B, :T] if out_w is not None else torch.empty_like(k) + u = out_u[:B, :T] if out_u is not None else torch.empty_like(v) + qg = torch.empty_like(q) if q is not None else None + if gk is not None: + kg = out_kg[:B, :T] if out_kg is not None else torch.empty_like(k) + else: + kg = None + + NT = triton.cdiv(T, BT) + grid = (NT, B * H) + _kda_recompute_wuk_kernel[grid]( + q=q, + k=k, + qg=qg, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return w, u, qg, kg diff --git a/tests/test_intracard_cp_sm90.py b/tests/test_intracard_cp_sm90.py new file mode 100644 index 00000000..2353a66e --- /dev/null +++ b/tests/test_intracard_cp_sm90.py @@ -0,0 +1,439 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# Licensed under the Apache License, Version 2.0. +"""Tests for SM90 intra-card CP: dispatch routing + numerical accuracy. + +Mirrors tests/test_intracard_cp.py (SM100 version) but targets the Hopper +(SM90) `kda_prefill_hopper_opt` / `kda_prefill_hopper_auto` path. + +Three reference levels: + - cuLA basic (kda_prefill_hopper) — same C++ kernel, no CP scheduling + → verifies CP scheduling is value-preserving + - cuLA opt with auto_cp=False — opt Python wrapper but CP disabled + → isolates the CP code paths + - FLA chunk_kda (cross-impl reference) → source of truth for end-to-end output + +The CP path is exercised through: + - kda_prefill_hopper_auto (router picks opt when shape benefits from CP) + - kda_prefill_hopper_opt(auto_cp=True) (force CP entry; bypasses router) +""" + +from __future__ import annotations + +import math +import pathlib +import sys + +import pytest +import torch + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from fla.ops.kda import chunk_kda as fla_chunk_kda # noqa: E402 +from fla.utils import assert_close # noqa: E402 + +from cula.kda import ( # noqa: E402 + kda_prefill_hopper, + kda_prefill_hopper_auto, + kda_prefill_hopper_opt, +) +from cula.kda.cp_context import _calc_cp_seqs # noqa: E402 +from cula.utils import get_device_sm_count # noqa: E402 + +BT, D = 64, 128 +DEVICE = "cuda" +DTYPE = torch.bfloat16 +LOWER_BOUND = -5.0 + +# Tolerances — same convention as tests/test_intracard_cp.py: +# * Same-kernel (CP scheduling only): torch.testing.assert_close +# (CP-on vs CP-off both go through cuLA kernels) +# * Cross-impl (vs FLA): fla.utils.assert_close(ratio=...) +ATOL_SAME_KERNEL = 1e-2 +RTOL_SAME_KERNEL = 1e-2 +RATIO_VS_FLA = 0.015 # bf16 cross-impl noise band (matches SM100 test) +RATIO_STRESS = 1e-6 # deterministic re-run: drift implies race + + +pytestmark = [ + pytest.mark.sm90_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +# ============================== Helpers ============================== + + +def _cu_from_seq_lens(seq_lens, device=DEVICE): + cu = [0] + for s in seq_lens: + cu.append(cu[-1] + s) + return torch.tensor(cu, dtype=torch.int32, device=device) + + +def make_varlen_inputs(seq_lens, H, *, use_h0=False, seed=42): + """Build varlen-packed B=1 inputs for kda_prefill_hopper_*.""" + total = sum(seq_lens) + N = len(seq_lens) + cu = _cu_from_seq_lens(seq_lens) + torch.manual_seed(seed) + q = torch.randn(1, total, H, D, dtype=DTYPE, device=DEVICE) + k = torch.randn(1, total, H, D, dtype=DTYPE, device=DEVICE) + v = torch.randn(1, total, H, D, dtype=DTYPE, device=DEVICE) + g = -torch.rand(1, total, H, D, dtype=torch.float32, device=DEVICE).abs() * 0.5 + beta = torch.randn(1, total, H, dtype=torch.float32, device=DEVICE).sigmoid().to(DTYPE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) * 0.1 if use_h0 else None + return q, k, v, g, beta, h0, A_log, dt_bias, cu + + +# ---- entry points under test ---- + + +def _common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu): + return dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=1.0 / math.sqrt(D), + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=LOWER_BOUND, + cu_seqlens=cu, + ) + + +def run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu): + return kda_prefill_hopper(**_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu)) + + +def run_cula_opt_no_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu): + return kda_prefill_hopper_opt( + **_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu), + auto_cp=False, + ) + + +def run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu): + """Force CP entry through opt wrapper.""" + return kda_prefill_hopper_opt( + **_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu), + auto_cp=True, + ) + + +def run_cula_auto(q, k, v, g, beta, h0, A_log, dt_bias, cu): + """Adaptive router — exercises the production entry point.""" + return kda_prefill_hopper_auto( + **_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu), + ) + + +def run_fla(q, k, v, g, beta, h0, A_log, dt_bias, cu): + """FLA reference. cuLA returns ht as [N, HV, V, K]; FLA's default layout + is [N, HV, K, V]. We pass ``transpose_state_layout=True`` so its output + matches cuLA's layout — no manual transpose needed before assert_close. + """ + return fla_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=1.0 / math.sqrt(D), + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=LOWER_BOUND, + cu_seqlens=cu.long(), + transpose_state_layout=True, + ) + + +# ---- assertions ---- + + +def _assert_same_kernel(name, actual, ref): + """torch.testing.assert_close with tight atol/rtol (CP-on vs CP-off use the + same C++ kernel; the only delta is per-chunk recurrence reordering).""" + if actual is None or ref is None: + assert actual is ref, f"{name}: one is None and the other isn't" + return + torch.testing.assert_close( + actual.float(), + ref.float(), + atol=ATOL_SAME_KERNEL, + rtol=RTOL_SAME_KERNEL, + msg=lambda m: f"{name}: {m}", + ) + + +def assert_cp_engages(cu, H): + """Fail fast if _calc_cp_seqs won't engage CP for this shape — without + that, the test silently checks CP-off vs CP-off. + """ + num_sms = get_device_sm_count(torch.device(DEVICE)) + use_cp, cp_cu, *_ = _calc_cp_seqs( + cu, + BT, + H, + num_sms, + raw_cu_seqlens_cpu=cu.cpu(), + ) + assert use_cp and cp_cu is not None, f"_calc_cp_seqs returned use_cp=False for cu={cu.tolist()} H={H}" + n_sub = int(cp_cu.numel() - 1) + raw_batch = int(cu.numel() - 1) + assert n_sub > raw_batch, f"CP didn't split: n_sub={n_sub} == raw_batch={raw_batch}" + + +# ====================== Dispatch path: CP vs no-CP ====================== +# Verifies kda_prefill_hopper_opt(auto_cp=True) routes through CP and matches +# the same-kernel no-CP baseline (kda_prefill_hopper). + +DISPATCH_CONFIGS = [ + # (seq_lens, H, use_h0) + ([32768], 4, False), + ([32768], 4, True), + ([65536], 4, True), + ([32768], 8, False), + ([65536], 8, True), + ([16384, 16384], 4, True), + ([28672, 4096], 4, True), + ([131072, 1024], 4, False), +] + + +@pytest.mark.parametrize("seq_lens,H,use_h0", DISPATCH_CONFIGS) +def test_cp_matches_basic_baseline(seq_lens, H, use_h0): + """CP-on (opt+auto_cp) output equals basic baseline (no-CP).""" + q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs( + seq_lens, + H, + use_h0=use_h0, + ) + assert_cp_engages(cu, H) + with torch.inference_mode(): + o_base, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) + o_cp, ht_cp = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) + _assert_same_kernel("o", o_cp, o_base) + _assert_same_kernel("ht", ht_cp, ht_base) + + +@pytest.mark.parametrize("seq_lens,H,use_h0", DISPATCH_CONFIGS) +def test_auto_router_matches_basic_baseline(seq_lens, H, use_h0): + """kda_prefill_hopper_auto output (whatever path it picks) equals basic baseline.""" + q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs( + seq_lens, + H, + use_h0=use_h0, + ) + with torch.inference_mode(): + o_base, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) + o_auto, ht_auto = run_cula_auto(q, k, v, g, beta, h0, A_log, dt_bias, cu) + _assert_same_kernel("o", o_auto, o_base) + _assert_same_kernel("ht", ht_auto, ht_base) + + +def test_cp_off_matches_basic_baseline(): + """opt with auto_cp=False must match basic (no CP, no fused-pre divergence).""" + seq_lens, H, use_h0 = [32768], 4, True + q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=use_h0) + with torch.inference_mode(): + o_base, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) + o_off, ht_off = run_cula_opt_no_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) + _assert_same_kernel("o", o_off, o_base) + _assert_same_kernel("ht", ht_off, ht_base) + + +# ====================== Cross-impl: CP vs FLA ====================== + +VS_FLA_CONFIGS = [ + ([32768], 4), + ([65536], 4), + ([32768], 8), + ([16384, 16384], 4), + ([28672, 4096], 4), + ([131072, 1024], 4), +] + + +# Irregular varlen lengths +IRREGULAR_VARLEN_CONFIGS = [ + ([1], 4), + ([63], 4), + ([64], 4), + ([65], 4), + ([129], 4), + ([1, 63, 64, 65, 129], 4), + ([1, 63, 64, 65, 129], 8), + ([129, 65, 64, 63, 1], 4), + ([1024, 1, 63, 65, 129], 4), + ([4096, 1, 63, 64, 65, 129], 4), + ([4096, 1, 63, 64, 65, 129], 8), + ([8192, 1, 31, 63, 65, 127, 129, 255], 4), + ([1] * 8 + [63] * 4 + [129] * 2, 4), + ([255, 257, 511, 513], 4), +] + + +@pytest.mark.parametrize("seq_lens,H", IRREGULAR_VARLEN_CONFIGS) +def test_irregular_varlen_vs_fla(seq_lens, H): + """Irregular varlen lengths""" + q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) + with torch.inference_mode(): + o_fla, ht_fla = run_fla(q, k, v, g, beta, None, A_log, dt_bias, cu) + o_opt, ht_opt = run_cula_opt_no_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) + assert_close(f"o (cu={cu.tolist()},H={H})", o_fla, o_opt, ratio=RATIO_VS_FLA) + assert_close(f"ht (cu={cu.tolist()},H={H})", ht_fla, ht_opt, ratio=RATIO_VS_FLA) + + +@pytest.mark.parametrize("seq_lens,H", IRREGULAR_VARLEN_CONFIGS) +def test_irregular_varlen_opt_matches_basic(seq_lens, H): + """Irregular varlen: opt path (may take fused gate+l2norm) equals basic baseline.""" + q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) + with torch.inference_mode(): + o_base, ht_base = run_cula_basic(q, k, v, g, beta, None, A_log, dt_bias, cu) + o_opt, ht_opt = run_cula_opt_no_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) + _assert_same_kernel("o", o_opt, o_base) + _assert_same_kernel("ht", ht_opt, ht_base) + + +@pytest.mark.parametrize("seq_lens,H", VS_FLA_CONFIGS) +def test_cp_vs_fla(seq_lens, H): + """CP output matches FLA chunk_kda reference (cross-impl).""" + q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) + assert_cp_engages(cu, H) + with torch.inference_mode(): + o_fla, ht_fla = run_fla(q, k, v, g, beta, None, A_log, dt_bias, cu) + o_cp, ht_cp = run_cula_opt_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) + assert_close(f"o (cu={cu.tolist()},H={H})", o_fla, o_cp, ratio=RATIO_VS_FLA) + assert_close(f"ht (cu={cu.tolist()},H={H})", ht_fla, ht_cp, ratio=RATIO_VS_FLA) + + +# ====================== Final state ht correctness ====================== +# Per-sequence ht must be independently correct for prefill→decode handoff. + +FINAL_STATE_CONFIGS = [ + ([65536], 4, False), + ([65536], 4, True), + ([65536, 16384], 4, True), + ([28672, 4096], 4, False), + ([131072, 1024], 4, True), +] + + +@pytest.mark.parametrize("seq_lens,H,use_h0", FINAL_STATE_CONFIGS) +def test_cp_final_state_per_seq(seq_lens, H, use_h0): + """Each sequence's ht matches basic baseline independently (no cross-leakage).""" + q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=use_h0) + assert_cp_engages(cu, H) + with torch.inference_mode(): + _, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) + _, ht_cp = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) + assert ht_cp is not None and ht_cp.shape == ht_base.shape, ( + f"shape mismatch: cp={tuple(ht_cp.shape)} base={tuple(ht_base.shape)}" + ) + for i in range(len(seq_lens)): + _assert_same_kernel(f"ht[{i}] (len={seq_lens[i]})", ht_cp[i], ht_base[i]) + + +# ====================== Stress: race / non-determinism ====================== +# CP's per-chunk preprocess + main kernel — re-running same inputs must +# produce bit-identical outputs (no race, no order-dependence). + +STRESS_ITERS = 50 + + +@pytest.mark.parametrize( + "seq_lens,H,use_h0", + [ + pytest.param([65536], 4, True, id="single-64K-H4-h0"), + pytest.param([65536, 4096], 4, True, id="multi-64K+4K-H4-h0"), + ], +) +def test_cp_stress_repeat(seq_lens, H, use_h0): + """Run CP N times; every iter must match the first (deterministic).""" + q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs( + seq_lens, + H, + use_h0=use_h0, + seed=20260516, + ) + assert_cp_engages(cu, H) + with torch.inference_mode(): + o_ref, ht_ref = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) + torch.cuda.synchronize() + for i in range(STRESS_ITERS): + o_i, ht_i = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) + torch.cuda.synchronize() + assert_close(f"iter {i} o", o_ref, o_i, ratio=RATIO_STRESS) + assert_close(f"iter {i} ht", ht_ref, ht_i, ratio=RATIO_STRESS) + + +# ====================== h0=None equivalence ====================== +# We patched cp_context.py so raw_h0=None synthesizes a zero pool. Verify the +# kernel result is numerically equivalent to passing an explicit zero h0. + + +def test_cp_h0_none_equiv_h0_zeros(): + """h0=None must produce identical ht to h0=zeros (no implicit init drift).""" + seq_lens, H = [65536, 4096], 4 + assert_cp_engages(_cu_from_seq_lens(seq_lens), H) + q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs( + seq_lens, + H, + use_h0=False, + seed=20260501, + ) + h0_zeros = torch.zeros(len(seq_lens), H, D, D, dtype=torch.float32, device=DEVICE) + with torch.inference_mode(): + o_none, ht_none = run_cula_opt_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) + o_zeros, ht_zeros = run_cula_opt_cp(q, k, v, g, beta, h0_zeros, A_log, dt_bias, cu) + torch.cuda.synchronize() + o_diff = (o_none.float() - o_zeros.float()).abs().max().item() + ht_diff = (ht_none.float() - ht_zeros.float()).abs().max().item() + assert o_diff < 1e-3, f"o: h0=None vs h0=zeros max abs diff {o_diff:.4e}" + assert ht_diff < 1e-4, f"ht: h0=None vs h0=zeros max abs diff {ht_diff:.4e}" + + +# ====================== CP bypass: shapes where _calc_cp_seqs returns False ====================== +# When CP heuristic says "don't split", auto_cp=True must produce bit-identical +# output to basic (because the kernel takes the same path). + +BYPASS_CONFIGS = [ + ([2048], 8), # H=8 single seq T<=2048 → no CP + ([16384], 64), # H=64 → CP never fires (per _calc_cp_seqs H>=64 branch) + ([4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096], 8), # native_grid 64 >> 16 + ([131072] + [1024] * 5, 8), # raw_batch big enough that native_grid > 16 +] + + +@pytest.mark.parametrize("seq_lens,H", BYPASS_CONFIGS) +def test_cp_bypass_matches_basic(seq_lens, H): + """When CP heuristic skips, auto_cp=True must be a no-op (same output as basic).""" + q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) + num_sms = get_device_sm_count(torch.device(DEVICE)) + use_cp, cp_cu, *_ = _calc_cp_seqs(cu, BT, H, num_sms, raw_cu_seqlens_cpu=cu.cpu()) + n_sub = int(cp_cu.numel() - 1) if cp_cu is not None else 0 + assert not use_cp or n_sub == len(seq_lens), ( + f"expected bypass for cu={cu.tolist()} H={H}, got use_cp={use_cp} n_sub={n_sub}" + ) + with torch.inference_mode(): + o_base, ht_base = run_cula_basic(q, k, v, g, beta, None, A_log, dt_bias, cu) + o_cp, ht_cp = run_cula_opt_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) + _assert_same_kernel("o", o_cp, o_base) + _assert_same_kernel("ht", ht_cp, ht_base) From 8d572b914daae319224677c61bb49bb77a4e8ec6 Mon Sep 17 00:00:00 2001 From: Longxmas <92327126+Longxmas@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:25:40 +0800 Subject: [PATCH 29/34] [KDA] KDA MTP decode: recurrent + KVBuffer chunkwise verify + flush (#96) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CuTe DSL kernels for KDA (Kimi Delta Attention) multi-token-prediction decode and speculative verify, in two families behind unified dispatchers. Recurrent verify — one kernel with register-resident gated-delta-rule state: - vk: single warp per CTA, each lane owns a K-slice with butterfly-shuffle K-reductions; the small-batch and single-token (T=1) path. - ws: 4-warp warp-specialized — warp 0 stages the shared q/k and gate in SMEM while the others run the V-tile recurrence; the large-batch path. KVBuffer chunkwise verify + flush — verify emits the output plus a compact per-token (u, k, g) buffer; flush rebuilds the rank-m state: - shuffle: token-parallel SIMT. - tensor_core: CuTe tensor-core GEMM, flat over T. Dispatch selects shuffle or tensor_core by batch / HV / T. Verify snapshots per-step state for rollback and commits the accepted step via the sglang state-scatter; both gate forms (safe-gate lower_bound and softplus) are supported. Adds benchmarks and unit tests. --- REPO_LAYOUT.md | 4 +- benchmarks/bench_kda_decode_mtp.py | 762 ++++++++++ cula/kda/__init__.py | 9 + cula/ops/__init__.py | 9 + cula/ops/kda/decode/mtp.py | 2163 +++++++++++++++++++++++++++ cula/ops/kda/decode/mtp_kvbuffer.py | 1950 ++++++++++++++++++++++++ tests/test_kda_decode_mtp.py | 837 +++++++++++ 7 files changed, 5733 insertions(+), 1 deletion(-) create mode 100644 benchmarks/bench_kda_decode_mtp.py create mode 100644 cula/ops/kda/decode/mtp.py create mode 100644 cula/ops/kda/decode/mtp_kvbuffer.py create mode 100644 tests/test_kda_decode_mtp.py diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 51f44511..81358296 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -34,8 +34,10 @@ cuLA/ │ │ │ ├── fwd_o.py # output (chunk_gla_fwd_o) │ │ │ ├── bwd_wy_dqkg.py# backward wy/dqkg fused (used by chunk_bwd) │ │ │ └── cp/ # SM100 intracard-CP: chunk_delta_h, pre_scan, merge -│ │ ├── decode/ # single-token decode +│ │ ├── decode/ # single-token + MTP decode │ │ │ ├── cute.py # kda_decode / fused_sigmoid_gating_delta_rule_update (CuTe DSL) +│ │ │ ├── mtp.py # kda_decode_mtp recurrent / recurrent_ws MTP verify (CuTe DSL) +│ │ │ ├── mtp_kvbuffer.py # KVBuffer chunkwise MTP verify (shuffle / tensor_core) + flush │ │ │ └── reference_fla.py │ │ └── experimental/sm100_fused/ # [exp] unwired fully-fused │ │ ├── kda_fully_fused_wip.py # KDAChunkwise (~6k lines) diff --git a/benchmarks/bench_kda_decode_mtp.py b/benchmarks/bench_kda_decode_mtp.py new file mode 100644 index 00000000..20599454 --- /dev/null +++ b/benchmarks/bench_kda_decode_mtp.py @@ -0,0 +1,762 @@ +"""KDA MTP decode benchmark — recurrent vs KVBuffer (chunkwise) verify CHAIN. + +Unified bench (supersedes the old forward-only bench_kda_decode_mtp and +bench_kda_kvbuffer). Variants, selectable via --only / --profile: + recurrent verify: vk / tri (official Triton), all writing T*d^2 states; + kvbuffer verify: shuffle (token-parallel) / tensor_core (CuTe tensor-core GEMM + form, flat-in-T), both writing the compact u-buffer; + forward-only baselines (no rollback cost, breakdown table only): kv / auto / loop. + +Chain: REC = recurrent verify (writes T·d² intermediate states) + commit; KVB = +kvbuffer verify (emit output + write a compact u-buffer) + flush (rank-m rebuild of +S_m). spd = REC / KVB. The commit uses the REAL sglang fused_mamba_state_scatter_with_mask +(from KDA_SCATTER_FILE) so the recurrent rollback cost is official code, not a model. + +Self-contained (inlines input/timing helpers). Triton recurrent baseline (numerical +check only) from KDA_TRITON_FILE; scatter commit from KDA_SCATTER_FILE. +""" + +import argparse +import importlib.util +import os + +import torch + +from cula.ops.kda.decode.cute import kda_decode +from cula.ops.kda.decode.mtp import ( + kda_decode_mtp, + kda_decode_mtp_recurrent, + kda_decode_mtp_recurrent_ws, +) +from cula.ops.kda.decode.mtp_kvbuffer import kda_flush_kvbuffer + +# shuffle-kvbuffer (token-parallel, structure B) is optional too. +try: + from cula.ops.kda.decode.mtp_kvbuffer import kda_decode_mtp_shuffle_kvbuffer + + _HAVE_SHUFFLE = True +except Exception: + _HAVE_SHUFFLE = False + +# tensor_core-kvbuffer (CuTe tensor-core, flat-in-T verify). +try: + from cula.ops.kda.decode.mtp_kvbuffer import kda_decode_mtp_tensor_core_kvbuffer + + _HAVE_TCORE = True +except Exception: + _HAVE_TCORE = False + + +def _load_from_file(path, attr): + """Load a single attribute from a standalone .py file via importlib.""" + spec = importlib.util.spec_from_file_location(f"_standalone_{attr}", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return getattr(mod, attr) + + +# Triton recurrent baseline (numerical check only). +_HAVE_TRITON, _TRITON_ERR = True, "" +fused_sigmoid_gating_delta_rule_update = None +try: + _f = os.environ.get("KDA_TRITON_FILE", "") + if _f and os.path.exists(_f): + fused_sigmoid_gating_delta_rule_update = _load_from_file(_f, "fused_sigmoid_gating_delta_rule_update") + else: + from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update, + ) +except Exception as e: + _HAVE_TRITON, _TRITON_ERR = False, repr(e) + +# Official sglang scatter commit (update_mamba_state_after_mtp_verify). +_HAVE_SCATTER, _SCATTER_ERR = True, "" +fused_mamba_state_scatter_with_mask = None +try: + _f = os.environ.get("KDA_SCATTER_FILE", "") + if _f and os.path.exists(_f): + fused_mamba_state_scatter_with_mask = _load_from_file(_f, "fused_mamba_state_scatter_with_mask") + else: + from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( + fused_mamba_state_scatter_with_mask, + ) +except Exception as e: + _HAVE_SCATTER, _SCATTER_ERR = False, repr(e) + + +def make_dense_inputs(N, T, H, HV, K, V, device, seed=42): + g = torch.Generator(device=device).manual_seed(seed) + bf16 = torch.bfloat16 + q = torch.randn(N, T, H, K, device=device, dtype=bf16, generator=g) + k = torch.randn(N, T, H, K, device=device, dtype=bf16, generator=g) + v = torch.randn(N, T, HV, V, device=device, dtype=bf16, generator=g) + a = (torch.randn(N, T, HV, K, device=device, dtype=torch.float32, generator=g) * 0.1).to(bf16) + b = torch.randn(N, T, HV, device=device, dtype=bf16, generator=g) + A_log = -torch.rand(HV, device=device, dtype=torch.float32, generator=g) * 2 + dt_bias = torch.randn(HV, K, device=device, dtype=torch.float32, generator=g) * 0.1 + state = torch.randn(N, HV, V, K, device=device, dtype=torch.float32, generator=g) * 0.01 + indices = torch.arange(N, device=device, dtype=torch.int32) + return q, k, v, a, b, A_log, dt_bias, state, indices + + +def to_triton_varlen(q, k, v, a, b): + N, T, H, K = q.shape + HV, V = v.shape[2], v.shape[3] + NT = N * T + q_t = q.reshape(1, NT, H, K).contiguous() + k_t = k.reshape(1, NT, H, K).contiguous() + v_t = v.reshape(1, NT, HV, V).contiguous() + a_t = a.reshape(1, NT, HV * K).contiguous() + b_t = b.reshape(1, NT, HV).contiguous() + cu_seqlens = torch.arange(0, (N + 1) * T, T, device=q.device, dtype=torch.int32) + return q_t, k_t, v_t, a_t, b_t, cu_seqlens + + +def make_triton_call( + qt, + kt, + vt, + at, + bt, + cu_seqlens, + A_log, + dt_bias, + state, + indices, + scale, + dsu, + inter_buf=None, + inter_idx=None, + cache_steps=None, +): + """Official sglang recurrent verify. In verify mode (inter_buf set) it writes the T·d² + intermediate_states_buffer, same rollback cost as our production recurrent_v.""" + + def call(): + return fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=at, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=qt, + k=kt, + v=vt, + b=bt, + initial_state_source=state, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + is_kda=True, + disable_state_update=dsu, + intermediate_states_buffer=inter_buf, + intermediate_state_indices=inter_idx, + cache_steps=cache_steps, + retrieve_parent_token=None, + lower_bound=None, + ) + + return call + + +def warmup(fn, n): + for _ in range(n): + fn() + torch.cuda.synchronize() + + +def t_graph_ms(fn, warmup_iters, rep, graph_calls=1): + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(warmup_iters): + fn() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(graph_calls): + fn() + for _ in range(10): + g.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(rep): + g.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / rep / graph_calls + + +_VK_BV = -1 +_ONLY = set() # empty = all variants + + +def _want(name): + return not _ONLY or name in _ONLY + + +def make_vk_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu, inter_buf=None): + """Production recurrent vk. In verify mode (inter_buf set) it writes the T·d² + intermediate_states_buffer — the rollback cost kvbuffer replaces with a u-buffer.""" + + def call(): + return kda_decode_mtp_recurrent( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + disable_state_update=dsu, + variant="vk", + bv=_VK_BV, + intermediate_states_buffer=inter_buf, + ) + + return call + + +def make_recurrent_ws_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu, inter_buf=None): + """Production warp-spec recurrent (recurrent_ws). In verify mode (inter_buf set) it also writes T*d^2 states.""" + + def call(): + return kda_decode_mtp_recurrent_ws( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + disable_state_update=dsu, + intermediate_states_buffer=inter_buf, + ) + + return call + + +def make_shuffle_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu, ubufs=None): + """shuffle-kvbuffer (token-parallel chunkwise, structure B) — target: verify latency ~flat in T. + tile_v / ilp_rows overridable via env KDA_SHUFFLE_TILE_V / KDA_SHUFFLE_ILP_ROWS (-1 = auto).""" + d_buf, k_buf, g_buf = ubufs if ubufs is not None else (None, None, None) + _tv = int(os.environ.get("KDA_SHUFFLE_TILE_V", "-1")) + _ilp = int(os.environ.get("KDA_SHUFFLE_ILP_ROWS", "-1")) + + def call(): + return kda_decode_mtp_shuffle_kvbuffer( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + disable_state_update=dsu, + emit_output=True, + d_buffer=d_buf, + k_buffer=k_buf, + g_buffer=g_buf, + tile_v=_tv, + ilp_rows=_ilp, + ) + + return call + + +def make_tcore_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu, ubufs=None): + """CuTe tensor-core tensor_core-kvbuffer. env KDA_TCORE_BV / KDA_TCORE_NUM_V_TILES (-1 = auto).""" + d_buf, k_buf, g_buf = ubufs if ubufs is not None else (None, None, None) + _bv = int(os.environ.get("KDA_TCORE_BV", "32")) + _num_v_tiles = int(os.environ.get("KDA_TCORE_NUM_V_TILES", "-1")) + + def call(): + return kda_decode_mtp_tensor_core_kvbuffer( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + disable_state_update=dsu, + emit_output=True, + d_buffer=d_buf, + k_buffer=k_buf, + g_buffer=g_buf, + bv=_bv, + num_v_tiles=_num_v_tiles, + ) + + return call + + +def make_kv_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu): + """Forward-only production kv (lane=V recurrent; no intermediate-state support).""" + state_kv = state.transpose(-2, -1).contiguous() # vk->kv once, outside timing + + def call(): + return kda_decode_mtp_recurrent( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state_kv, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + disable_state_update=dsu, + variant="kv", + ) + + return call + + +def make_auto_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu, inter_buf=None): + """kda_decode_mtp dispatch (recurrent vk).""" + + def call(): + return kda_decode_mtp( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + disable_state_update=dsu, + state_layout="vk", + intermediate_states_buffer=inter_buf, + ) + + return call + + +def make_loop_call(q, k, v, a, b, A_log, dt_bias, state, indices, scale, dsu): + """Per-token kda_decode loop baseline (slices pre-cut; kda_decode always writes state).""" + N, T = q.shape[0], q.shape[1] + HV, V = v.shape[2], v.shape[3] + qs = [q[:, t].unsqueeze(1).contiguous() for t in range(T)] + ks = [k[:, t].unsqueeze(1).contiguous() for t in range(T)] + vs = [v[:, t].unsqueeze(1).contiguous() for t in range(T)] + as_ = [a[:, t].unsqueeze(1).contiguous() for t in range(T)] + bs = [b[:, t].unsqueeze(1).contiguous() for t in range(T)] + st = state.clone().contiguous() + o = torch.empty(N, T, HV, V, device=q.device, dtype=torch.bfloat16) + + def call(): + for t in range(T): + o_t = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=qs[t], + k=ks[t], + v=vs[t], + a=as_[t], + b=bs[t], + initial_state_source=st, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + o[:, t] = o_t.squeeze(1) + return o + + return call + + +# ---- verify-chain components: commit (recurrent rollback) & flush (kvbuffer) ---- +def make_scatter_commit_call(state_pool, inter_buf, m, N, T, HV, V, K): + """Recurrent rollback via the OFFICIAL sglang fused_mamba_state_scatter_with_mask: + gather each request's accepted-step state from the intermediate cache into the pool + (num_layers=1; step = m-1 for all requests).""" + dst = state_pool.view(1, N, HV, V, K) # [layers, cache, *state] + src = inter_buf.view(1, N, T, HV, V, K) # [layers, req, step, *state] + dst_idx = torch.arange(N, device=state_pool.device, dtype=torch.int32) + step_idx = torch.full((N,), m - 1, device=state_pool.device, dtype=torch.int32) + + def call(): + fused_mamba_state_scatter_with_mask(dst, src, dst_idx, step_idx) + return state_pool + + return call + + +def make_gather_commit_call(state_pool, inter_buf, m): + """Recurrent rollback, strided gather model: copy inter_buf[:,m-1] (a T-strided view) + into the pool. Less coalesced than the official kernel — kept for sensitivity only.""" + midx = m - 1 + + def call(): + state_pool.copy_(inter_buf[:, midx]) + return state_pool + + return call + + +def make_flush_call(state_pool, indices, ubufs, m): + """KVBuffer flush: read the compact u-buffer, rank-m rebuild S_m (no recompute).""" + d_b, k_b, g_b = ubufs + + def call(): + return kda_flush_kvbuffer(state_pool, indices, d_b, k_b, g_b, m) + + return call + + +def _accept_len(T, accept, N=0): + if accept == "full": + return T + if accept == "half": + return max(1, (T + 1) // 2) + if accept == "one": + return 1 + if accept == "random": + # Deterministic per-(N,T) accept length in [1,T] (real serving is per-req variable). + g = torch.Generator().manual_seed(1000 * N + T) + return int(torch.randint(1, T + 1, (1,), generator=g).item()) + return max(1, min(int(accept), T)) + + +def _profile_one(args, DSU, device): + """Run ONE method's kernel in a loop so ncu can wrap it. Shape = (batch_sizes[0], Ts[0]).""" + N, T = args.batch_sizes[0], args.Ts[0] + q, k, v, a, b, A_log, dt_bias, state0, indices = make_dense_inputs(N, T, args.H, args.HV, args.K, args.V, device) + scale = args.K**-0.5 + m = _accept_len(T, args.accept, N) + inter_buf = torch.empty(N, T, args.HV, args.V, args.K, dtype=torch.float32, device=device) + ubufs = ( + torch.empty(N, T, args.HV, args.V, dtype=torch.float32, device=device), + torch.empty(N, T, args.HV, args.K, dtype=torch.float32, device=device), + torch.empty(N, T, args.HV, args.K, dtype=torch.float32, device=device), + ) + p = args.profile + if p == "recurrent": + fn = make_vk_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, inter_buf) + elif p == "recurrent_ws": + fn = make_recurrent_ws_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, inter_buf) + elif p == "shuffle": + fn = make_shuffle_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs) + elif p == "tensor_core": + fn = make_tcore_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs) + elif p == "triton": + qt, kt, vt, at, bt, cu = to_triton_varlen(q, k, v, a, b) + tri_idx = torch.arange(N, device=device, dtype=torch.int32) + fn = make_triton_call( + qt, kt, vt, at, bt, cu, A_log, dt_bias, state0.clone(), indices, scale, DSU, inter_buf, tri_idx, T + ) + elif p == "commit": + make_vk_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, inter_buf)() + fn = make_scatter_commit_call(state0.clone(), inter_buf, m, N, T, args.HV, args.V, args.K) + elif p == "recurrent_kv": + fn = make_kv_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU) + elif p == "auto": + fn = make_auto_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU) + elif p == "loop": + fn = make_loop_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU) + elif p == "flush": + make_shuffle_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs)() + fn = make_flush_call(state0.clone(), indices, ubufs, m) + for _ in range(5): + fn() + torch.cuda.synchronize() + for _ in range(args.profile_iters): + fn() + torch.cuda.synchronize() + print(f"profiled {p} N={N} T={T} HV={args.HV} m={m} iters={args.profile_iters}") + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8]) + ap.add_argument("--Ts", type=int, nargs="+", default=[2, 3, 4, 6, 8]) + ap.add_argument("--H", type=int, default=32) + ap.add_argument("--HV", type=int, default=32) + ap.add_argument("--K", type=int, default=128) + ap.add_argument("--V", type=int, default=128) + ap.add_argument("--rep", type=int, default=300) + ap.add_argument("--warmup", type=int, default=5, help="warmup iters before each timed segment") + ap.add_argument( + "--graph-calls", + type=int, + default=20, + help="ops per CUDA graph to amortize fixed launch overhead at small batch " + "(N<16; N>=16 uses 1). needs idempotent dsu=1.", + ) + ap.add_argument( + "--dsu", + type=int, + default=1, + choices=[0, 1], + help="disable_state_update; 1=forward-only (idempotent, default), 0=write state", + ) + ap.add_argument("--vk-bv", type=int, default=-1, choices=[-1, 8, 16, 32]) + ap.add_argument( + "--accept", default="full", help="chain accept length m: full(=T)/half/one/random/; drives commit/flush." + ) + ap.add_argument( + "--commit", + default="scatter", + choices=["scatter", "gather"], + help="recurrent commit model: scatter=official sglang " + "fused_mamba_state_scatter_with_mask (coalesced N·d², default); " + "gather=strided copy (sensitivity). kvbuffer flush always counted.", + ) + ap.add_argument( + "--only", + nargs="+", + default=[], + choices=["recurrent", "recurrent_ws", "triton", "shuffle", "tensor_core", "recurrent_kv", "auto", "loop"], + help="restrict check/timing to these verify variants (default: all). REC/spd columns show n/a for skipped baselines.", + ) + ap.add_argument("--check", action="store_true", help="numerical check only, no timing") + ap.add_argument("--atol", type=float, default=5e-2) + ap.add_argument( + "--profile", + default="", + choices=[ + "", + "recurrent", + "recurrent_ws", + "shuffle", + "tensor_core", + "triton", + "commit", + "flush", + "recurrent_kv", + "auto", + "loop", + ], + help="ncu profile mode: run one method's kernel in a loop (uses batch-sizes[0], Ts[0])", + ) + ap.add_argument("--profile-iters", type=int, default=20, help="kernel launches in the profiled loop") + args = ap.parse_args() + + global _VK_BV + _VK_BV = args.vk_bv + global _ONLY + _ONLY = set(args.only) + DSU = bool(args.dsu) + device = "cuda" + if args.profile: + _profile_one(args, DSU, device) + return + print(f"GPU: {torch.cuda.get_device_name()}") + print( + f"shape H={args.H} HV={args.HV} K={args.K} V={args.V} dsu={DSU} shuffle_impl={_HAVE_SHUFFLE} tensor_core_impl={_HAVE_TCORE}" + ) + + # ---------------- numerical check (vs Triton recurrent) ---------------- + if not _HAVE_TRITON: + print(f"[warn] Triton baseline unavailable ({_TRITON_ERR}); skipping numerical check.") + else: + print(f"\n=== numerical check (max|Δ| vs Triton recurrent, threshold {args.atol}) ===") + print( + f"{'N':>4} {'T':>3} | {'Δ recurrent':>10} | {'Δ recurrent_ws':>10} | {'Δ shuffle':>10} | {'Δ tensor_core':>10} | flag" + ) + for N in args.batch_sizes: + for T in args.Ts: + q, k, v, a, b, A_log, dt_bias, state0, indices = make_dense_inputs( + N, T, args.H, args.HV, args.K, args.V, device + ) + scale = args.K**-0.5 + qt, kt, vt, at, bt, cu = to_triton_varlen(q, k, v, a, b) + o_tri = make_triton_call(qt, kt, vt, at, bt, cu, A_log, dt_bias, state0.clone(), indices, scale, True)() + o_tri = o_tri.reshape(N, T, args.HV, args.V) + d_recurrent = float("nan") + if _want("recurrent"): + o_recurrent = make_vk_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, True)() + d_recurrent = (o_recurrent - o_tri).abs().max().item() + d_recurrent_ws = float("nan") + if _want("recurrent_ws"): + o_recurrent_ws = make_recurrent_ws_call( + q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, True + )() + d_recurrent_ws = (o_recurrent_ws - o_tri).abs().max().item() + d_shuffle = float("nan") + if _HAVE_SHUFFLE and _want("shuffle"): + o_shuffle = make_shuffle_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, True)() + d_shuffle = (o_shuffle - o_tri).abs().max().item() + d_tensor_core = float("nan") + if _HAVE_TCORE and _want("tensor_core"): + o_tensor_core = make_tcore_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, True)() + d_tensor_core = (o_tensor_core - o_tri).abs().max().item() + cand = [x for x in (d_recurrent, d_recurrent_ws, d_shuffle, d_tensor_core) if x == x] + flag = ("OK" if max(cand) < args.atol else "DIFF!") if cand else "n/a" + print( + f"{N:>4} {T:>3} | {d_recurrent:>10.2e} | {d_recurrent_ws:>10.2e} | {d_shuffle:>10.2e} | {d_tensor_core:>10.2e} | {flag}" + ) + + if args.check: + return + + _timing_verify_chain(args, DSU, device) + + +def _timing_verify_chain(args, DSU, device): + """Fair spec-decode verify CHAIN (each segment timed in its own CUDA graph, summed). All verify + kernels run dsu=1 + verify-mode: recurrent vk/triton write the T·d² intermediate states, + kvbuffer writes its compact u-buffer. REC = recurrent verify + commit; KVB = kvbuffer verify + + flush. spd_recurrent = REC/KVB vs production recurrent; spd_bf = official triton REC chain + / kvbuffer KVB chain. Prints chain totals + speedups first, per-segment breakdown after.""" + + def us(x): + return f"{x * 1e3:.1f}" if x else "n/a" + + def rat(a_, b_): + return f"{a_ / b_:.2f}x" if (a_ and b_) else "n/a" + + if args.commit == "scatter" and not _HAVE_SCATTER: + raise RuntimeError( + f"commit=scatter needs the official sglang kernel; set KDA_SCATTER_FILE to " + f"mamba_state_scatter_triton.py (load error: {_SCATTER_ERR})" + ) + + # ---- measure every segment for every (N, T) into `results` ---- + results = [] + for N in args.batch_sizes: + for T in args.Ts: + q, k, v, a, b, A_log, dt_bias, state0, indices = make_dense_inputs(N, T, args.H, args.HV, args.K, args.V, device) + scale = args.K**-0.5 + m = _accept_len(T, args.accept, N) + gc = 1 if N >= 16 else args.graph_calls # amortize launch overhead at small batch + inter_buf = torch.empty(N, T, args.HV, args.V, args.K, dtype=torch.float32, device=device) + ubufs = ( + torch.empty(N, T, args.HV, args.V, dtype=torch.float32, device=device), + torch.empty(N, T, args.HV, args.K, dtype=torch.float32, device=device), + torch.empty(N, T, args.HV, args.K, dtype=torch.float32, device=device), + ) + tg = {} + + def time_seg(fn): + warmup(fn, args.warmup) + return t_graph_ms(fn, args.warmup, args.rep, gc) + + # recurrent verify (dsu=1, writes T·d² states) + commit + if _want("recurrent"): + tg["recurrent_v"] = time_seg( + make_vk_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, inter_buf) + ) + if _want("recurrent_ws"): + tg["recurrent_ws_v"] = time_seg( + make_recurrent_ws_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, inter_buf) + ) + if _want("recurrent") or _want("recurrent_ws") or _want("triton"): + if args.commit == "scatter": + fn_cmt = make_scatter_commit_call(state0.clone(), inter_buf, m, N, T, args.HV, args.V, args.K) + else: + fn_cmt = make_gather_commit_call(state0.clone(), inter_buf, m) + tg["cmt"] = time_seg(fn_cmt) + # kvbuffer verify (dsu=1, writes u-buffer) + flush + if _want("shuffle") or _want("tensor_core"): + # flush needs a populated u-buffer: run one kvbuffer verify first to fill it + if _HAVE_SHUFFLE and _want("shuffle"): + make_shuffle_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs)() + elif _HAVE_TCORE and _want("tensor_core"): + make_tcore_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs)() + tg["flush"] = time_seg(make_flush_call(state0.clone(), indices, ubufs, m)) + if _HAVE_SHUFFLE and _want("shuffle"): + tg["shuffle_v"] = time_seg( + make_shuffle_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs) + ) + if _HAVE_TCORE and _want("tensor_core"): + tg["tensor_core_v"] = time_seg( + make_tcore_call(q, k, v, a, b, A_log, dt_bias, state0.clone(), indices, scale, DSU, ubufs) + ) + # official triton recurrent verify (dsu=1, writes T·d² states) + if _HAVE_TRITON and _want("triton"): + qt, kt, vt, at, bt, cu = to_triton_varlen(q, k, v, a, b) + tri_inter = torch.empty(N, T, args.HV, args.V, args.K, dtype=torch.float32, device=device) + tri_idx = torch.arange(N, device=device, dtype=torch.int32) + tg["triton_v"] = time_seg( + make_triton_call( + qt, kt, vt, at, bt, cu, A_log, dt_bias, state0.clone(), indices, scale, DSU, tri_inter, tri_idx, T + ) + ) + + r = {"N": N, "T": T, "m": m, "tg": tg} + + def _sum(av, bv): + return tg[av] + tg[bv] if (av in tg and bv in tg) else None + + r["REC_recurrent"] = _sum("recurrent_v", "cmt") + r["REC_recurrent_ws"] = _sum("recurrent_ws_v", "cmt") + r["KVB_shuffle"] = _sum("shuffle_v", "flush") + r["KVB_tensor_core"] = _sum("tensor_core_v", "flush") + r["REC_triton"] = _sum("triton_v", "cmt") + results.append(r) + + # ---- table 1: chain totals + speedups ---- + print(f"\n=== verify-CHAIN total latency (us) + speedup — accept m={args.accept} commit={args.commit} ===") + print(" REC_* = recurrent verify (writes T·d² states) + commit; KVB_* = kvbuffer verify (u-buffer) + flush") + print( + " spd_(recurrent/recurrent_ws/shuffle/tensor_core) = REC_triton (official triton) / (REC_recurrent/REC_recurrent_ws/KVB_shuffle/KVB_tensor_core) -- chain speedup over triton" + ) + hdr = ( + f"{'N':>4} {'T':>3} {'m':>3} | {'REC_recurrent':>7} {'REC_recurrent_ws':>7} {'REC_triton':>7} | {'KVB_shuffle':>11} {'KVB_tensor_core':>9} | " + f"{'spd_recurrent':>7} {'spd_recurrent_ws':>7} {'spd_shuffle':>11} {'spd_tensor_core':>9}" + ) + print(hdr) + print("-" * len(hdr)) + for r in results: + print( + f"{r['N']:>4} {r['T']:>3} {r['m']:>3} | {us(r['REC_recurrent']):>7} {us(r['REC_recurrent_ws']):>7} {us(r['REC_triton']):>7} | " + f"{us(r['KVB_shuffle']):>11} {us(r['KVB_tensor_core']):>9} | " + f"{rat(r['REC_triton'], r['REC_recurrent']):>7} {rat(r['REC_triton'], r['REC_recurrent_ws']):>7} {rat(r['REC_triton'], r['KVB_shuffle']):>11} {rat(r['REC_triton'], r['KVB_tensor_core']):>9}" + ) + + # ---- table 2: per-segment breakdown ---- + print("\n=== per-segment breakdown (us) — verify kernels + shared commit/flush ===") + hdr2 = f"{'N':>4} {'T':>3} | {'recurrent_v':>6} {'recurrent_ws_v':>7} {'triton_v':>6} | {'shuffle_v':>9} {'tensor_core_v':>7} | {'cmt':>5} {'flush':>6}" + print(hdr2) + print("-" * len(hdr2)) + for r in results: + tg = r["tg"] + print( + f"{r['N']:>4} {r['T']:>3} | {us(tg.get('recurrent_v')):>6} {us(tg.get('recurrent_ws_v')):>7} {us(tg.get('triton_v')):>6} | " + f"{us(tg.get('shuffle_v')):>9} {us(tg.get('tensor_core_v')):>7} | " + f"{us(tg.get('cmt')):>5} {us(tg.get('flush')):>6}" + ) + + +if __name__ == "__main__": + main() diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index f0d21e62..0b61b13e 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -17,6 +17,9 @@ __all__ = [ "chunk_kda", "kda_decode", + "kda_decode_mtp", + "kda_decode_mtp_recurrent", + "kda_decode_mtp_recurrent_ws", "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", "kda_prefill_hopper_opt", @@ -27,6 +30,12 @@ "chunk_kda": ("cula.kda.chunk", "chunk_kda"), "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), + "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), + "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), + "kda_decode_mtp_recurrent_ws": ( + "cula.ops.kda.decode.mtp", + "kda_decode_mtp_recurrent_ws", + ), "fused_sigmoid_gating_delta_rule_update": ( "cula.ops.kda.decode.cute", "fused_sigmoid_gating_delta_rule_update", diff --git a/cula/ops/__init__.py b/cula/ops/__init__.py index 8332f620..60c60de0 100644 --- a/cula/ops/__init__.py +++ b/cula/ops/__init__.py @@ -14,12 +14,21 @@ __all__ = [ "kda_decode", + "kda_decode_mtp", + "kda_decode_mtp_recurrent", + "kda_decode_mtp_recurrent_ws", "fused_sigmoid_gating_delta_rule_update", "linear_attention_decode", ] _LAZY = { "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), + "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), + "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), + "kda_decode_mtp_recurrent_ws": ( + "cula.ops.kda.decode.mtp", + "kda_decode_mtp_recurrent_ws", + ), "fused_sigmoid_gating_delta_rule_update": ( "cula.ops.kda.decode.cute", "fused_sigmoid_gating_delta_rule_update", diff --git a/cula/ops/kda/decode/mtp.py b/cula/ops/kda/decode/mtp.py new file mode 100644 index 00000000..7e2f8484 --- /dev/null +++ b/cula/ops/kda/decode/mtp.py @@ -0,0 +1,2163 @@ +"""CuTe DSL KDA MTP decode + +Recurrent KDA MTP verify/decode kernels. ``kda_decode_mtp`` dispatches to the +single-warp ``vk`` (lane=K, Triton-identical K-reduce; production verify variant) and +``kv`` (lane=V) kernels. KDA's decay gate ``g_t in R^K`` is per-K-channel (``beta`` is +a per-(head, token) scalar). State is register-resident across the T tokens +(full-warp-shuffle K-reduce, DECAY-FIRST recurrence). An ``intermediate_states_buffer`` ([N,T,HV,V,K] vk) +snapshots per-token post-states to GMEM for spec-decode rollback; +``disable_state_update`` skips the final write-back. + +Math per token t (decay-first, per-channel g): + g_t = exp(-exp(A_log) * softplus(a_t + dt_bias)) # (K,) per-channel + S <- S * diag(g_t) # step 1 (per channel) + s = S @ k_norm # step 2 (reduce K) + v_new = sigmoid(b_t) * (v_t - s) # step 3 + S += v_new (x) k_norm # step 4 (rank-1, raw k) + o_t = S @ (l2norm(q_t) * scale) # step 5 (reduce K) +""" + +import logging + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from cula.ops.kda.decode.cute import ( + NUM_THREADS, + TILE_K, + _canonicalize_state_layout, + _get_cached_stream, + _normalize_A_log, + _normalize_dt_bias, + _normalize_state_indices, + _normalize_state_source, + _prepare_output_tensor, +) + +logger = logging.getLogger(__name__) + +# vec_size = 4 -> 32 threads/group = a full warp, 4 groups (warps) per block. +VEC_SIZE_MTP = 4 + +_compiled_mtp_recurrent_ws_kernels: dict[tuple, object] = {} + + +def _normalize_mtp_a(a: torch.Tensor, *, N: int, T: int, HV: int, K: int) -> torch.Tensor: + """Normalize `a` to the compile-time dense MTP shape (N, T, HV, K).""" + if a.dim() == 4 and tuple(a.shape) == (N, T, HV, K): + return a + if a.dim() == 3 and tuple(a.shape) == (N, T, HV * K): + return a.view(N, T, HV, K) + raise ValueError(f"Unexpected a shape for MTP dense: {tuple(a.shape)}; expected {(N, T, HV, K)}") + + +# Valid V-tile sizes {8,16,32,64}: each a multiple of NUM_WARPS (4) so V_PER_WARP +_MTP_TILE_V_CHOICES = (8, 16, 32, 64) + + +def _select_mtp_config( + N: int, + HV: int, + V: int, + T: int, + *, + disable_state_update: bool = False, +) -> tuple[int, int, bool]: + work_units = N * HV + + if work_units <= 64: + tile_v, ilp_rows, use_smem_v = 8, 2, False + elif work_units <= 128: + tile_v, ilp_rows, use_smem_v = 16, 4, False + elif work_units <= 448: + if T <= 2: + tile_v, ilp_rows, use_smem_v = 16, 2, False + else: + tile_v, ilp_rows, use_smem_v = 32, 4, False + elif work_units <= 1024: + tile_v, ilp_rows, use_smem_v = 32, 4, False + else: + # Large batches: ilp capped at 4, so (64, 4, True) uniformly. + tile_v, ilp_rows, use_smem_v = 64, 4, True + + tile_v = min(tile_v, V) + while tile_v > _MTP_TILE_V_CHOICES[0] and V % tile_v != 0: + tile_v //= 2 + + # Legality backstop: ilp=4 requires (tile_v//4) % 4 == 0, i.e. tile_v % 16 == 0 + if ilp_rows == 4 and tile_v % 16 != 0: + ilp_rows = 2 + + return tile_v, ilp_rows, use_smem_v + + +def _select_mtp_tile_v(N: int, HV: int, V: int, T: int) -> int: + return _select_mtp_config(N, HV, V, T)[0] + + +@cute.jit +def fma_pair(a1, a2, b1, b2, c1, c2): + # FMA two pairs: (a1*b1+c1, a2*b2+c2). + result1 = a1 * b1 + c1 + result2 = a2 * b2 + c2 + return result1, result2 + + +@cute.kernel +def kda_verify_kernel_mtp_recurrent_ws( + h0_source: cute.Tensor, # [pool_size * HV, V, K] fp32, K-last (VK layout) + intermediate_states: cute.Tensor, # [N*T*HV, V, K] fp32 snapshot cache (or dummy) + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + A_log: cute.Tensor, # [HV] fp32 (per-channel decay) + a: cute.Tensor, # [N, T, HV, K] (per-channel decay input) + dt_bias: cute.Tensor, # [HV, K] (per-channel decay bias) + q: cute.Tensor, # [N, T, H, K] + k: cute.Tensor, # [N, T, H, K] + v: cute.Tensor, # [N, T, HV, V] + b: cute.Tensor, # [N, T, HV] (update-gate logit) + o: cute.Tensor, # [N, T, HV, V] output + h0_indices: cute.Tensor, # [N] int32 (state-pool slot per sequence; <0 = pad) + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + use_smem_v: cutlass.Constexpr[bool], + cache_intermediate_states: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # vec_size=4 -> threads_per_group=32 (full warp), 4 groups (one per warp). + threads_per_group: cutlass.Constexpr[int] = K // vec_size # 32 + num_groups: cutlass.Constexpr[int] = 4 + lane_in_group = lane_id % threads_per_group + group_idx = warp_idx + + batch_idx, _, _ = cute.arch.block_idx() + + # Decode the flat CTA index into (i_n sequence, i_hv value-head, i_v V-tile). + i_v = batch_idx % num_v_tiles + tmp = batch_idx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) # GVA: HV//H value-heads share one q/k head + + cache_idx = h0_indices[i_n] + + # exp(A_log) is per-head, shared across all K channels — hoist once. + r_A_log = cutlass.Float32(A_log[i_hv]) + r_exp_A = cute.exp(r_A_log, fastmath=fast_math) + + # SMEM broadcast buffers (warp 0 -> all warps). sG is [T, K] (per-channel); + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sK = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T,)), 16) + + # use_smem_v (Stage C): preload the v-tile into SMEM + accumulate outputs for a + # coalesced merged writeback. Allocated last/conditionally so off-path offsets stay put. + if cutlass.const_expr(use_smem_v): + sVdata = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, tile_v), stride=(tile_v, 1)), 16) + sOutput = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout((T, tile_v), stride=(tile_v, 1)), 16) + + # Per-lane registers: r_g = this lane's vec_size channels of g; r_h = up to 8 + # V-rows of state (only ilp_rows used), each row spanning 32 lanes over K=128. + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_g = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_h = cute.make_rmem_tensor(cute.make_layout((8, vec_size), stride=(vec_size, 1)), cutlass.Float32) + r_q_bf16 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + r_k_bf16 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + + if cache_idx >= 0: + k_start = lane_in_group * vec_size # this lane's first K channel + rows_per_group: cutlass.Constexpr[int] = tile_v // num_groups + flat_state_idx = cache_idx * HV + i_hv # row in [pool*HV, V, K] + + # ---- Phase 1a: all 4 warps compute the per-K-channel decay gate ---- + g_ch = warp_idx * threads_per_group + lane_in_group + for i_t in cutlass.range_constexpr(T): + x = cutlass.Float32(a[i_n, i_t, i_hv, g_ch]) + cutlass.Float32(dt_bias[i_hv, g_ch]) + if cutlass.const_expr(use_lower_bound): + # safe gate: g = lower_bound * sigmoid(exp(A_log) * x) + sigmoid_ax = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-r_exp_A * x, fastmath=fast_math)) + sG[(i_t, g_ch)] = cute.exp(lower_bound * sigmoid_ax, fastmath=fast_math) + else: + beta_x = softplus_beta * x + exp_beta_x = cute.exp(beta_x, fastmath=fast_math) + softplus_val = (cutlass.Float32(1.0) / softplus_beta) * cute.log( + cutlass.Float32(1.0) + exp_beta_x, fastmath=fast_math + ) + use_softplus = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) + softplus_x = use_softplus * softplus_val + (cutlass.Float32(1.0) - use_softplus) * x + sG[(i_t, g_ch)] = cute.exp(-r_exp_A * softplus_x, fastmath=fast_math) + + # ============ Phase 1b: warp 0 q/k+beta, warps 1-3 state prefetch ============ + if warp_idx == 0: + # Warp 0 computes q/k/g/beta for all T tokens, broadcasts via SMEM. + for i_t in cutlass.range_constexpr(T): + q_tile = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, i_t, i_h, lane_in_group)) + k_tile = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, i_t, i_h, lane_in_group)) + cute.autovec_copy(q_tile, r_q_bf16) + cute.autovec_copy(k_tile, r_k_bf16) + for i in cutlass.range_constexpr(vec_size): + r_q[i] = cutlass.Float32(r_q_bf16[i]) + r_k[i] = cutlass.Float32(r_k_bf16[i]) + + if cutlass.const_expr(use_qk_l2norm): + sum_q = 0.0 + sum_k = 0.0 + for i in cutlass.range_constexpr(vec_size): + sum_q += r_q[i] * r_q[i] + sum_k += r_k[i] * r_k[i] + # Full-warp reduction (32 lanes x vec_size=4 = all 128 K). + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=offset, mask=-1, mask_and_clamp=31) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=offset, mask=-1, mask_and_clamp=31) + inv_norm_q_scaled = cute.rsqrt(sum_q + 1e-6, fastmath=fast_math) * scale + inv_norm_k = cute.rsqrt(sum_k + 1e-6, fastmath=fast_math) + for i in cutlass.range_constexpr(vec_size): + r_q[i] = r_q[i] * inv_norm_q_scaled + r_k[i] = r_k[i] * inv_norm_k + else: + for i in cutlass.range_constexpr(vec_size): + r_q[i] = r_q[i] * scale + + # vec_size=4 -> warp 0's 32 lanes cover all 128 K channels. + for i in cutlass.range_constexpr(vec_size): + sQ[(i_t, k_start + i)] = r_q[i] + sK[(i_t, k_start + i)] = r_k[i] + + # Update gate beta is a per-(head, token) scalar (warp-uniform). + r_b = cutlass.Float32(b[i_n, i_t, i_hv]) + r_beta = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-r_b, fastmath=fast_math)) + sBeta[i_t] = r_beta + + # Preload the v-tile into SMEM: warp 0 covers tile-local cols 0..31, + # warps 1-3 the rest (tidx each col written once). + if cutlass.const_expr(use_smem_v): + if tidx < tile_v: + v_global_idx = i_v * tile_v + tidx + if v_global_idx < V: + sVdata[(i_t, tidx)] = cutlass.Float32(v[i_n, i_t, i_hv, v_global_idx]) + else: + # Warps 1-3: prefetch the first ILP set of state rows into registers, + # overlapping the h-state DRAM latency with warp 0's Phase 1 compute. + v_base_prefetch = i_v * tile_v + group_idx * rows_per_group + if cutlass.const_expr(ilp_rows == 4): + # Prefetch 4 h-state rows (4 independent load streams). + v_pf_d = v_base_prefetch + 3 + if v_pf_d < V: + pf_a = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_base_prefetch, lane_in_group), + ) + pf_b = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_base_prefetch + 1, lane_in_group), + ) + pf_c = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_base_prefetch + 2, lane_in_group), + ) + pf_d = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_base_prefetch + 3, lane_in_group), + ) + cute.autovec_copy(pf_a, cute.slice_(r_h, (0, None))) + cute.autovec_copy(pf_b, cute.slice_(r_h, (1, None))) + cute.autovec_copy(pf_c, cute.slice_(r_h, (2, None))) + cute.autovec_copy(pf_d, cute.slice_(r_h, (3, None))) + elif cutlass.const_expr(ilp_rows == 2): + v_pf_b = v_base_prefetch + 1 + if v_pf_b < V: + pf_a = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_base_prefetch, lane_in_group), + ) + pf_b = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_base_prefetch + 1, lane_in_group), + ) + cute.autovec_copy(pf_a, cute.slice_(r_h, (0, None))) + cute.autovec_copy(pf_b, cute.slice_(r_h, (1, None))) + + # Warps 1-3 cover the tile-local v columns warp 0 can't reach + # (tidx 32..127); same tidx each column written once. + if cutlass.const_expr(use_smem_v): + for i_t in cutlass.range_constexpr(T): + if tidx < tile_v: + v_global_idx = i_v * tile_v + tidx + if v_global_idx < V: + sVdata[(i_t, tidx)] = cutlass.Float32(v[i_n, i_t, i_hv, v_global_idx]) + + # Publish warp 0's SMEM writes (q/k/g/beta + preloaded v) to all warps + # before the recurrence reads them. + cute.arch.barrier() + + # ============ Recurrence: ilp_rows == 2 (process 2 V-rows together) === + if cutlass.const_expr(ilp_rows == 2): + half_rows: cutlass.Constexpr[int] = rows_per_group // 2 + + for row_pair in cutlass.range_constexpr(half_rows): + v_idx_a = i_v * tile_v + group_idx * rows_per_group + row_pair * 2 + v_idx_b = v_idx_a + 1 + + if v_idx_b < V: + # Load state for both rows. Warps 1-3 reuse the Phase-1 + # prefetch on the first pair; everyone else loads in place. + if warp_idx == 0 or row_pair > 0: + h_tile_a = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_a, lane_in_group), + ) + h_tile_b = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_b, lane_in_group), + ) + cute.autovec_copy(h_tile_a, cute.slice_(r_h, (0, None))) + cute.autovec_copy(h_tile_b, cute.slice_(r_h, (1, None))) + + for i_t in cutlass.range_constexpr(T): + # Read warp-0-staged q/k/g for this token (shared by both rows). + sQ_tile = cute.local_tile(sQ, (1, vec_size), (i_t, lane_in_group)) + sK_tile = cute.local_tile(sK, (1, vec_size), (i_t, lane_in_group)) + sG_tile = cute.local_tile(sG, (1, vec_size), (i_t, lane_in_group)) + cute.autovec_copy(sQ_tile, r_q) + cute.autovec_copy(sK_tile, r_k) + cute.autovec_copy(sG_tile, r_g) + r_beta = sBeta[i_t] + + # Step 1: per-channel decay (KDA: r_g[i], not a scalar). + for i in cutlass.range_constexpr(vec_size): + r_h[0, i] = r_h[0, i] * r_g[i] + r_h[1, i] = r_h[1, i] * r_g[i] + + # Step 2: s = (decayed S) @ k_norm (reduce over K). + sum_hk_a = 0.0 + sum_hk_b = 0.0 + for i in cutlass.range_constexpr(vec_size): + sum_hk_a += r_h[0, i] * r_k[i] + sum_hk_b += r_h[1, i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_hk_a += cute.arch.shuffle_sync_bfly(sum_hk_a, offset=offset, mask=-1, mask_and_clamp=31) + sum_hk_b += cute.arch.shuffle_sync_bfly(sum_hk_b, offset=offset, mask=-1, mask_and_clamp=31) + + # Step 3: delta rule. v from SMEM (preloaded) or GMEM. + if cutlass.const_expr(use_smem_v): + v_local_a = v_idx_a - i_v * tile_v + r_v_a = sVdata[(i_t, v_local_a)] + r_v_b = sVdata[(i_t, v_local_a + 1)] + else: + r_v_a = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_a]) + r_v_b = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_b]) + v_new_a = (r_v_a - sum_hk_a) * r_beta + v_new_b = (r_v_b - sum_hk_b) * r_beta + + # Step 4: rank-1 update with raw k (decay already applied). + for i in cutlass.range_constexpr(vec_size): + r_h[0, i] += r_k[i] * v_new_a + r_h[1, i] += r_k[i] * v_new_b + + # Stage D: snapshot post-token state, sequence-indexed + # (flat_idx = i_n*T*HV + i_t*HV + i_hv), race-free before step 5. + if cutlass.const_expr(cache_intermediate_states): + flat_idx = i_n * T * HV + i_t * HV + i_hv + inter_a = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_a, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (0, None)), inter_a) + inter_b = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_b, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (1, None)), inter_b) + + # Step 5: o = S_new @ q_scaled (reduce over K). + sum_hq_a = 0.0 + sum_hq_b = 0.0 + for i in cutlass.range_constexpr(vec_size): + sum_hq_a += r_h[0, i] * r_q[i] + sum_hq_b += r_h[1, i] * r_q[i] + for offset in [16, 8, 4, 2, 1]: + sum_hq_a += cute.arch.shuffle_sync_bfly(sum_hq_a, offset=offset, mask=-1, mask_and_clamp=31) + sum_hq_b += cute.arch.shuffle_sync_bfly(sum_hq_b, offset=offset, mask=-1, mask_and_clamp=31) + + # Reduction result is identical on all lanes -> lane 0 + # writes. To SMEM (merged flush at kernel end) or GMEM. + if lane_in_group == 0: + if cutlass.const_expr(use_smem_v): + vla = v_idx_a - i_v * tile_v + sOutput[(i_t, vla)] = cutlass.BFloat16(sum_hq_a) + sOutput[(i_t, vla + 1)] = cutlass.BFloat16(sum_hq_b) + else: + o[(i_n, i_t, i_hv, v_idx_a)] = cutlass.BFloat16(sum_hq_a) + o[(i_n, i_t, i_hv, v_idx_b)] = cutlass.BFloat16(sum_hq_b) + + # Write final state for both rows back to the pool (once). + if cutlass.const_expr(not disable_state_update): + h_tile_out_a = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_a, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (0, None)), h_tile_out_a) + h_tile_out_b = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_b, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (1, None)), h_tile_out_b) + + # ============ Recurrence: ilp_rows == 4 (process 4 V-rows together) === + # Steps 1+2 fused (decay then h@k) and 4+5 fused (rank-1 then h@q), with + # double accumulators (halve the K-reduce FFMA chain) + packed F32x2 FMA on + # SM100. Per-channel decay r_g[i]/r_g[i+1] loaded from sG. + elif cutlass.const_expr(ilp_rows == 4): + quarter_rows: cutlass.Constexpr[int] = rows_per_group // 4 + + for row_quad in cutlass.range_constexpr(quarter_rows): + v_idx_a = i_v * tile_v + group_idx * rows_per_group + row_quad * 4 + v_idx_b = v_idx_a + 1 + v_idx_c = v_idx_a + 2 + v_idx_d = v_idx_a + 3 + + if v_idx_d < V: + if warp_idx == 0 or row_quad > 0: + h_tile_a = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_a, lane_in_group), + ) + h_tile_b = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_b, lane_in_group), + ) + h_tile_c = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_c, lane_in_group), + ) + h_tile_d = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_d, lane_in_group), + ) + cute.autovec_copy(h_tile_a, cute.slice_(r_h, (0, None))) + cute.autovec_copy(h_tile_b, cute.slice_(r_h, (1, None))) + cute.autovec_copy(h_tile_c, cute.slice_(r_h, (2, None))) + cute.autovec_copy(h_tile_d, cute.slice_(r_h, (3, None))) + + for i_t in cutlass.range_constexpr(T): + # Warp-0-staged q/k/g for this token (shared by all 4 rows). + sQ_tile = cute.local_tile(sQ, (1, vec_size), (i_t, lane_in_group)) + sK_tile = cute.local_tile(sK, (1, vec_size), (i_t, lane_in_group)) + sG_tile = cute.local_tile(sG, (1, vec_size), (i_t, lane_in_group)) + cute.autovec_copy(sQ_tile, r_q) + cute.autovec_copy(sK_tile, r_k) + cute.autovec_copy(sG_tile, r_g) + r_beta = sBeta[i_t] + + # Steps 1+2 fused: per-channel decay then h@k. + sum_hk_a = cutlass.Float32(0.0) + sum_hk_a2 = cutlass.Float32(0.0) + sum_hk_b = cutlass.Float32(0.0) + sum_hk_b2 = cutlass.Float32(0.0) + sum_hk_c = cutlass.Float32(0.0) + sum_hk_c2 = cutlass.Float32(0.0) + sum_hk_d = cutlass.Float32(0.0) + sum_hk_d2 = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(0, vec_size, 2): + # Step 1: per-channel decay (KDA: r_g[i]/r_g[i+1]). + r_h[0, i] = r_h[0, i] * r_g[i] + r_h[0, i + 1] = r_h[0, i + 1] * r_g[i + 1] + r_h[1, i] = r_h[1, i] * r_g[i] + r_h[1, i + 1] = r_h[1, i + 1] * r_g[i + 1] + r_h[2, i] = r_h[2, i] * r_g[i] + r_h[2, i + 1] = r_h[2, i + 1] * r_g[i + 1] + r_h[3, i] = r_h[3, i] * r_g[i] + r_h[3, i + 1] = r_h[3, i + 1] * r_g[i + 1] + # Step 2: h@k, two channels per step (packed on SM100). + if cutlass.const_expr(use_packed_fma): + sum_hk_a, sum_hk_a2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[0, i], r_h[0, i + 1]), + src_b=(r_k[i], r_k[i + 1]), + src_c=(sum_hk_a, sum_hk_a2), + ) + sum_hk_b, sum_hk_b2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[1, i], r_h[1, i + 1]), + src_b=(r_k[i], r_k[i + 1]), + src_c=(sum_hk_b, sum_hk_b2), + ) + sum_hk_c, sum_hk_c2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[2, i], r_h[2, i + 1]), + src_b=(r_k[i], r_k[i + 1]), + src_c=(sum_hk_c, sum_hk_c2), + ) + sum_hk_d, sum_hk_d2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[3, i], r_h[3, i + 1]), + src_b=(r_k[i], r_k[i + 1]), + src_c=(sum_hk_d, sum_hk_d2), + ) + else: + sum_hk_a, sum_hk_a2 = fma_pair( + r_h[0, i], r_h[0, i + 1], r_k[i], r_k[i + 1], sum_hk_a, sum_hk_a2 + ) + sum_hk_b, sum_hk_b2 = fma_pair( + r_h[1, i], r_h[1, i + 1], r_k[i], r_k[i + 1], sum_hk_b, sum_hk_b2 + ) + sum_hk_c, sum_hk_c2 = fma_pair( + r_h[2, i], r_h[2, i + 1], r_k[i], r_k[i + 1], sum_hk_c, sum_hk_c2 + ) + sum_hk_d, sum_hk_d2 = fma_pair( + r_h[3, i], r_h[3, i + 1], r_k[i], r_k[i + 1], sum_hk_d, sum_hk_d2 + ) + sum_hk_a = sum_hk_a + sum_hk_a2 + sum_hk_b = sum_hk_b + sum_hk_b2 + sum_hk_c = sum_hk_c + sum_hk_c2 + sum_hk_d = sum_hk_d + sum_hk_d2 + + # Full-warp reduction for all 4 h@k dot products. + for offset in [16, 8, 4, 2, 1]: + sum_hk_a += cute.arch.shuffle_sync_bfly(sum_hk_a, offset=offset, mask=-1, mask_and_clamp=31) + sum_hk_b += cute.arch.shuffle_sync_bfly(sum_hk_b, offset=offset, mask=-1, mask_and_clamp=31) + sum_hk_c += cute.arch.shuffle_sync_bfly(sum_hk_c, offset=offset, mask=-1, mask_and_clamp=31) + sum_hk_d += cute.arch.shuffle_sync_bfly(sum_hk_d, offset=offset, mask=-1, mask_and_clamp=31) + + # Step 3: delta rule for all 4 rows. v from SMEM or GMEM. + if cutlass.const_expr(use_smem_v): + v_local_a = v_idx_a - i_v * tile_v + r_v_a = sVdata[(i_t, v_local_a)] + r_v_b = sVdata[(i_t, v_local_a + 1)] + r_v_c = sVdata[(i_t, v_local_a + 2)] + r_v_d = sVdata[(i_t, v_local_a + 3)] + else: + r_v_a = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_a]) + r_v_b = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_b]) + r_v_c = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_c]) + r_v_d = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_d]) + v_new_a = (r_v_a - sum_hk_a) * r_beta + v_new_b = (r_v_b - sum_hk_b) * r_beta + v_new_c = (r_v_c - sum_hk_c) * r_beta + v_new_d = (r_v_d - sum_hk_d) * r_beta + + # Steps 4+5 FUSED: rank-1 update with raw k (step 4) then + # h@q (step 5), per row. Double accumulators again. + sum_hq_a = cutlass.Float32(0.0) + sum_hq_a2 = cutlass.Float32(0.0) + sum_hq_b = cutlass.Float32(0.0) + sum_hq_b2 = cutlass.Float32(0.0) + sum_hq_c = cutlass.Float32(0.0) + sum_hq_c2 = cutlass.Float32(0.0) + sum_hq_d = cutlass.Float32(0.0) + sum_hq_d2 = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(0, vec_size, 2): + if cutlass.const_expr(use_packed_fma): + r_h[0, i], r_h[0, i + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(v_new_a, v_new_a), + src_c=(r_h[0, i], r_h[0, i + 1]), + ) + r_h[1, i], r_h[1, i + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(v_new_b, v_new_b), + src_c=(r_h[1, i], r_h[1, i + 1]), + ) + r_h[2, i], r_h[2, i + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(v_new_c, v_new_c), + src_c=(r_h[2, i], r_h[2, i + 1]), + ) + r_h[3, i], r_h[3, i + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(v_new_d, v_new_d), + src_c=(r_h[3, i], r_h[3, i + 1]), + ) + sum_hq_a, sum_hq_a2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[0, i], r_h[0, i + 1]), + src_b=(r_q[i], r_q[i + 1]), + src_c=(sum_hq_a, sum_hq_a2), + ) + sum_hq_b, sum_hq_b2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[1, i], r_h[1, i + 1]), + src_b=(r_q[i], r_q[i + 1]), + src_c=(sum_hq_b, sum_hq_b2), + ) + sum_hq_c, sum_hq_c2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[2, i], r_h[2, i + 1]), + src_b=(r_q[i], r_q[i + 1]), + src_c=(sum_hq_c, sum_hq_c2), + ) + sum_hq_d, sum_hq_d2 = cute.arch.fma_packed_f32x2( + src_a=(r_h[3, i], r_h[3, i + 1]), + src_b=(r_q[i], r_q[i + 1]), + src_c=(sum_hq_d, sum_hq_d2), + ) + else: + r_h[0, i], r_h[0, i + 1] = fma_pair( + r_k[i], r_k[i + 1], v_new_a, v_new_a, r_h[0, i], r_h[0, i + 1] + ) + r_h[1, i], r_h[1, i + 1] = fma_pair( + r_k[i], r_k[i + 1], v_new_b, v_new_b, r_h[1, i], r_h[1, i + 1] + ) + r_h[2, i], r_h[2, i + 1] = fma_pair( + r_k[i], r_k[i + 1], v_new_c, v_new_c, r_h[2, i], r_h[2, i + 1] + ) + r_h[3, i], r_h[3, i + 1] = fma_pair( + r_k[i], r_k[i + 1], v_new_d, v_new_d, r_h[3, i], r_h[3, i + 1] + ) + sum_hq_a, sum_hq_a2 = fma_pair( + r_h[0, i], r_h[0, i + 1], r_q[i], r_q[i + 1], sum_hq_a, sum_hq_a2 + ) + sum_hq_b, sum_hq_b2 = fma_pair( + r_h[1, i], r_h[1, i + 1], r_q[i], r_q[i + 1], sum_hq_b, sum_hq_b2 + ) + sum_hq_c, sum_hq_c2 = fma_pair( + r_h[2, i], r_h[2, i + 1], r_q[i], r_q[i + 1], sum_hq_c, sum_hq_c2 + ) + sum_hq_d, sum_hq_d2 = fma_pair( + r_h[3, i], r_h[3, i + 1], r_q[i], r_q[i + 1], sum_hq_d, sum_hq_d2 + ) + sum_hq_a = sum_hq_a + sum_hq_a2 + sum_hq_b = sum_hq_b + sum_hq_b2 + sum_hq_c = sum_hq_c + sum_hq_c2 + sum_hq_d = sum_hq_d + sum_hq_d2 + + # Full-warp reduction for all 4 h@q dot products. + for offset in [16, 8, 4, 2, 1]: + sum_hq_a += cute.arch.shuffle_sync_bfly(sum_hq_a, offset=offset, mask=-1, mask_and_clamp=31) + sum_hq_b += cute.arch.shuffle_sync_bfly(sum_hq_b, offset=offset, mask=-1, mask_and_clamp=31) + sum_hq_c += cute.arch.shuffle_sync_bfly(sum_hq_c, offset=offset, mask=-1, mask_and_clamp=31) + sum_hq_d += cute.arch.shuffle_sync_bfly(sum_hq_d, offset=offset, mask=-1, mask_and_clamp=31) + + # Reduction result is identical on all lanes -> lane 0 + # writes. To SMEM (merged flush at kernel end) or GMEM. + if lane_in_group == 0: + if cutlass.const_expr(use_smem_v): + vla = v_idx_a - i_v * tile_v + sOutput[(i_t, vla)] = cutlass.BFloat16(sum_hq_a) + sOutput[(i_t, vla + 1)] = cutlass.BFloat16(sum_hq_b) + sOutput[(i_t, vla + 2)] = cutlass.BFloat16(sum_hq_c) + sOutput[(i_t, vla + 3)] = cutlass.BFloat16(sum_hq_d) + else: + o[(i_n, i_t, i_hv, v_idx_a)] = cutlass.BFloat16(sum_hq_a) + o[(i_n, i_t, i_hv, v_idx_b)] = cutlass.BFloat16(sum_hq_b) + o[(i_n, i_t, i_hv, v_idx_c)] = cutlass.BFloat16(sum_hq_c) + o[(i_n, i_t, i_hv, v_idx_d)] = cutlass.BFloat16(sum_hq_d) + + # Stage D: snapshot post-token state (sequence-indexed), + # last here since fused 4+5 means r_h is final only now. + if cutlass.const_expr(cache_intermediate_states): + flat_idx = i_n * T * HV + i_t * HV + i_hv + inter_a = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_a, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (0, None)), inter_a) + inter_b = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_b, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (1, None)), inter_b) + inter_c = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_c, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (2, None)), inter_c) + inter_d = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_d, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (3, None)), inter_d) + + # Write final state for all 4 rows back to the pool (once). + if cutlass.const_expr(not disable_state_update): + h_tile_out_a = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_a, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (0, None)), h_tile_out_a) + h_tile_out_b = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_b, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (1, None)), h_tile_out_b) + h_tile_out_c = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_c, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (2, None)), h_tile_out_c) + h_tile_out_d = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_d, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (3, None)), h_tile_out_d) + + # ============ Merged output writeback (use_smem_v only) ============ + # Barrier publishes all groups' disjoint lane-0 sOutput writes, then all 128 + # threads flush sOutput -> o (one tile-local column each, all T tokens) so the + # GMEM writes coalesce. Inside `cache_idx >= 0` so the barrier never deadlocks. + if cutlass.const_expr(use_smem_v): + cute.arch.barrier() + v_tile_base = i_v * tile_v + for t_idx in cutlass.range_constexpr(T): + if tidx < tile_v: + v_global = v_tile_base + tidx + if v_global < V: + o[(i_n, t_idx, i_hv, v_global)] = sOutput[(t_idx, tidx)] + + +@cute.jit +def run_kda_verify_kernel_mtp_recurrent_ws( + h0_source: cute.Tensor, + intermediate_states: cute.Tensor, + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + vec_size: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + use_smem_v: cutlass.Constexpr[bool], + cache_intermediate_states: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], + stream: cuda.CUstream, +): + """Host-side launcher: grid = N * HV * num_v_tiles, block = 128 (4 warps).""" + n_indices = h0_indices.layout.shape[0] + v_dim = h0_source.layout.shape[1] + k_dim = h0_source.layout.shape[2] + + num_v_tiles = cute.ceil_div(v_dim, tile_v) + grid_size = n_indices * HV * num_v_tiles + + smem_bytes = ( + 4 * T * (k_dim + 8) # sQ + + 4 * T * (k_dim + 8) # sK + + 4 * T * (k_dim + 8) # sG (per-channel) + + 4 * T # sBeta + + 128 # alignment slack + ) + if cutlass.const_expr(use_smem_v): + smem_bytes += 4 * T * tile_v # sVdata (fp32) + smem_bytes += 2 * T * tile_v # sOutput (bf16) + + kda_verify_kernel_mtp_recurrent_ws( + h0_source, + intermediate_states, + vec_size, + num_v_tiles, + tile_v, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + h0_indices, + softplus_beta, + softplus_threshold, + scale, + HV, + T, + H, + K, + V, + use_qk_l2norm, + disable_state_update, + ilp_rows, + use_packed_fma, + use_smem_v, + cache_intermediate_states, + fast_math, + use_lower_bound, + lower_bound, + ).launch( + grid=(grid_size, 1, 1), + block=[NUM_THREADS, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +def _dlp_qkv(_t, _dyn): + # dyn-stride: K-contiguous strided view -> dynamic-layout tensor (no copy); + # contiguous input keeps the compact (byte-identical) descriptor. + if _dyn: + return from_dlpack(_t, assumed_align=16).mark_layout_dynamic(leading_dim=3) + return from_dlpack(_t, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=_t.dim_order()) + + +def _get_compiled_mtp_recurrent_ws_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + scale, + use_qk_l2norm, + disable_state_update, + softplus_beta, + softplus_threshold, + tile_v, + ilp_rows, + use_packed_fma, + use_smem_v, + cache_intermediate_states, + opt_level=3, + fast_math=True, + use_lower_bound=False, + lower_bound=0.0, + dyn_stride=False, +): + """Get or lazily compile the warp-spec MTP kernel for one shape/config. + + ``opt_level`` (``--opt-level``) and ``fast_math`` are part of the cache key. + """ + key = ( + T, + H, + HV, + K, + V, + scale, + use_qk_l2norm, + disable_state_update, + softplus_beta, + softplus_threshold, + tile_v, + ilp_rows, + use_packed_fma, + use_smem_v, + cache_intermediate_states, + opt_level, + fast_math, + use_lower_bound, + lower_bound, + dyn_stride, + ) + if key in _compiled_mtp_recurrent_ws_kernels: + return _compiled_mtp_recurrent_ws_kernels[key] + + q = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, T, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, T, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + # Warp-spec kernel uses the flat 3D state view [pool*HV, V, K] (VK layout). + h0_source = torch.zeros(pool_size * HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + if cache_intermediate_states: + intermediate_states = torch.zeros(N * T * HV, V, K, dtype=torch.float32, device="cuda") + else: + intermediate_states = torch.zeros(1, 1, 1, dtype=torch.float32, device="cuda") + + # dynamic-N (flashinfer-aligned): batch + pool axes dynamic -> one cubin per shape config. + q_tensor = _dlp_qkv(q, dyn_stride) + k_tensor = _dlp_qkv(k, dyn_stride) + v_tensor = _dlp_qkv(v, dyn_stride) + a_tensor = from_dlpack(a, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()) + b_tensor = from_dlpack(b, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=b.dim_order()) + A_log_tensor = from_dlpack(A_log, assumed_align=16) + dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) + h0_source_tensor = from_dlpack(h0_source, assumed_align=16).mark_compact_shape_dynamic( + mode=0, stride_order=h0_source.dim_order() + ) + h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16).mark_layout_dynamic() + o_tensor = from_dlpack(o, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=o.dim_order()) + intermediate_states_tensor = from_dlpack(intermediate_states, assumed_align=16) + if cache_intermediate_states: + intermediate_states_tensor = intermediate_states_tensor.mark_compact_shape_dynamic( + mode=0, stride_order=intermediate_states.dim_order() + ) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + compiled_kernel = cute.compile( + run_kda_verify_kernel_mtp_recurrent_ws, + h0_source_tensor, + intermediate_states_tensor, + A_log_tensor, + a_tensor, + dt_bias_tensor, + q_tensor, + k_tensor, + v_tensor, + b_tensor, + o_tensor, + h0_indices_tensor, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + tile_v=tile_v, + vec_size=VEC_SIZE_MTP, + use_qk_l2norm=use_qk_l2norm, + disable_state_update=disable_state_update, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + use_smem_v=use_smem_v, + cache_intermediate_states=cache_intermediate_states, + fast_math=fast_math, + use_lower_bound=use_lower_bound, + lower_bound=lower_bound, + stream=stream, + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + + _compiled_mtp_recurrent_ws_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA MTP warp-spec kernel compiled: " + f"N={N}, T={T}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, " + f"tile_v={tile_v}, ilp_rows={ilp_rows}, use_packed_fma={use_packed_fma}, " + f"use_smem_v={use_smem_v}, cache_intermediate_states={cache_intermediate_states}" + ) + return compiled_kernel + + +def kda_decode_mtp_recurrent_ws( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + state_layout: str = "vk", + tile_v: int | None = None, + ilp_rows: int | None = None, + disable_state_update: bool = False, + use_packed_fma: bool | None = None, + use_smem_v: bool | None = None, + intermediate_states_buffer: torch.Tensor | None = None, + lower_bound: float | None = None, +) -> torch.Tensor: + N, T, H, K = q.shape + HV = v.shape[2] + V = v.shape[3] + + if scale is None: + scale = K**-0.5 + else: + assert scale > 0, f"scale must be positive, got {scale}" + + assert K == TILE_K, f"KDA MTP (ws) kernel requires K={TILE_K}, got {K}" + + if tile_v is None or ilp_rows is None or use_smem_v is None: + sel_tile_v, sel_ilp_rows, sel_use_smem_v = _select_mtp_config(N, HV, V, T, disable_state_update=disable_state_update) + write_bound_verify = intermediate_states_buffer is not None and N >= 8 and V % 16 == 0 + if tile_v is None: + tile_v = 16 if write_bound_verify else sel_tile_v + if ilp_rows is None: + if write_bound_verify and T > 4 and N * HV >= _WS_WORK_UNIT_THRESHOLD: + ilp_rows = 2 + else: + ilp_rows = sel_ilp_rows + if ilp_rows == 4 and tile_v % 16 != 0: + ilp_rows = 2 + if use_smem_v is None: + use_smem_v = sel_use_smem_v + + if ilp_rows not in (2, 4): + raise NotImplementedError(f"kda_decode_mtp_recurrent_ws implements ilp_rows in {{2, 4}}, got {ilp_rows}") + + # packed F32x2 FMA exists only on SM100+ (Blackwell） + if use_packed_fma is None: + major, _ = torch.cuda.get_device_capability(q.device) + use_packed_fma = major >= 10 + # The packed path only exists in the ilp=4 kernel branch; ilp=2 is scalar. + if ilp_rows != 4: + use_packed_fma = False + + state_layout = _canonicalize_state_layout(state_layout) + if state_layout != "vk": + raise NotImplementedError(f"kda_decode_mtp_recurrent_ws only supports state_layout='vk'; got {state_layout!r}") + + assert tile_v % 4 == 0, f"KDA MTP (ws) requires tile_v % 4 == 0, got tile_v={tile_v}" + assert V % tile_v == 0, f"KDA MTP (ws) requires V % tile_v == 0, got V={V}, tile_v={tile_v}" + + rows_per_group = tile_v // 4 + assert rows_per_group % ilp_rows == 0, ( + f"ilp_rows={ilp_rows} requires (tile_v//4) divisible by {ilp_rows}, got tile_v={tile_v} (tile_v//4={rows_per_group})" + ) + + # State is token-independent: reuse the single-token normalizer/validator. + h0_source, pool_size, state_layout_is_kv = _normalize_state_source( + initial_state_source, + N=N, + HV=HV, + K=K, + V=V, + device=q.device, + state_layout=state_layout, + ) + assert not state_layout_is_kv # guaranteed by the vk-only guard above + + a = _normalize_mtp_a(a, N=N, T=T, HV=HV, K=K) + if b.dim() != 3 or tuple(b.shape) != (N, T, HV): + raise ValueError(f"Unexpected b shape for MTP dense: {tuple(b.shape)}; expected {(N, T, HV)}") + + o = _prepare_output_tensor(q, out, (N, T, HV, V)) + + _dyn_ws = ( + not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()) + and q.stride(-1) == 1 + and k.stride(-1) == 1 + and v.stride(-1) == 1 + ) + q = q if (_dyn_ws or q.is_contiguous()) else q.contiguous() + k = k if (_dyn_ws or k.is_contiguous()) else k.contiguous() + v = v if (_dyn_ws or v.is_contiguous()) else v.contiguous() + a = a if a.is_contiguous() else a.contiguous() + b = b if b.is_contiguous() else b.contiguous() + + A_log = _normalize_A_log(A_log, HV) + dt_bias = _normalize_dt_bias(dt_bias, HV, K) + initial_state_indices = _normalize_state_indices(initial_state_indices, N=N, pool_size=pool_size, device=q.device) + + # Flatten the VK state pool [pool, HV, V, K] -> [pool*HV, V, K] + h0_source_flat = h0_source.view(pool_size * HV, V, K) + + # Stage D: resolve the snapshot cache. + cache_intermediate_states = intermediate_states_buffer is not None + if cache_intermediate_states: + if intermediate_states_buffer.dtype != torch.float32: + raise ValueError(f"intermediate_states_buffer must be float32, got {intermediate_states_buffer.dtype}") + expected_buf_shape = (N, T, HV, V, K) + if tuple(intermediate_states_buffer.shape) != expected_buf_shape: + raise ValueError( + f"intermediate_states_buffer shape {tuple(intermediate_states_buffer.shape)} " + f"!= expected {expected_buf_shape} ([N, T, HV, V, K] vk / K-last)" + ) + intermediate_states_flat = intermediate_states_buffer.view(N * T * HV, V, K) + else: + intermediate_states_flat = torch.empty(1, 1, 1, dtype=torch.float32, device=q.device) + + stream = _get_cached_stream(q.device) + + compiled_kernel = _get_compiled_mtp_recurrent_ws_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + disable_state_update=disable_state_update, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + tile_v=tile_v, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + use_smem_v=use_smem_v, + cache_intermediate_states=cache_intermediate_states, + use_lower_bound=lower_bound is not None, + lower_bound=(0.0 if lower_bound is None else float(lower_bound)), + dyn_stride=_dyn_ws, + ) + + compiled_kernel( + h0_source_flat, + intermediate_states_flat, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + initial_state_indices, + stream, + ) + + return o + + +# ============================================================================ +# recurrent kernel (1-warp/program):kv layout(lane=V)+ vk layout(lane=K) +# ============================================================================ + + +WARP_BV = 32 +VEC_SIZE = 4 + +_compiled_mtp_recurrent_kernels: dict[tuple, object] = {} + + +@cute.kernel +def kda_mtp_recurrent_kernel( + h0_source: cute.Tensor, # [pool*HV, K, V] fp32 (kv, V-last) + A_log: cute.Tensor, # [HV] fp32 + a: cute.Tensor, # [N, T, HV, K] + dt_bias: cute.Tensor, # [HV, K] + q: cute.Tensor, # [N, T, H, K] + k: cute.Tensor, # [N, T, H, K] + v: cute.Tensor, # [N, T, HV, V] + b: cute.Tensor, # [N, T, HV] + o: cute.Tensor, # [N, T, HV, V] + h0_indices: cute.Tensor, # [N] int32 + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + k_split: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + lane = tidx + + bidx, _, _ = cute.arch.block_idx() + i_v = bidx % num_v_tiles # flat CTA -> (i_n, i_hv, i_v V-block) + tmp = bidx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv]), fastmath=fast_math) # per-head, shared across T + + # SMEM-broadcast q/k/g (shared across V-cols on K dim); XOR swizzle staggers k_split segments across banks. + smem_k = K + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout((smem_k,), stride=(1,)), 16) + sK = smem.allocate_tensor(cutlass.Float32, cute.make_layout((smem_k,), stride=(1,)), 16) + sG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((smem_k,), stride=(1,)), 16) + + # k_split lanes split one V-col's K (each holds k_per_lane), butterfly-merged after reduce. + k_per_lane = K // k_split + v_local = lane % BV + k_part = lane // BV + k_off = k_part * k_per_lane + + r_h = cute.make_rmem_tensor(cute.make_layout((k_per_lane,), stride=(1,)), cutlass.Float32) + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_q_bf16 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + r_k_bf16 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + + v_global = i_v * BV + v_local # global V-col this lane serves + k_start = lane * vec_size # prep: full warp, 32 lanes x 4 = all 128 K + + # constexpr k_split decisions hoisted to top level so they stay python + # constants inside the cache_idx>=0 block (else reboxed to Int32 -> error). + ks_single = cutlass.const_expr(k_split == 1) + ks_log2 = cutlass.const_expr(k_split.bit_length() - 1) + if cache_idx >= 0: + flat_state_idx = cache_idx * HV + i_hv + for j in cutlass.range_constexpr(k_per_lane): + r_h[j] = cutlass.Float32(h0_source[flat_state_idx, k_off + j, v_global]) + + for i_t in cutlass.range_constexpr(T): + q_tile = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, i_t, i_h, lane)) + k_tile = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, i_t, i_h, lane)) + cute.autovec_copy(q_tile, r_q_bf16) + cute.autovec_copy(k_tile, r_k_bf16) + for i in cutlass.range_constexpr(vec_size): + r_q[i] = cutlass.Float32(r_q_bf16[i]) + r_k[i] = cutlass.Float32(r_k_bf16[i]) + + if cutlass.const_expr(use_qk_l2norm): + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(vec_size): + sum_q += r_q[i] * r_q[i] + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=offset, mask=-1, mask_and_clamp=31) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=offset, mask=-1, mask_and_clamp=31) + inv_q = cute.rsqrt(sum_q + 1e-6, fastmath=fast_math) * scale + inv_k = cute.rsqrt(sum_k + 1e-6, fastmath=fast_math) + for i in cutlass.range_constexpr(vec_size): + r_q[i] = r_q[i] * inv_q + r_k[i] = r_k[i] * inv_k + else: + for i in cutlass.range_constexpr(vec_size): + r_q[i] = r_q[i] * scale + + for i in cutlass.range_constexpr(vec_size): + kk = k_start + i + sw = kk ^ (kk // k_per_lane) # XOR swizzle SMEM write addr (a/dt_bias read GMEM with raw kk) + x = cutlass.Float32(a[i_n, i_t, i_hv, kk]) + cutlass.Float32(dt_bias[i_hv, kk]) + if cutlass.const_expr(use_lower_bound): + # safe gate: g = lower_bound * sigmoid(exp(A_log) * x) + sigmoid_ax = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-r_exp_A * x, fastmath=fast_math)) + sG[sw] = cute.exp(lower_bound * sigmoid_ax, fastmath=fast_math) + else: + beta_x = softplus_beta * x + exp_bx = cute.exp(beta_x, fastmath=fast_math) + sp_val = (cutlass.Float32(1.0) / softplus_beta) * cute.log( + cutlass.Float32(1.0) + exp_bx, fastmath=fast_math + ) + use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) + sp_x = use_sp * sp_val + (cutlass.Float32(1.0) - use_sp) * x + sG[sw] = cute.exp(-r_exp_A * sp_x, fastmath=fast_math) + sQ[sw] = r_q[i] + sK[sw] = r_k[i] + + r_beta = cutlass.Float32(1.0) / ( + cutlass.Float32(1.0) + cute.exp(-cutlass.Float32(b[i_n, i_t, i_hv]), fastmath=fast_math) + ) + + cute.arch.barrier() # publish prep's SMEM writes before recurrence reads + + r_v = cutlass.Float32(v[i_n, i_t, i_hv, v_global]) + # fused decay + s partial. + s = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(k_per_lane): + sw = j if ks_single else (k_off + j) ^ k_part # XOR swizzle read addr = swz(k_off+j) + r_h[j] = r_h[j] * sG[sw] + s += r_h[j] * sK[sw] + for st in cutlass.range_constexpr(ks_log2): + s += cute.arch.shuffle_sync_bfly(s, offset=BV << st, mask=-1, mask_and_clamp=31) + v_new = (r_v - s) * r_beta + o_val = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(k_per_lane): + sw = j if ks_single else (k_off + j) ^ k_part # XOR swizzle read addr + r_h[j] = r_h[j] + sK[sw] * v_new + o_val += r_h[j] * sQ[sw] + for st in cutlass.range_constexpr(ks_log2): + o_val += cute.arch.shuffle_sync_bfly(o_val, offset=BV << st, mask=-1, mask_and_clamp=31) + o[(i_n, i_t, i_hv, v_global)] = cutlass.BFloat16(o_val) + + cute.arch.barrier() + + if cutlass.const_expr(not disable_state_update): + flat_state_idx = cache_idx * HV + i_hv + for j in cutlass.range_constexpr(k_per_lane): + h0_source[(flat_state_idx, k_off + j, v_global)] = r_h[j] + + +@cute.jit +def run_kda_mtp_recurrent_kernel( + h0_source: cute.Tensor, + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + vec_size: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + k_split: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], + stream: cuda.CUstream, +): + n_indices = h0_indices.layout.shape[0] + num_v_tiles = cute.ceil_div(V, BV) + grid_size = n_indices * HV * num_v_tiles + + smem_bytes = 3 * K * 4 + 256 # sQ + sK + sG + + kda_mtp_recurrent_kernel( + h0_source, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + h0_indices, + vec_size, + num_v_tiles, + BV, + k_split, + softplus_beta, + softplus_threshold, + scale, + HV, + T, + H, + K, + V, + use_qk_l2norm, + disable_state_update, + fast_math, + use_lower_bound, + lower_bound, + ).launch( + grid=(grid_size, 1, 1), + block=[32, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +def _get_compiled_mtp_recurrent_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + BV, + k_split, + scale, + use_qk_l2norm, + disable_state_update, + softplus_beta, + softplus_threshold, + opt_level=3, + fast_math=True, + use_lower_bound=False, + lower_bound=0.0, +): + key = ( + T, + H, + HV, + K, + V, + BV, + k_split, + scale, + use_qk_l2norm, + disable_state_update, + softplus_beta, + softplus_threshold, + opt_level, + fast_math, + use_lower_bound, + lower_bound, + ) + if key in _compiled_mtp_recurrent_kernels: + return _compiled_mtp_recurrent_kernels[key] + + q = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, T, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, T, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + h0_source = torch.zeros(pool_size * HV, K, V, dtype=torch.float32, device="cuda") # kv + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + + # dynamic-N (flashinfer-aligned): batch + pool axes dynamic. + q_t = from_dlpack(q, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=q.dim_order()) + k_t = from_dlpack(k, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=k.dim_order()) + v_t = from_dlpack(v, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=v.dim_order()) + a_t = from_dlpack(a, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()) + b_t = from_dlpack(b, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=b.dim_order()) + o_t = from_dlpack(o, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=o.dim_order()) + A_log_t = from_dlpack(A_log, assumed_align=16) + dt_bias_t = from_dlpack(dt_bias, assumed_align=16) + h0_source_t = from_dlpack(h0_source, assumed_align=16).mark_compact_shape_dynamic( + mode=0, stride_order=h0_source.dim_order() + ) + h0_indices_t = from_dlpack(h0_indices, assumed_align=16).mark_layout_dynamic() + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + compiled_kernel = cute.compile( + run_kda_mtp_recurrent_kernel, + h0_source_t, + A_log_t, + a_t, + dt_bias_t, + q_t, + k_t, + v_t, + b_t, + o_t, + h0_indices_t, + vec_size=VEC_SIZE, + BV=BV, + k_split=k_split, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + use_qk_l2norm=use_qk_l2norm, + disable_state_update=disable_state_update, + fast_math=fast_math, + use_lower_bound=use_lower_bound, + lower_bound=lower_bound, + stream=stream, + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + + _compiled_mtp_recurrent_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA MTP small-batch kernel compiled: " + f"N={N}, T={T}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, BV={BV}, " + f"k_split={k_split}, opt_level={opt_level}, fast_math={fast_math}" + ) + return compiled_kernel + + +_KV_CTAS_PER_SM = {1: 8, 2: 12, 4: 16} + + +def _select_k_split(work_units, V, num_sms): + waves1 = work_units * (V // 32) / (num_sms * _KV_CTAS_PER_SM[1]) + for ks, thresh in ((4, 0.3), (2, 0.6)): + vcols = 32 // ks + if V % vcols == 0 and waves1 < thresh: + return ks + return 1 + + +def kda_decode_mtp_recurrent( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + disable_state_update: bool = False, + variant: str = "kv", + bv: int = WARP_BV, + k_split: int = 1, + opt_level: int = 3, + fast_math: bool = True, + intermediate_states_buffer: torch.Tensor | None = None, + lower_bound: float | None = None, +) -> torch.Tensor: + assert variant in ("kv", "vk"), f"variant only supports 'kv'/'vk',got {variant!r}" + N, T, H, K = q.shape + HV = v.shape[2] + V = v.shape[3] + + if scale is None: + scale = K**-0.5 + else: + assert scale > 0, f"scale must be positive, got {scale}" + + assert K == TILE_K, f"KDA MTP (recurrent) requires K={TILE_K}, got {K}" + assert K % VEC_SIZE == 0 and K // VEC_SIZE == 32, f"recurrent assumes K//vec_size==32, got K={K}, vec_size={VEC_SIZE}" + + if variant == "kv": + state_layout = "kv" + assert bv == WARP_BV, f"recurrent(kv) supports 1 warp,bv must be {WARP_BV},got {bv}" + if k_split <= 0: + num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count + k_split = _select_k_split(N * HV, V, num_sms) + assert k_split in (1, 2, 4), f"k_split only supports 1/2/4 or <=0(auto),got {k_split}" + assert bv % k_split == 0 and K % k_split == 0, ( + f"requires bv%k_split==0 and K%k_split==0, got bv={bv}, K={K}, k_split={k_split}" + ) + vcols = bv // k_split + assert V % vcols == 0, f"recurrent(kv) requires V % (bv//k_split) == 0, got V={V}, vcols={vcols}" + else: # vk + state_layout = "vk" + if bv <= 0: + num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count + bv = _select_vk_bv(N * HV, V, num_sms) + assert bv in (8, 16, 32), f"vk bv only supports 8/16/32 or <=0(auto),got {bv}" + assert V % bv == 0, f"vk requires V % bv == 0, got V={V}, bv={bv}" + + h0_source, pool_size, _ = _normalize_state_source( + initial_state_source, + N=N, + HV=HV, + K=K, + V=V, + device=q.device, + state_layout=state_layout, + ) + + a = _normalize_mtp_a(a, N=N, T=T, HV=HV, K=K) + if b.dim() != 3 or tuple(b.shape) != (N, T, HV): + raise ValueError(f"Unexpected b shape for MTP dense: {tuple(b.shape)}; expected {(N, T, HV)}") + + o = _prepare_output_tensor(q, out, (N, T, HV, V)) + + # dyn-stride (vk only): keep K-contiguous strided q/k/v views as-is and + # compile the dynamic-layout kernel variant instead of copying. Auto: + # contiguous inputs keep the compact (byte-identical) kernel. + _dyn_vk = ( + variant == "vk" + and not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()) + and q.stride(-1) == 1 + and k.stride(-1) == 1 + and v.stride(-1) == 1 + ) + q = q if (_dyn_vk or q.is_contiguous()) else q.contiguous() + k = k if (_dyn_vk or k.is_contiguous()) else k.contiguous() + v = v if (_dyn_vk or v.is_contiguous()) else v.contiguous() + a = a if a.is_contiguous() else a.contiguous() + b = b if b.is_contiguous() else b.contiguous() + + A_log = _normalize_A_log(A_log, HV) + dt_bias = _normalize_dt_bias(dt_bias, HV, K) + initial_state_indices = _normalize_state_indices(initial_state_indices, N=N, pool_size=pool_size, device=q.device) + + stream = _get_cached_stream(q.device) + + cache_intermediate_states = intermediate_states_buffer is not None + if cache_intermediate_states: + if variant != "vk": + raise NotImplementedError("intermediate_states_buffer only supported for variant='vk'") + if intermediate_states_buffer.dtype != torch.float32: + raise ValueError(f"intermediate_states_buffer must be float32, got {intermediate_states_buffer.dtype}") + if tuple(intermediate_states_buffer.shape) != (N, T, HV, V, K): + raise ValueError( + f"intermediate_states_buffer shape {tuple(intermediate_states_buffer.shape)} != expected {(N, T, HV, V, K)} ([N,T,HV,V,K] vk)" + ) + intermediate_states_flat = intermediate_states_buffer.view(N * T * HV, V, K) + else: + intermediate_states_flat = torch.empty(1, 1, 1, dtype=torch.float32, device=q.device) + + if variant == "kv": + h0_source_flat = h0_source.view(pool_size * HV, K, V) # kv + compiled_kernel = _get_compiled_mtp_recurrent_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + vcols, + k_split, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + disable_state_update=disable_state_update, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + opt_level=opt_level, + fast_math=fast_math, + use_lower_bound=lower_bound is not None, + lower_bound=(0.0 if lower_bound is None else float(lower_bound)), + ) + else: # vk + h0_source_flat = h0_source.view(pool_size * HV, V, K) # vk + compiled_kernel = _get_compiled_mtp_vk_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + bv, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + disable_state_update=disable_state_update, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + opt_level=opt_level, + fast_math=fast_math, + cache_intermediate_states=cache_intermediate_states, + use_lower_bound=lower_bound is not None, + lower_bound=(0.0 if lower_bound is None else float(lower_bound)), + dyn_stride=_dyn_vk, + ) + + if variant == "vk": + compiled_kernel( + h0_source_flat, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + intermediate_states_flat, + initial_state_indices, + stream, + ) + else: + compiled_kernel( + h0_source_flat, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + initial_state_indices, + stream, + ) + + return o + + +@cute.kernel +def kda_mtp_recurrent_vk_kernel( + h0_source: cute.Tensor, # [pool*HV, V, K] fp32 (vk) + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + intermediate_states: cute.Tensor, + h0_indices: cute.Tensor, + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + cache_intermediate_states: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + lane = tidx # 1 warp = 32 lanes + + bidx, _, _ = cute.arch.block_idx() + i_v = bidx % num_v_tiles + tmp = bidx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv]), fastmath=fast_math) + + # lane t holds vec_size contiguous K (K[4t:4t+4]) x all BV V-cols; r_h[vv*vec_size+c]=state[i_v*BV+vv, vec_size*lane+c]. + r_h = cute.make_rmem_tensor(cute.make_layout((BV * vec_size,), stride=(1,)), cutlass.Float32) + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_g = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_vbf = [ + cute.make_rmem_tensor(cute.make_layout((BV,), stride=(1,)), cutlass.BFloat16) for _ in range(2) + ] # v: bf16 double-buffer + r_red = cute.make_rmem_tensor( + cute.make_layout((BV,), stride=(1,)), cutlass.Float32 + ) # ILP: BV reduce partials, batched butterfly + r_gx = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) # gate: x=a+dtb + r_gexp = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) # gate: exp(beta_x) + r_h4 = cute.make_rmem_tensor( + cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32 + ) # float4 temp buffer (state load/store) + # ===== 2-stage software-pipeline double-buffer: prefetch token t+1's q/k/a/b while computing token t ===== + r_qbf = [cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) for _ in range(2)] + r_kbf = [cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) for _ in range(2)] + r_abf = [cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) for _ in range(2)] + r_bbf = [cute.make_rmem_tensor(cute.make_layout((1,), stride=(1,)), cutlass.Float32) for _ in range(2)] + r_dtb = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) # dt_bias + + # ===== state load (contiguous + float4: lane t takes K[4t:4t+4] ===== + if cache_idx >= 0: + flat_state_idx = cache_idx * HV + i_hv + for vv in cutlass.range_constexpr(BV): + v_global = i_v * BV + vv + # local_tile 3rd coord = lane, tile=vec_size -> contiguous K -> autovec float4 + h_tile = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_global, lane)) + cute.autovec_copy(h_tile, r_h4) + for c in cutlass.range_constexpr(vec_size): + r_h[vv * vec_size + c] = r_h4[c] + + for c in cutlass.range_constexpr(vec_size): # dt_bias loaded once outside loop (contiguous K[4t:4t+4]) + r_dtb[c] = cutlass.Float32(dt_bias[i_hv, vec_size * lane + c]) + + # prefetch token 0's q/k/a/b into stage 0 (pipeline fill). + q_t0 = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, 0, i_h, lane)) + k_t0 = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, 0, i_h, lane)) + cute.autovec_copy(q_t0, r_qbf[0]) + cute.autovec_copy(k_t0, r_kbf[0]) + a_t0 = cute.local_tile(a, (1, 1, 1, vec_size), (i_n, 0, i_hv, lane)) + cute.autovec_copy(a_t0, r_abf[0]) + v_t0 = cute.local_tile(v, (1, 1, 1, BV), (i_n, 0, i_hv, i_v)) + cute.autovec_copy(v_t0, r_vbf[0]) + r_bbf[0][0] = cutlass.Float32(b[i_n, 0, i_hv]) + + for i_t in cutlass.range_constexpr(T): + cur = i_t % 2 + # ===== prefetch t+1's q/k/a/b ===== + if cutlass.const_expr(i_t + 1 < T): + nxt = (i_t + 1) % 2 + q_tn = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, i_t + 1, i_h, lane)) + k_tn = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, i_t + 1, i_h, lane)) + cute.autovec_copy(q_tn, r_qbf[nxt]) + cute.autovec_copy(k_tn, r_kbf[nxt]) + a_tn = cute.local_tile(a, (1, 1, 1, vec_size), (i_n, i_t + 1, i_hv, lane)) + cute.autovec_copy(a_tn, r_abf[nxt]) + v_tn = cute.local_tile(v, (1, 1, 1, BV), (i_n, i_t + 1, i_hv, i_v)) + cute.autovec_copy(v_tn, r_vbf[nxt]) + r_bbf[nxt][0] = cutlass.Float32(b[i_n, i_t + 1, i_hv]) + + # ===== prep: read q/k + gate<->l2norm cross-pipe interleave ===== + for c in cutlass.range_constexpr(vec_size): + r_q[c] = cutlass.Float32(r_qbf[cur][c]) + r_k[c] = cutlass.Float32(r_kbf[cur][c]) + + # gate stage 1: x=a+dtb + for c in cutlass.range_constexpr(vec_size): + r_gx[c] = cutlass.Float32(r_abf[cur][c]) + r_dtb[c] # x = a + dt_bias + if cutlass.const_expr(not use_lower_bound): + for c in cutlass.range_constexpr(vec_size): + r_gexp[c] = cute.exp(softplus_beta * r_gx[c], fastmath=fast_math) # exp(beta_x) + + if cutlass.const_expr(use_qk_l2norm): + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + sum_q += r_q[c] * r_q[c] + sum_k += r_k[c] * r_k[c] + for off in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=off, mask=-1, mask_and_clamp=31) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=off, mask=-1, mask_and_clamp=31) + inv_q = cute.rsqrt(sum_q + 1e-6, fastmath=fast_math) * scale + inv_k = cute.rsqrt(sum_k + 1e-6, fastmath=fast_math) + for c in cutlass.range_constexpr(vec_size): + r_q[c] = r_q[c] * inv_q + r_k[c] = r_k[c] * inv_k + else: + for c in cutlass.range_constexpr(vec_size): + r_q[c] = r_q[c] * scale + + # gate stage 2: finalize per-channel decay r_g + if cutlass.const_expr(use_lower_bound): + # safe gate: g = lower_bound * sigmoid(exp(A_log) * x) + for c in cutlass.range_constexpr(vec_size): + sigmoid_ax = cutlass.Float32(1.0) / ( + cutlass.Float32(1.0) + cute.exp(-r_exp_A * r_gx[c], fastmath=fast_math) + ) + r_g[c] = cute.exp(lower_bound * sigmoid_ax, fastmath=fast_math) + else: + for c in cutlass.range_constexpr(vec_size): + beta_x = softplus_beta * r_gx[c] + sp_val = (cutlass.Float32(1.0) / softplus_beta) * cute.log( + cutlass.Float32(1.0) + r_gexp[c], fastmath=fast_math + ) + use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) + r_g[c] = use_sp * sp_val + (cutlass.Float32(1.0) - use_sp) * r_gx[c] # stash sp_x + for c in cutlass.range_constexpr(vec_size): + r_g[c] = cute.exp(-r_exp_A * r_g[c], fastmath=fast_math) # final exp (batched) + + r_beta = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-r_bbf[cur][0], fastmath=fast_math)) + + # ===== recurrence (fused: decay+h@k in one pass / update+h@q in one pass) ===== + for vv in cutlass.range_constexpr(BV): + sv = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + r_h[vv * vec_size + c] = r_h[vv * vec_size + c] * r_g[c] # decay: h *= exp(g) (per K) + sv += r_h[vv * vec_size + c] * r_k[c] # s = sum_k h*k_norm + r_red[vv] = sv + for off in [16, 8, 4, 2, 1]: + for vv in cutlass.range_constexpr(BV): + r_red[vv] = r_red[vv] + cute.arch.shuffle_sync_bfly(r_red[vv], offset=off, mask=-1, mask_and_clamp=31) + for vv in cutlass.range_constexpr(BV): + v_new = (cutlass.Float32(r_vbf[cur][vv]) - r_red[vv]) * r_beta # v_new = beta*(v - s) + ovv = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + r_h[vv * vec_size + c] = r_h[vv * vec_size + c] + r_k[c] * v_new # rank-1 update: h += k*v_new + ovv += r_h[vv * vec_size + c] * r_q[c] # o = sum_k h*q_scaled (partial) + r_red[vv] = ovv + for off in [16, 8, 4, 2, 1]: + for vv in cutlass.range_constexpr(BV): + r_red[vv] = r_red[vv] + cute.arch.shuffle_sync_bfly(r_red[vv], offset=off, mask=-1, mask_and_clamp=31) + for vv in cutlass.range_constexpr(BV): + o[(i_n, i_t, i_hv, i_v * BV + vv)] = cutlass.BFloat16(r_red[vv]) + if cutlass.const_expr(cache_intermediate_states): # Stage-D snapshot: post-token-t state + flat_idx = i_n * T * HV + i_t * HV + i_hv + for vv in cutlass.range_constexpr(BV): + for c in cutlass.range_constexpr(vec_size): + r_h4[c] = r_h[vv * vec_size + c] + inter_tile = cute.local_tile(intermediate_states, (1, 1, vec_size), (flat_idx, i_v * BV + vv, lane)) + cute.autovec_copy(r_h4, inter_tile) + + # ===== epilogue: write state back ===== + if cutlass.const_expr(not disable_state_update): + flat_state_idx = cache_idx * HV + i_hv + for vv in cutlass.range_constexpr(BV): + v_global = i_v * BV + vv + for c in cutlass.range_constexpr(vec_size): + r_h4[c] = r_h[vv * vec_size + c] + h_out = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_global, lane)) + cute.autovec_copy(r_h4, h_out) + + +@cute.jit +def run_kda_mtp_recurrent_vk_kernel( + h0_source: cute.Tensor, + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + intermediate_states: cute.Tensor, + h0_indices: cute.Tensor, + vec_size: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + cache_intermediate_states: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], + stream: cuda.CUstream, +): + """lane=K vk launcher:grid = N*HV*(V//BV),block = 32(1 warp)。无 SMEM。""" + n_indices = h0_indices.layout.shape[0] + num_v_tiles = cute.ceil_div(V, BV) + grid_size = n_indices * HV * num_v_tiles + + kda_mtp_recurrent_vk_kernel( + h0_source, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + intermediate_states, + h0_indices, + vec_size, + num_v_tiles, + BV, + softplus_beta, + softplus_threshold, + scale, + HV, + T, + H, + K, + V, + use_qk_l2norm, + disable_state_update, + cache_intermediate_states, + fast_math, + use_lower_bound, + lower_bound, + ).launch( + grid=(grid_size, 1, 1), + block=[32, 1, 1], + smem=0, + stream=stream, + ) + + +_compiled_mtp_vk_kernels: dict[tuple, object] = {} + + +def _get_compiled_mtp_vk_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + BV, + scale, + use_qk_l2norm, + disable_state_update, + softplus_beta, + softplus_threshold, + opt_level=3, + fast_math=True, + cache_intermediate_states=False, + use_lower_bound=False, + lower_bound=0.0, + dyn_stride=False, +): + key = ( + T, + H, + HV, + K, + V, + BV, + scale, + use_qk_l2norm, + disable_state_update, + cache_intermediate_states, + softplus_beta, + softplus_threshold, + opt_level, + fast_math, + use_lower_bound, + lower_bound, + dyn_stride, + ) + if key in _compiled_mtp_vk_kernels: + return _compiled_mtp_vk_kernels[key] + + q = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, T, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, T, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + h0_source = torch.zeros(pool_size * HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + if cache_intermediate_states: + intermediate_states = torch.zeros(N * T * HV, V, K, dtype=torch.float32, device="cuda") + else: + intermediate_states = torch.empty(1, 1, 1, dtype=torch.float32, device="cuda") + + # dynamic-N: mark the batch axis (dim 0) dynamic so one cubin serves all N. + # Explicit stride_order: at N=1/T=1 the size-1 dims make auto-deduction ambiguous. + if dyn_stride: + # dyn-stride: q/k/v arrive as K-contiguous strided views (the caller + # skipped the contiguous copy). Mark shape AND strides dynamic with + # the innermost K axis static so vectorized loads stay legal. + q_t = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=3) + k_t = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=3) + v_t = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=3) + else: + q_t = from_dlpack(q, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=q.dim_order()) + k_t = from_dlpack(k, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=k.dim_order()) + v_t = from_dlpack(v, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=v.dim_order()) + a_t = from_dlpack(a, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()) + b_t = from_dlpack(b, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=b.dim_order()) + o_t = from_dlpack(o, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=o.dim_order()) + A_log_t = from_dlpack(A_log, assumed_align=16) + dt_bias_t = from_dlpack(dt_bias, assumed_align=16) + h0_source_t = from_dlpack(h0_source, assumed_align=16).mark_compact_shape_dynamic( + mode=0, stride_order=h0_source.dim_order() + ) + h0_indices_t = from_dlpack(h0_indices, assumed_align=16).mark_layout_dynamic() + intermediate_states_t = from_dlpack(intermediate_states, assumed_align=16) + if cache_intermediate_states: + intermediate_states_t = intermediate_states_t.mark_compact_shape_dynamic( + mode=0, stride_order=intermediate_states.dim_order() + ) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + compiled_kernel = cute.compile( + run_kda_mtp_recurrent_vk_kernel, + h0_source_t, + A_log_t, + a_t, + dt_bias_t, + q_t, + k_t, + v_t, + b_t, + o_t, + intermediate_states_t, + h0_indices_t, + vec_size=VEC_SIZE, + BV=BV, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + use_qk_l2norm=use_qk_l2norm, + disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, + fast_math=fast_math, + use_lower_bound=use_lower_bound, + lower_bound=lower_bound, + stream=stream, + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + + _compiled_mtp_vk_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA MTP small-batch VK(lane=K) kernel compiled: " + f"N={N}, T={T}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, BV={BV}, " + f"opt_level={opt_level}, fast_math={fast_math}" + ) + return compiled_kernel + + +def _select_vk_bv(work_units, V, num_sms): + waves32 = work_units * (V // 32) / (num_sms * 12) + if V % 8 == 0 and waves32 < 3.0: + return 8 + return 32 + + +# T>4 recurrent dispatch: below this work-unit (N*HV) count the single-warp vk kernel +# still beats warp-spec at high T; at/above it vk hits the DRAM-bandwidth wall (kernel bench). +_WS_WORK_UNIT_THRESHOLD = 2048 + + +def kda_decode_mtp( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + state_layout: str = "vk", + disable_state_update: bool = False, + intermediate_states_buffer: torch.Tensor | None = None, + lower_bound: float | None = None, +) -> torch.Tensor: + common = dict( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + out=out, + disable_state_update=disable_state_update, + lower_bound=lower_bound, + intermediate_states_buffer=intermediate_states_buffer, + ) + if state_layout == "kv": + return kda_decode_mtp_recurrent(**common, variant="kv", k_split=-1) # k_split auto + T = q.shape[1] + work_units = q.shape[0] * v.shape[2] # N * HV + # N*HV >= _WS_WORK_UNIT_THRESHOLD: single-warp vk is DRAM-bandwidth-bound (worse under dyn-stride) -> warp-spec. + if T < 4 or work_units < _WS_WORK_UNIT_THRESHOLD: + return kda_decode_mtp_recurrent(**common, variant="vk", bv=-1) # bv auto + return kda_decode_mtp_recurrent_ws(**common, state_layout="vk") diff --git a/cula/ops/kda/decode/mtp_kvbuffer.py b/cula/ops/kda/decode/mtp_kvbuffer.py new file mode 100644 index 00000000..ae24abbb --- /dev/null +++ b/cula/ops/kda/decode/mtp_kvbuffer.py @@ -0,0 +1,1950 @@ +"""CuTe DSL KDA MTP decode — KVBuffer / chunkwise parallel-verification variant. + +KVBuffer paper's chunkwise verify form (https://arxiv.org/abs/2605.19049) as a new +operator vs the recurrent vk/kv ops in ``kda_decode_mtp.py``. The T draft tokens +are treated as ONE chunk: per-token outputs come from the FIXED input state S0 plus a +small T×T intra-chunk correction, and the state is updated once at the end — the +S0-matvecs are independent across tokens (no length-T serial chain), the latency angle +at small batch. Infra (grid N*HV*(V//BV), 1 warp/CTA, lane=K, float4 loads, butterfly +reduce-over-K) mirrors the production vk kernel for apples-to-apples comparison. + +Chunkwise math (state S0[v,k], decay-first; matches the recurrent op): + g_t[k] = exp(-exp(A_log) * softplus(a_t[k] + dt_bias[k])) # per channel + b_t[k] = prod_{i<=t} g_i[k] # cumulative decay + kdec_t = k_norm_t * b_t ; qdec_t = q_scaled_t * b_t + r(t,i) = prod_{it} g_j # suffix-decayed key + S_T[v,k]= b_{T-1}[k] * S0[v,k] + sum_i u_i[v] ksuf_i[k] # full accept + +Numerical form: every decay factor is an ORDERED product bounded by 1 — there is +no division by the cumulative gate product (which can underflow to 0 in fp32 +under unbounded softplus gates). The op is valid for both softplus and safe-gate +models. +The scratch stores raw (u_i, k_i, g_i) per token — the same triplet as ReplaySSM's +(d, k, g) ring — and the flush rebuilds S_m with descending suffix products. +""" + +import logging + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from cula.ops.kda.decode.cute import ( + TILE_K, + _get_cached_stream, + _normalize_A_log, + _normalize_dt_bias, + _normalize_state_indices, + _normalize_state_source, + _prepare_output_tensor, +) +from cula.ops.kda.decode.mtp import ( + VEC_SIZE, + _normalize_mtp_a, +) + +logger = logging.getLogger(__name__) + + +# tile_v by WU=N*HV: <=32->16, <256->32, >=256->64 (H200 sweep). +def _select_kvb_tile_v(V, N, HV): + """work-unit (N*HV) dependent tile_v. Returns the first candidate that divides V.""" + wu = N * HV + if wu <= 32: + order = (16, 32, 8, 64) + elif wu < 256: + order = (32, 64, 16, 8) + else: + order = (64, 32, 16, 8) + for tv in order: + if V % tv == 0: + return tv + return 8 + + +# flush BV = smallest tile (DRAM-latency bound; bv=8 > bv=32 ~18% at large N*HV). +def _select_flush_bv(V): + for bv in (8, 16, 32): + if V % bv == 0: + return bv + raise ValueError(f"V={V} must be divisible by 8, 16 or 32") + + +# flush kernel: rank-m rebuild of S_m from compact (u,k,g) scratch (Phase-D, lane=K+vk): +# S_m[v,k] = prod_{j= 0: + flat_state_idx = cache_idx * HV + i_hv + m_n = m_buf[i_n] # this request's accept length (runtime; 1 <= m_n <= T) + + r_acc = cute.make_rmem_tensor(cute.make_layout((BV * vec_size,), stride=(1,)), cutlass.Float32) + r_h4 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_suf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_g = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + + # Overflow-safe rebuild: descending suffix products (all factors <=1, no division). + for c in cutlass.range_constexpr(vec_size): + r_suf[c] = cutlass.Float32(1.0) + for j in cutlass.range_constexpr(BV * vec_size): + r_acc[j] = cutlass.Float32(0.0) + for tt in cutlass.range_constexpr(T): + i_i = T - 1 - tt + if i_i < m_n: + k_tile = cute.local_tile(k_buf, (1, 1, 1, vec_size), (i_n, i_i, i_hv, lane)) + cute.autovec_copy(k_tile, r_k) + g_tile = cute.local_tile(g_buf, (1, 1, 1, vec_size), (i_n, i_i, i_hv, lane)) + cute.autovec_copy(g_tile, r_g) + for vv in cutlass.range_constexpr(BV): + uval = cutlass.Float32(d_buf[i_n, i_i, i_hv, i_v * BV + vv]) + for c in cutlass.range_constexpr(vec_size): + r_acc[vv * vec_size + c] += uval * r_k[c] * r_suf[c] + for c in cutlass.range_constexpr(vec_size): + r_suf[c] = r_suf[c] * r_g[c] + + # S_m = prefix * S0 + acc, write back (contiguous float4) + for vv in cutlass.range_constexpr(BV): + v_global = i_v * BV + vv + h_tile = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_global, lane)) + cute.autovec_copy(h_tile, r_h4) + for c in cutlass.range_constexpr(vec_size): + r_h4[c] = r_suf[c] * r_h4[c] + r_acc[vv * vec_size + c] + h_out = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_global, lane)) + cute.autovec_copy(r_h4, h_out) + + +@cute.jit +def run_kda_flush_kvbuffer_vk_kernel( + h0_source: cute.Tensor, + d_buf: cute.Tensor, + k_buf: cute.Tensor, + g_buf: cute.Tensor, + h0_indices: cute.Tensor, + m_buf: cute.Tensor, + vec_size: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + stream: cuda.CUstream, +): + n_indices = h0_indices.layout.shape[0] + num_v_tiles = cute.ceil_div(V, BV) + grid_size = n_indices * HV * num_v_tiles + kda_flush_kvbuffer_vk_kernel( + h0_source, + d_buf, + k_buf, + g_buf, + h0_indices, + m_buf, + vec_size, + num_v_tiles, + BV, + HV, + T, + K, + V, + ).launch(grid=(grid_size, 1, 1), block=[32, 1, 1], smem=0, stream=stream) + + +_compiled_flush_kvbuffer_kernels: dict[tuple, object] = {} + + +def _get_compiled_flush_kvbuffer_kernel(N, T, HV, K, V, pool_size, BV, opt_level=3): + key = (N, T, HV, K, V, pool_size, BV, opt_level) + if key in _compiled_flush_kvbuffer_kernels: + return _compiled_flush_kvbuffer_kernels[key] + + h0_source = torch.zeros(pool_size * HV, V, K, dtype=torch.float32, device="cuda") + d_buf = torch.zeros(N, T, HV, V, dtype=torch.float32, device="cuda") + k_buf = torch.zeros(N, T, HV, K, dtype=torch.float32, device="cuda") + g_buf = torch.zeros(N, T, HV, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + m_buf = torch.zeros(N, dtype=torch.int32, device="cuda") + + compiled = cute.compile( + run_kda_flush_kvbuffer_vk_kernel, + from_dlpack(h0_source, assumed_align=16), + from_dlpack(d_buf, assumed_align=16), + from_dlpack(k_buf, assumed_align=16), + from_dlpack(g_buf, assumed_align=16), + from_dlpack(h0_indices, assumed_align=16), + from_dlpack(m_buf, assumed_align=16), + vec_size=VEC_SIZE, + BV=BV, + HV=HV, + T=T, + K=K, + V=V, + stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + _compiled_flush_kvbuffer_kernels[key] = compiled + logger.info(f"CuTe DSL KDA flush KVBuffer kernel compiled: N={N}, T={T}, HV={HV}, K={K}, V={V}, BV={BV}") + return compiled + + +def kda_flush_kvbuffer( + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + d_buffer: torch.Tensor, + k_buffer: torch.Tensor, + g_buffer: torch.Tensor, + accept_len, # int (broadcast to all N) OR per-request [N] int tensor; each in [1, T] + bv: int = -1, + opt_level: int = 3, +) -> torch.Tensor: + N, T, HV, V = d_buffer.shape + K = k_buffer.shape[3] + if isinstance(accept_len, torch.Tensor): + assert accept_len.numel() == N, f"per-request accept_len must have N={N} entries, got {accept_len.numel()}" + m_buf = accept_len.to(device=d_buffer.device, dtype=torch.int32).contiguous() + else: + m = int(accept_len) + assert 1 <= m <= T, f"accept_len must be in [1,{T}], got {m}" + m_buf = torch.full((N,), m, dtype=torch.int32, device=d_buffer.device) + + if bv <= 0: + bv = _select_flush_bv(V) + assert bv in (8, 16, 32) and V % bv == 0, f"flush bv must be 8/16/32 and divide V, got bv={bv}, V={V}" + + h0_source, pool_size, _ = _normalize_state_source( + initial_state_source, + N=N, + HV=HV, + K=K, + V=V, + device=initial_state_source.device, + state_layout="vk", + ) + initial_state_indices = _normalize_state_indices( + initial_state_indices, N=N, pool_size=pool_size, device=initial_state_source.device + ) + stream = _get_cached_stream(initial_state_source.device) + + h0_source_flat = h0_source.view(pool_size * HV, V, K) + compiled = _get_compiled_flush_kvbuffer_kernel(N, T, HV, K, V, pool_size, bv, opt_level=opt_level) + compiled(h0_source_flat, d_buffer, k_buffer, g_buffer, initial_state_indices, m_buf, stream) + return initial_state_source + + +# =========================================================================== +# MULTILAYER_FLUSH_PATCH: all-layers batched flush, dynamic-N (2D grid x=layer-grid, y=layer). +@cute.kernel +def kda_flush_kvbuffer_vk_ml_kernel( + h0_source: cute.Tensor, + d_buf: cute.Tensor, + k_buf: cute.Tensor, + g_buf: cute.Tensor, + h0_indices: cute.Tensor, + m_buf: cute.Tensor, + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + lane = tidx + + bx, i_l, _ = cute.arch.block_idx() + i_v = bx % num_v_tiles + tmp = bx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + + cache_idx = h0_indices[i_n] + if cache_idx >= 0: + flat_state_idx = cache_idx * HV + i_hv + m_n = m_buf[i_n] + + r_acc = cute.make_rmem_tensor(cute.make_layout((BV * vec_size,), stride=(1,)), cutlass.Float32) + r_h4 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_suf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_g = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + + # Overflow-safe rebuild via descending suffix products (see single-layer flush). + for c in cutlass.range_constexpr(vec_size): + r_suf[c] = cutlass.Float32(1.0) + for j in cutlass.range_constexpr(BV * vec_size): + r_acc[j] = cutlass.Float32(0.0) + for tt in cutlass.range_constexpr(T): + i_i = T - 1 - tt + if i_i < m_n: + k_tile = cute.local_tile(k_buf, (1, 1, 1, 1, vec_size), (i_l, i_n, i_i, i_hv, lane)) + cute.autovec_copy(k_tile, r_k) + g_tile = cute.local_tile(g_buf, (1, 1, 1, 1, vec_size), (i_l, i_n, i_i, i_hv, lane)) + cute.autovec_copy(g_tile, r_g) + for vv in cutlass.range_constexpr(BV): + uval = cutlass.Float32(d_buf[i_l, i_n, i_i, i_hv, i_v * BV + vv]) + for c in cutlass.range_constexpr(vec_size): + r_acc[vv * vec_size + c] += uval * r_k[c] * r_suf[c] + for c in cutlass.range_constexpr(vec_size): + r_suf[c] = r_suf[c] * r_g[c] + + for vv in cutlass.range_constexpr(BV): + v_global = i_v * BV + vv + h_tile = cute.local_tile(h0_source, (1, 1, 1, vec_size), (i_l, flat_state_idx, v_global, lane)) + cute.autovec_copy(h_tile, r_h4) + for c in cutlass.range_constexpr(vec_size): + r_h4[c] = r_suf[c] * r_h4[c] + r_acc[vv * vec_size + c] + h_out = cute.local_tile(h0_source, (1, 1, 1, vec_size), (i_l, flat_state_idx, v_global, lane)) + cute.autovec_copy(r_h4, h_out) + + +@cute.jit +def run_kda_flush_kvbuffer_vk_ml_kernel( + h0_source: cute.Tensor, + d_buf: cute.Tensor, + k_buf: cute.Tensor, + g_buf: cute.Tensor, + h0_indices: cute.Tensor, + m_buf: cute.Tensor, + vec_size: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + stream: cuda.CUstream, +): + L = h0_source.layout.shape[0] + n_indices = h0_indices.layout.shape[0] + num_v_tiles = cute.ceil_div(V, BV) + gx = n_indices * HV * num_v_tiles + kda_flush_kvbuffer_vk_ml_kernel( + h0_source, + d_buf, + k_buf, + g_buf, + h0_indices, + m_buf, + vec_size, + num_v_tiles, + BV, + HV, + T, + K, + V, + ).launch(grid=(gx, L, 1), block=[32, 1, 1], smem=0, stream=stream) + + +_compiled_flush_kvbuffer_ml_kernels: dict[tuple, object] = {} + + +def _get_compiled_flush_kvbuffer_ml_kernel( + L, T, HV, K, V, pool_size, kvb_pool, BV, h0_source, d_buf, k_buf, g_buf, h0_indices, m_buf, opt_level=3 +): + # Trace on the tensors passed in. N not in key (index layout-dynamic). + key = (L, T, HV, K, V, pool_size, kvb_pool, BV, opt_level) + if key in _compiled_flush_kvbuffer_ml_kernels: + return _compiled_flush_kvbuffer_ml_kernels[key] + + compiled = cute.compile( + run_kda_flush_kvbuffer_vk_ml_kernel, + from_dlpack(h0_source, assumed_align=16), + from_dlpack(d_buf, assumed_align=16), + from_dlpack(k_buf, assumed_align=16), + from_dlpack(g_buf, assumed_align=16), + from_dlpack(h0_indices, assumed_align=16).mark_layout_dynamic(), + from_dlpack(m_buf, assumed_align=16).mark_layout_dynamic(), + vec_size=VEC_SIZE, + BV=BV, + HV=HV, + T=T, + K=K, + V=V, + stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + _compiled_flush_kvbuffer_ml_kernels[key] = compiled + logger.info(f"CuTe DSL KDA flush KVBuffer ML(dyn-N) kernel compiled: L={L}, T={T}, HV={HV}, K={K}, V={V}, BV={BV}") + return compiled + + +def kda_flush_kvbuffer_all_layers( + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + d_buffer: torch.Tensor, + k_buffer: torch.Tensor, + g_buffer: torch.Tensor, + accept_len, + bv: int = -1, + opt_level: int = 3, +) -> torch.Tensor: + L, kvb_pool, T, HV, V = d_buffer.shape + K = k_buffer.shape[4] + N = initial_state_indices.shape[0] + if isinstance(accept_len, torch.Tensor): + assert accept_len.numel() == N, f"per-request accept_len must have N={N} entries, got {accept_len.numel()}" + m_buf = accept_len.to(device=d_buffer.device, dtype=torch.int32).contiguous() + else: + m = int(accept_len) + assert 1 <= m <= T, f"accept_len must be in [1,{T}], got {m}" + m_buf = torch.full((N,), m, dtype=torch.int32, device=d_buffer.device) + + if bv <= 0: + bv = _select_flush_bv(V) + assert bv in (8, 16, 32) and V % bv == 0, f"flush bv must be 8/16/32 and divide V, got bv={bv}, V={V}" + + pool_size = initial_state_source.shape[1] + h0_source_flat = initial_state_source.view(L, pool_size * HV, V, K) + idx = _normalize_state_indices(initial_state_indices, N=N, pool_size=pool_size, device=initial_state_source.device) + stream = _get_cached_stream(initial_state_source.device) + + compiled = _get_compiled_flush_kvbuffer_ml_kernel( + L, + T, + HV, + K, + V, + pool_size, + kvb_pool, + bv, + h0_source_flat, + d_buffer, + k_buffer, + g_buffer, + idx, + m_buf, + opt_level=opt_level, + ) + compiled(h0_source_flat, d_buffer, k_buffer, g_buffer, idx, m_buf, stream) + return initial_state_source + + +# --------------------------------------------------------------------------- +# shuffle-kvbuffer: token-parallel chunkwise verify (structure B). UT-transform +# W = L^{-1} diag(beta) makes the consumer solve dependence-free: u = W @ (v - S0 kdec). +# --------------------------------------------------------------------------- +@cute.kernel +def kda_mtp_shuffle_kvbuffer_kernel( + h0_source: cute.Tensor, # [pool*HV, V, K] fp32 (vk) + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + d_buf: cute.Tensor, # [N, T, HV, V] fp32 + k_buf: cute.Tensor, # [N, T, HV, K] fp32 raw normalized key k_t + g_buf: cute.Tensor, # [N, T, HV, K] fp32 per-step gate g_t + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + emit_output: cutlass.Constexpr[bool], + write_ubuf: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + num_warps: cutlass.Constexpr[int] = 4 + + bidx, _, _ = cute.arch.block_idx() + i_v = bidx % num_v_tiles + tmp = bidx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv]), fastmath=fast_math) + + # SMEM. sKdec/sQdec double as staging for k_norm/q_scaled between Stage 1 and 2. + smem = cutlass.utils.SmemAllocator() + sKdec = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sKn = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sQdec = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sBrun = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, K), stride=(K + 8, 1)), 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T,)), 16) + sA = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, T), stride=(T, 1)), 16) + sP = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, T), stride=(T, 1)), 16) + sW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, T), stride=(T, 1)), 16) + + r_qbf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + r_kbf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + r_qf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_kf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_dtb = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_tmp = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_h = cute.make_rmem_tensor(cute.make_layout((ilp_rows, vec_size), stride=(vec_size, 1)), cutlass.Float32) + # r_part: ilp_rows*T batched partials (Skdec, then reused as x = v - Skdec, then Sqdec). + r_part = cute.make_rmem_tensor(cute.make_layout((ilp_rows, T), stride=(T, 1)), cutlass.Float32) + r_u = cute.make_rmem_tensor(cute.make_layout((ilp_rows, T), stride=(T, 1)), cutlass.Float32) + # Stage-3 pair partials: ceil(T*T/4) per warp. + ppw: cutlass.Constexpr[int] = (T * T + num_warps - 1) // num_warps + r_red = cute.make_rmem_tensor(cute.make_layout((ppw,), stride=(1,)), cutlass.Float32) + + if cache_idx >= 0: + k_start = lane_id * vec_size + rows_per_group: cutlass.Constexpr[int] = tile_v // num_warps + flat_state_idx = cache_idx * HV + i_hv + + # ---- Stage 1: token-parallel gating/l2norm (warp w owns tokens w, w+4, ...) ---- + for c in cutlass.range_constexpr(vec_size): + r_dtb[c] = cutlass.Float32(dt_bias[i_hv, k_start + c]) + tokens_per_warp: cutlass.Constexpr[int] = (T + num_warps - 1) // num_warps + for tt in cutlass.range_constexpr(tokens_per_warp): + t_tok = tt * num_warps + warp_idx + if t_tok < T: + q_tile = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, t_tok, i_h, lane_id)) + k_tile = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, t_tok, i_h, lane_id)) + cute.autovec_copy(q_tile, r_qbf) + cute.autovec_copy(k_tile, r_kbf) + for c in cutlass.range_constexpr(vec_size): + r_qf[c] = cutlass.Float32(r_qbf[c]) + r_kf[c] = cutlass.Float32(r_kbf[c]) + + if cutlass.const_expr(use_qk_l2norm): + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + sum_q += r_qf[c] * r_qf[c] + sum_k += r_kf[c] * r_kf[c] + for off in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=off, mask=-1, mask_and_clamp=31) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=off, mask=-1, mask_and_clamp=31) + inv_q = cute.rsqrt(sum_q + 1e-6, fastmath=fast_math) * scale + inv_k = cute.rsqrt(sum_k + 1e-6, fastmath=fast_math) + for c in cutlass.range_constexpr(vec_size): + r_qf[c] = r_qf[c] * inv_q + r_kf[c] = r_kf[c] * inv_k + else: + for c in cutlass.range_constexpr(vec_size): + r_qf[c] = r_qf[c] * scale + + # gate g_t per channel; stage k_norm/q_scaled (decay applied in Stage 2) + for c in cutlass.range_constexpr(vec_size): + x = cutlass.Float32(a[i_n, t_tok, i_hv, k_start + c]) + r_dtb[c] + if cutlass.const_expr(use_lower_bound): + sigmoid_ax = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-r_exp_A * x, fastmath=fast_math)) + sG[t_tok, k_start + c] = cute.exp(lower_bound * sigmoid_ax, fastmath=fast_math) + else: + beta_x = softplus_beta * x + exp_bx = cute.exp(beta_x, fastmath=fast_math) + sp_val = (cutlass.Float32(1.0) / softplus_beta) * cute.log( + cutlass.Float32(1.0) + exp_bx, fastmath=fast_math + ) + use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) + sp_x = use_sp * sp_val + (cutlass.Float32(1.0) - use_sp) * x + sG[t_tok, k_start + c] = cute.exp(-r_exp_A * sp_x, fastmath=fast_math) + sKdec[t_tok, k_start + c] = r_kf[c] + sQdec[t_tok, k_start + c] = r_qf[c] + if lane_id == 0: + sBeta[t_tok] = cutlass.Float32(1.0) / ( + cutlass.Float32(1.0) + cute.exp(-cutlass.Float32(b[i_n, t_tok, i_hv]), fastmath=fast_math) + ) + cute.arch.barrier() + + # ---- Stage 2: K-parallel prefix-product scan (thread = one channel). + kc = tidx # requires K == 128 == block size + b_run_s = cutlass.Float32(1.0) + for i_t in cutlass.range_constexpr(T): + kn = sKdec[i_t, kc] + g_t = sG[i_t, kc] + b_run_s = b_run_s * g_t + sKdec[i_t, kc] = kn * b_run_s + sKn[i_t, kc] = kn + sBrun[i_t, kc] = b_run_s + if cutlass.const_expr(write_ubuf): + if i_v == 0: + k_buf[i_n, i_t, i_hv, kc] = kn # raw key (was k/b_run) + g_buf[i_n, i_t, i_hv, kc] = g_t # per-step gate (was b_run) + cute.arch.barrier() + + # ---- Stage 3: (t,i)-parallel A/P, T^2 pairs round-robined over 4 warps, + # ONE batched butterfly per warp. Pair p: p < T*(T-1)/2 -> A, else P. ---- + for j in cutlass.range_constexpr(ppw): + r_red[j] = cutlass.Float32(0.0) + p_ctr = 0 + for i_t in cutlass.range_constexpr(T): + for i_i in cutlass.range_constexpr(i_t): # A[t,i], i no cross-lane sync needed. ---- + if warp_idx == 0: + if lane_id < T: + for i_t in cutlass.range_constexpr(T): + eq = cutlass.Float32(1.0) if lane_id == i_t else cutlass.Float32(0.0) + acc_w = eq + for i_i in cutlass.range_constexpr(i_t): + acc_w -= sA[i_t, i_i] * sW[i_i, lane_id] + sW[i_t, lane_id] = sBeta[i_t] * acc_w + cute.arch.barrier() + + # ---- Stage 4: consumer (4 warp groups over V rows), zero serial deps. ---- + n_row_groups: cutlass.Constexpr[int] = rows_per_group // ilp_rows + for rg in cutlass.range_constexpr(n_row_groups): + v_base = i_v * tile_v + warp_idx * rows_per_group + rg * ilp_rows + for r in cutlass.range_constexpr(ilp_rows): + h_tile = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_base + r, lane_id)) + cute.autovec_copy(h_tile, cute.slice_(r_h, (r, None))) + # all T Skdec_t for all ilp_rows rows in ONE batched butterfly + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + s = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + s += r_h[r, c] * sKdec[i_t, k_start + c] + r_part[r, i_t] = s + for off in [16, 8, 4, 2, 1]: + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + r_part[r, i_t] += cute.arch.shuffle_sync_bfly(r_part[r, i_t], offset=off, mask=-1, mask_and_clamp=31) + # x = v - Skdec (r_part reused), then u = W @ x (token-parallel, no dep chain) + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + r_part[r, i_t] = cutlass.Float32(v[i_n, i_t, i_hv, v_base + r]) - r_part[r, i_t] + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + acc = cutlass.Float32(0.0) + for i_i in cutlass.range_constexpr(i_t + 1): + acc += sW[i_t, i_i] * r_part[r, i_i] + r_u[r, i_t] = acc + if cutlass.const_expr(write_ubuf): + if lane_id == 0: + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + d_buf[i_n, i_t, i_hv, v_base + r] = r_u[r, i_t] + # o_t = Sqdec_t + sum_{i<=t} P[t,i] u_i (Sqdec batched butterfly into r_part) + if cutlass.const_expr(emit_output): + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + s = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + s += r_h[r, c] * sQdec[i_t, k_start + c] * sBrun[i_t, k_start + c] + r_part[r, i_t] = s + for off in [16, 8, 4, 2, 1]: + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + r_part[r, i_t] += cute.arch.shuffle_sync_bfly( + r_part[r, i_t], offset=off, mask=-1, mask_and_clamp=31 + ) + for r in cutlass.range_constexpr(ilp_rows): + for i_t in cutlass.range_constexpr(T): + ov = r_part[r, i_t] + for i_i in cutlass.range_constexpr(i_t + 1): + ov += sP[i_t, i_i] * r_u[r, i_i] + if lane_id == 0: + o[(i_n, i_t, i_hv, v_base + r)] = cutlass.BFloat16(ov) + # final state S_T[v,k] = b_{T-1}[k]*S0[v,k] + sum_t u_t k_t[k]*suf(t)[k], + # suf(t) = prod_{j>t} g_j accumulated descending (bounded <= 1; the + # running product ends as the full prefix for the S0 term). + if cutlass.const_expr(not disable_state_update): + for r in cutlass.range_constexpr(ilp_rows): + for c in cutlass.range_constexpr(vec_size): + acc = cutlass.Float32(0.0) + suf = cutlass.Float32(1.0) + for tt in cutlass.range_constexpr(T): + i_t = T - 1 - tt + acc += r_u[r, i_t] * sKn[i_t, k_start + c] * suf + suf = suf * sG[i_t, k_start + c] + r_tmp[c] = suf * r_h[r, c] + acc + h_out = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_base + r, lane_id)) + cute.autovec_copy(r_tmp, h_out) + + +@cute.jit +def run_kda_mtp_shuffle_kvbuffer_kernel( + h0_source: cute.Tensor, + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + d_buf: cute.Tensor, + k_buf: cute.Tensor, + g_buf: cute.Tensor, + vec_size: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + emit_output: cutlass.Constexpr[bool], + write_ubuf: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], + stream: cuda.CUstream, +): + """shuffle-kvbuffer launcher: grid = N*HV*(V//tile_v), block = 128 (4 warps).""" + n_indices = h0_indices.layout.shape[0] + num_v_tiles = cute.ceil_div(V, tile_v) + grid_size = n_indices * HV * num_v_tiles + smem_bytes = ( + 5 * 4 * T * (K + 8) # sKdec/sKn/sQdec/sG/sBrun + + 4 * T # sBeta + + 3 * 4 * T * T # sA/sP/sW + + 256 # alignment slack + ) + kda_mtp_shuffle_kvbuffer_kernel( + h0_source, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + h0_indices, + d_buf, + k_buf, + g_buf, + vec_size, + num_v_tiles, + tile_v, + ilp_rows, + softplus_beta, + softplus_threshold, + scale, + HV, + T, + H, + K, + V, + use_qk_l2norm, + disable_state_update, + emit_output, + write_ubuf, + fast_math, + use_lower_bound, + lower_bound, + ).launch(grid=(grid_size, 1, 1), block=[128, 1, 1], smem=smem_bytes, stream=stream) + + +_compiled_mtp_shuffle_kvbuffer_kernels: dict[tuple, object] = {} + + +def _dlp_qkv(_t, _dyn): + # dyn-stride: K-contiguous strided view -> dynamic-layout tensor (no copy); + # contiguous input keeps the compact (byte-identical) descriptor. + if _dyn: + return from_dlpack(_t, assumed_align=16).mark_layout_dynamic(leading_dim=3) + return from_dlpack(_t, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=_t.dim_order()) + + +def _get_compiled_mtp_shuffle_kvbuffer_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + tile_v, + ilp_rows, + scale, + use_qk_l2norm, + disable_state_update, + emit_output, + write_ubuf, + softplus_beta, + softplus_threshold, + opt_level=3, + fast_math=True, + use_lower_bound=False, + lower_bound=0.0, + dyn_stride=False, +): + key = ( + T, + H, + HV, + K, + V, + tile_v, + ilp_rows, + scale, + use_qk_l2norm, + disable_state_update, + emit_output, + write_ubuf, + softplus_beta, + softplus_threshold, + opt_level, + fast_math, + use_lower_bound, + lower_bound, + dyn_stride, + ) + if key in _compiled_mtp_shuffle_kvbuffer_kernels: + return _compiled_mtp_shuffle_kvbuffer_kernels[key] + + q = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, T, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, T, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + h0_source = torch.zeros(pool_size * HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + d_buf = torch.zeros(N, T, HV, V, dtype=torch.float32, device="cuda") + k_buf = torch.zeros(N, T, HV, K, dtype=torch.float32, device="cuda") + g_buf = torch.zeros(N, T, HV, K, dtype=torch.float32, device="cuda") + + compiled_kernel = cute.compile( + run_kda_mtp_shuffle_kvbuffer_kernel, + from_dlpack(h0_source, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=h0_source.dim_order()), + from_dlpack(A_log, assumed_align=16), + from_dlpack(a, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()), + from_dlpack(dt_bias, assumed_align=16), + _dlp_qkv(q, dyn_stride), + _dlp_qkv(k, dyn_stride), + _dlp_qkv(v, dyn_stride), + from_dlpack(b, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=b.dim_order()), + from_dlpack(o, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=o.dim_order()), + from_dlpack(h0_indices, assumed_align=16).mark_layout_dynamic(), + from_dlpack(d_buf, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=d_buf.dim_order()), + from_dlpack(k_buf, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=k_buf.dim_order()), + from_dlpack(g_buf, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=g_buf.dim_order()), + vec_size=VEC_SIZE, + tile_v=tile_v, + ilp_rows=ilp_rows, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + use_qk_l2norm=use_qk_l2norm, + disable_state_update=disable_state_update, + emit_output=emit_output, + write_ubuf=write_ubuf, + fast_math=fast_math, + use_lower_bound=use_lower_bound, + lower_bound=lower_bound, + stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + _compiled_mtp_shuffle_kvbuffer_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA MTP shuffle-KVBuffer kernel compiled: " + f"N={N}, T={T}, HV={HV}, K={K}, V={V}, tile_v={tile_v}, ilp_rows={ilp_rows}, " + f"opt_level={opt_level}, fast_math={fast_math}" + ) + return compiled_kernel + + +def _select_shuffle_kvb_ilp_rows(tile_v, T): + """Largest ilp_rows in {4,2,1} dividing rows_per_group with ilp_rows*T <= 16 — the consumer + holds two (ilp_rows, T) fp32 register arrays (r_part + r_u), so cap their footprint.""" + rows_per_group = tile_v // 4 + for r in (4, 2, 1): + if rows_per_group % r == 0 and r * T <= 16: + return r + return 1 + + +def kda_decode_mtp_shuffle_kvbuffer( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + disable_state_update: bool = True, + emit_output: bool = True, + d_buffer: torch.Tensor | None = None, + k_buffer: torch.Tensor | None = None, + g_buffer: torch.Tensor | None = None, + tile_v: int = -1, + ilp_rows: int = -1, + opt_level: int = 3, + fast_math: bool = True, + lower_bound: float | None = None, +) -> torch.Tensor: + """KDA MTP shuffle-KVBuffer verify (token-parallel chunkwise; flush reuses kda_flush_kvbuffer).""" + N, T, H, K = q.shape + HV = v.shape[2] + V = v.shape[3] + write_ubuf = d_buffer is not None + + if scale is None: + scale = K**-0.5 + else: + assert scale > 0, f"scale must be positive, got {scale}" + + assert K == TILE_K, f"shuffle-kvbuffer requires K={TILE_K}, got {K}" + assert K == 128, f"shuffle-kvbuffer Stage-2 scan maps 128 threads to K channels; needs K=128, got {K}" + assert T <= 32, f"shuffle-kvbuffer W-build uses one lane per token column; needs T<=32, got {T}" + + if tile_v <= 0: + tile_v = _select_kvb_tile_v(V, N, HV) + assert V % tile_v == 0, f"shuffle-kvbuffer requires V % tile_v == 0, got V={V}, tile_v={tile_v}" + assert tile_v % 4 == 0, f"shuffle-kvbuffer requires tile_v % 4 == 0 (4 warps), got {tile_v}" + rows_per_group = tile_v // 4 + if ilp_rows <= 0: + ilp_rows = _select_shuffle_kvb_ilp_rows(tile_v, T) + assert rows_per_group % ilp_rows == 0, ( + f"shuffle-kvbuffer requires (tile_v/4) % ilp_rows == 0, got tile_v={tile_v}, ilp_rows={ilp_rows}" + ) + + h0_source, pool_size, _ = _normalize_state_source( + initial_state_source, + N=N, + HV=HV, + K=K, + V=V, + device=q.device, + state_layout="vk", + ) + + a = _normalize_mtp_a(a, N=N, T=T, HV=HV, K=K) + if b.dim() != 3 or tuple(b.shape) != (N, T, HV): + raise ValueError(f"Unexpected b shape for MTP dense: {tuple(b.shape)}; expected {(N, T, HV)}") + + o = _prepare_output_tensor(q, out, (N, T, HV, V)) + + _dyn_kvb = ( + not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()) + and q.stride(-1) == 1 + and k.stride(-1) == 1 + and v.stride(-1) == 1 + ) + q = q if (_dyn_kvb or q.is_contiguous()) else q.contiguous() + k = k if (_dyn_kvb or k.is_contiguous()) else k.contiguous() + v = v if (_dyn_kvb or v.is_contiguous()) else v.contiguous() + a = a if a.is_contiguous() else a.contiguous() + b = b if b.is_contiguous() else b.contiguous() + + A_log = _normalize_A_log(A_log, HV) + dt_bias = _normalize_dt_bias(dt_bias, HV, K) + initial_state_indices = _normalize_state_indices(initial_state_indices, N=N, pool_size=pool_size, device=q.device) + + if write_ubuf: + if tuple(d_buffer.shape) != (N, T, HV, V): + raise ValueError(f"d_buffer shape must be {(N, T, HV, V)}, got {tuple(d_buffer.shape)}") + if tuple(k_buffer.shape) != (N, T, HV, K) or tuple(g_buffer.shape) != (N, T, HV, K): + raise ValueError(f"k_buffer/g_buffer shape must be {(N, T, HV, K)}") + d_buf, k_buf, g_buf = d_buffer, k_buffer, g_buffer + else: + d_buf = torch.empty(N, T, HV, V, dtype=torch.float32, device=q.device) + k_buf = torch.empty(N, T, HV, K, dtype=torch.float32, device=q.device) + g_buf = torch.empty(N, T, HV, K, dtype=torch.float32, device=q.device) + + stream = _get_cached_stream(q.device) + + h0_source_flat = h0_source.view(pool_size * HV, V, K) + compiled_kernel = _get_compiled_mtp_shuffle_kvbuffer_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + tile_v, + ilp_rows, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + disable_state_update=disable_state_update, + emit_output=emit_output, + write_ubuf=write_ubuf, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + opt_level=opt_level, + fast_math=fast_math, + use_lower_bound=lower_bound is not None, + lower_bound=(0.0 if lower_bound is None else float(lower_bound)), + dyn_stride=_dyn_kvb, + ) + compiled_kernel( + h0_source_flat, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + initial_state_indices, + d_buf, + k_buf, + g_buf, + stream, + ) + return o + + +# =========================================================================== +# tensor_core-kvbuffer (CuTe tensor-core, flat-in-T): every reduction on warp-level +# mma.sync.m16n8k8.tf32 (llvm.inline_asm wrapper); verify = the BT=8 stacked kernel below. +# +# mma.sync m16n8k8 fragment mapping (PTX ISA), gid = lane>>2, tig = lane&3: +# A row-major [16,8]: a0=A[gid][tig] a1=A[gid+8][tig] a2=A[gid][tig+4] a3=A[gid+8][tig+4] +# B col-major [8,8]: b0=B[tig][gid] b1=B[tig+4][gid] +# C/D [16,8] f32: c0=C[gid][2tig] c1=C[gid][2tig+1] c2=C[gid+8][2tig] c3=C[gid+8][2tig+1] +# =========================================================================== + +from cutlass._mlir.dialects import arith as _arith # noqa: E402 +from cutlass._mlir.dialects import llvm as _llvm # noqa: E402 +from cutlass.cutlass_dsl import T as _T # noqa: E402 +from cutlass.cutlass_dsl import dsl_user_op # noqa: E402 + + +@dsl_user_op +def _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """One mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32; returns (d0, d1, d2, d3). + + a*/b* are Float32 values reinterpreted as tf32 (raw f32 bits; HW ignores the low + mantissa bits — same truncation semantics as Triton's tf32 dots).""" + f32 = _T.f32() + i32 = _T.i32() + + def _bits(v): + vv = v.ir_value(loc=loc, ip=ip) if hasattr(v, "ir_value") else v + return _arith.bitcast(i32, vv, loc=loc, ip=ip) + + def _f(v): + return v.ir_value(loc=loc, ip=ip) if hasattr(v, "ir_value") else v + + res_ty = _llvm.StructType.get_literal([f32, f32, f32, f32]) + res = _llvm.inline_asm( + res_ty, + [_bits(a0), _bits(a1), _bits(a2), _bits(a3), _bits(b0), _bits(b1), _f(c0), _f(c1), _f(c2), _f(c3)], + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {$0,$1,$2,$3}, {$4,$5,$6,$7}, {$8,$9}, {$10,$11,$12,$13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(_llvm.extractvalue(f32, res, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(_llvm.extractvalue(f32, res, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(_llvm.extractvalue(f32, res, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(_llvm.extractvalue(f32, res, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def _tf32_lo(v, *, loc=None, ip=None): + """Residual v - tf32(v): the low-13-mantissa-bit part of an fp32, as Float32.""" + i32 = _T.i32() + f32 = _T.f32() + vv = v.ir_value(loc=loc, ip=ip) if hasattr(v, "ir_value") else v + bits = _arith.bitcast(i32, vv, loc=loc, ip=ip) + mask = _arith.constant(i32, -8192, loc=loc, ip=ip) # 0xFFFFE000: zero low 13 mantissa bits + hi_bits = _arith.andi(bits, mask, loc=loc, ip=ip) + hi = _arith.bitcast(f32, hi_bits, loc=loc, ip=ip) + lo = _arith.subf(vv, hi, loc=loc, ip=ip) + return cutlass.Float32(lo) + + +@dsl_user_op +def _mma_m16n8k8_3xtf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """3xTF32-emulated m16n8k8 GEMM (~fp32 accuracy). 3 tf32 MMA passes: + hi*hi + hi*lo + lo*hi, lo = x - tf32(x). ~3x the HMMA of one tf32 mma.""" + a0l = _tf32_lo(a0) + a1l = _tf32_lo(a1) + a2l = _tf32_lo(a2) + a3l = _tf32_lo(a3) + b0l = _tf32_lo(b0) + b1l = _tf32_lo(b1) + c0, c1, c2, c3 = _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3) + c0, c1, c2, c3 = _mma_m16n8k8_tf32(a0, a1, a2, a3, b0l, b1l, c0, c1, c2, c3) + c0, c1, c2, c3 = _mma_m16n8k8_tf32(a0l, a1l, a2l, a3l, b0, b1, c0, c1, c2, c3) + return c0, c1, c2, c3 + + +_compiled_tensor_core_kvbuffer_kernels: dict[tuple, object] = {} + + +def _get_compiled_tensor_core_kvbuffer_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + bv, + num_v_tiles, + scale, + use_qk_l2norm, + disable_state_update, + emit_output, + write_ubuf, + softplus_beta, + softplus_threshold, + opt_level=3, + fast_math=True, + use_lower_bound=False, + lower_bound=0.0, + dyn_stride=False, +): + key = ( + T, + H, + HV, + K, + V, + bv, + num_v_tiles, + scale, + use_qk_l2norm, + disable_state_update, + emit_output, + write_ubuf, + softplus_beta, + softplus_threshold, + opt_level, + fast_math, + use_lower_bound, + lower_bound, + dyn_stride, + ) + if key in _compiled_tensor_core_kvbuffer_kernels: + return _compiled_tensor_core_kvbuffer_kernels[key] + + q = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(N, T, H, K, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + a = torch.zeros(N, T, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, T, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, T, HV, V, dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + h0_source = torch.zeros(pool_size * HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + d_buf = torch.zeros(N, T, HV, V, dtype=torch.float32, device="cuda") + k_buf = torch.zeros(N, T, HV, K, dtype=torch.float32, device="cuda") + g_buf = torch.zeros(N, T, HV, K, dtype=torch.float32, device="cuda") + + run_fn = run_kda_mtp_tensor_core_kvbuffer_kernel + compiled_kernel = cute.compile( + run_fn, + from_dlpack(h0_source, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=h0_source.dim_order()), + from_dlpack(A_log, assumed_align=16), + from_dlpack(a, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()), + from_dlpack(dt_bias, assumed_align=16), + _dlp_qkv(q, dyn_stride), + _dlp_qkv(k, dyn_stride), + _dlp_qkv(v, dyn_stride), + from_dlpack(b, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=b.dim_order()), + from_dlpack(o, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=o.dim_order()), + from_dlpack(h0_indices, assumed_align=16).mark_layout_dynamic(), + from_dlpack(d_buf, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=d_buf.dim_order()), + from_dlpack(k_buf, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=k_buf.dim_order()), + from_dlpack(g_buf, assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=g_buf.dim_order()), + vec_size=VEC_SIZE, + BV=bv, + num_v_tiles=num_v_tiles, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + HV=HV, + T=T, + H=H, + K=K, + V=V, + use_qk_l2norm=use_qk_l2norm, + disable_state_update=disable_state_update, + emit_output=emit_output, + write_ubuf=write_ubuf, + fast_math=fast_math, + use_lower_bound=use_lower_bound, + lower_bound=lower_bound, + stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), + options=f"--enable-tvm-ffi --opt-level {opt_level}", + ) + _compiled_tensor_core_kvbuffer_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA MTP tensor_core-KVBuffer (tensor-core mma) kernel compiled: " + f"N={N}, T={T}, HV={HV}, K={K}, V={V}, BV={bv}, num_v_tiles={num_v_tiles}, opt_level={opt_level}" + ) + return compiled_kernel + + +def kda_decode_mtp_tensor_core_kvbuffer( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + disable_state_update: bool = True, + emit_output: bool = True, + d_buffer: torch.Tensor | None = None, + k_buffer: torch.Tensor | None = None, + g_buffer: torch.Tensor | None = None, + bv: int = 32, + num_v_tiles: int = -1, + opt_level: int = 3, + fast_math: bool = True, + lower_bound: float | None = None, +) -> torch.Tensor: + """KDA MTP decode — CuTe tensor-core kvbuffer VERIFY (port of the Triton gemm op).""" + N, T, H, K = q.shape + HV = v.shape[2] + V = v.shape[3] + write_ubuf = d_buffer is not None + + if scale is None: + scale = K**-0.5 + assert K == TILE_K == 128, f"tensor_core-kvbuffer requires K=128, got {K}" + assert T <= 8, f"tensor_core-kvbuffer (BT stacked) needs T<=8, got {T}" + assert bv == 32, f"tensor_core-kvbuffer (BT) requires bv=32 (one n-tile per warp), got {bv}" + assert V % bv == 0 and bv % 16 == 0, f"bv must divide V and be 16-aligned, got {bv}" + if num_v_tiles <= 0: + # auto: split V across CTAs until the grid reaches ~512 (fills H200's 132 SMs + # at small batch); producer redundancy per extra slice is negligible. + num_v_tiles = 1 + while num_v_tiles < V // bv and N * HV * num_v_tiles < 512: + num_v_tiles *= 2 + assert (V // bv) % num_v_tiles == 0, f"num_v_tiles must divide V//bv, got num_v_tiles={num_v_tiles}" + + h0_source, pool_size, _ = _normalize_state_source( + initial_state_source, + N=N, + HV=HV, + K=K, + V=V, + device=q.device, + state_layout="vk", + ) + a = _normalize_mtp_a(a, N=N, T=T, HV=HV, K=K) + if b.dim() != 3 or tuple(b.shape) != (N, T, HV): + raise ValueError(f"Unexpected b shape for MTP dense: {tuple(b.shape)}; expected {(N, T, HV)}") + o = _prepare_output_tensor(q, out, (N, T, HV, V)) + _dyn_kvb = ( + not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()) + and q.stride(-1) == 1 + and k.stride(-1) == 1 + and v.stride(-1) == 1 + ) + q = q if (_dyn_kvb or q.is_contiguous()) else q.contiguous() + k = k if (_dyn_kvb or k.is_contiguous()) else k.contiguous() + v = v if (_dyn_kvb or v.is_contiguous()) else v.contiguous() + a = a if a.is_contiguous() else a.contiguous() + b = b if b.is_contiguous() else b.contiguous() + A_log = _normalize_A_log(A_log, HV) + dt_bias = _normalize_dt_bias(dt_bias, HV, K) + initial_state_indices = _normalize_state_indices(initial_state_indices, N=N, pool_size=pool_size, device=q.device) + + if write_ubuf: + if tuple(d_buffer.shape) != (N, T, HV, V): + raise ValueError(f"d_buffer shape must be {(N, T, HV, V)}, got {tuple(d_buffer.shape)}") + if tuple(k_buffer.shape) != (N, T, HV, K) or tuple(g_buffer.shape) != (N, T, HV, K): + raise ValueError(f"k_buffer/g_buffer shape must be {(N, T, HV, K)}") + d_buf, k_buf, g_buf = d_buffer, k_buffer, g_buffer + else: + d_buf = torch.empty(N, T, HV, V, dtype=torch.float32, device=q.device) + k_buf = torch.empty(N, T, HV, K, dtype=torch.float32, device=q.device) + g_buf = torch.empty(N, T, HV, K, dtype=torch.float32, device=q.device) + + stream = _get_cached_stream(q.device) + h0_source_flat = h0_source.view(pool_size * HV, V, K) + compiled_kernel = _get_compiled_tensor_core_kvbuffer_kernel( + N, + T, + H, + HV, + K, + V, + pool_size, + bv, + num_v_tiles, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + disable_state_update=disable_state_update, + emit_output=emit_output, + write_ubuf=write_ubuf, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + opt_level=opt_level, + fast_math=fast_math, + use_lower_bound=lower_bound is not None, + lower_bound=(0.0 if lower_bound is None else float(lower_bound)), + dyn_stride=_dyn_kvb, + ) + compiled_kernel( + h0_source_flat, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + initial_state_indices, + d_buf, + k_buf, + g_buf, + stream, + ) + return o + + +# --------------------------------------------------------------------------- +# BT=8 stacked variant of the tensor_core kernel (T <= 8). mma.sync m16n8k8 has a +# hard M=16, so instead of padding tokens to 16 the spare 8 M-rows carry a +# SECOND matrix — pad waste becomes a ~2x instruction saving: +# P3: [kdec; qdec] @ kinv^T -> A (top) and P (bottom) in one GEMM chain +# P4: Neumann inverse in plain fp32 (precision); L_s is strictly-lower 8x8 so +# L_s^8 = 0 -> inv = (I+L_s)(I+L_s^2)(I+L_s^4), exactly 3 doubling steps +# P5: [kdec; qdec] @ S0^T -> Skdec + Sqdec together; u = inv @ (beta*x) on +# tensor cores; o-combine P@u in exact fp32 from SMEM (16 FMA/lane) +# Requires BV=32 (4 n-tiles = 1 per warp, keeps barriers warp-uniform). +# --------------------------------------------------------------------------- +BT = 8 + + +@cute.kernel +def kda_mtp_tensor_core_kvbuffer_kernel( + h0_source: cute.Tensor, + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + d_buf: cute.Tensor, + k_buf: cute.Tensor, + g_buf: cute.Tensor, + vec_size: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + emit_output: cutlass.Constexpr[bool], + write_ubuf: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + gid = lane_id // 4 + tig = lane_id % 4 + + num_warps: cutlass.Constexpr[int] = 4 + bidx, _, _ = cute.arch.block_idx() + i_v = bidx % num_v_tiles + tmp = bidx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + r_exp_A = cute.exp(cutlass.Float32(A_log[i_hv]), fastmath=fast_math) + + smem = cutlass.utils.SmemAllocator() + # stacked feature maps: rows 0..7 = kdec(tokens, pad-zeroed), rows 8..15 = qdec + sKQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout((2 * BT, K), stride=(K + 4, 1)), 16) + # suffix-decayed keys ksuf_t = kn_t * prod_{j>t} g_j (bounded; replaces kinv) + sKsuf = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, K), stride=(K + 8, 1)), 16) + sG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, K), stride=(K + 8, 1)), 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT,)), 16) + sBlast = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + # P3 cross-warp partial tiles: row = warp*16 + stacked-row + sPart = smem.allocate_tensor(cutlass.Float32, cute.make_layout((4 * 16, 12), stride=(12, 1)), 16) + sL = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, BT), stride=(BT + 1, 1)), 16) + sP = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, BT), stride=(BT + 1, 1)), 16) + sInv = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, BT), stride=(BT + 1, 1)), 16) + sLp = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, BT), stride=(BT + 1, 1)), 16) + sX = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, BV), stride=(BV + 1, 1)), 16) + sU = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, BV), stride=(BV + 1, 1)), 16) + sS0 = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BV, K), stride=(K + 4, 1)), 16) + + r_qbf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + r_kbf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.BFloat16) + r_qf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_kf = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_s = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + # P2a pair partials: ceil(2*T*T/4) per warp + ppw_tc: cutlass.Constexpr[int] = (2 * T * T + num_warps - 1) // num_warps + r_red = cute.make_rmem_tensor(cute.make_layout((ppw_tc,), stride=(1,)), cutlass.Float32) + + if cache_idx >= 0: + k_start = lane_id * vec_size + flat_state_idx = cache_idx * HV + i_hv + + # ---- P1: token-parallel l2norm + staging (k_norm -> sKQ top, q_scaled -> bottom) ---- + tokens_per_warp: cutlass.Constexpr[int] = (T + num_warps - 1) // num_warps + for tt in cutlass.range_constexpr(tokens_per_warp): + t_tok = tt * num_warps + warp_idx + if t_tok < T: + q_tile = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, t_tok, i_h, lane_id)) + k_tile = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, t_tok, i_h, lane_id)) + cute.autovec_copy(q_tile, r_qbf) + cute.autovec_copy(k_tile, r_kbf) + for c in cutlass.range_constexpr(vec_size): + r_qf[c] = cutlass.Float32(r_qbf[c]) + r_kf[c] = cutlass.Float32(r_kbf[c]) + if cutlass.const_expr(use_qk_l2norm): + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for c in cutlass.range_constexpr(vec_size): + sum_q += r_qf[c] * r_qf[c] + sum_k += r_kf[c] * r_kf[c] + for off in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=off, mask=-1, mask_and_clamp=31) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=off, mask=-1, mask_and_clamp=31) + inv_q = cute.rsqrt(sum_q + 1e-6, fastmath=fast_math) * scale + inv_k = cute.rsqrt(sum_k + 1e-6, fastmath=fast_math) + for c in cutlass.range_constexpr(vec_size): + r_qf[c] = r_qf[c] * inv_q + r_kf[c] = r_kf[c] * inv_k + else: + for c in cutlass.range_constexpr(vec_size): + r_qf[c] = r_qf[c] * scale + # gate g_t per channel into sG (decay applied in P2) + for c in cutlass.range_constexpr(vec_size): + x = cutlass.Float32(a[i_n, t_tok, i_hv, k_start + c]) + cutlass.Float32(dt_bias[i_hv, k_start + c]) + if cutlass.const_expr(use_lower_bound): + sigmoid_ax = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-r_exp_A * x, fastmath=fast_math)) + sG[t_tok, k_start + c] = cute.exp(lower_bound * sigmoid_ax, fastmath=fast_math) + else: + beta_x = softplus_beta * x + exp_bx = cute.exp(beta_x, fastmath=fast_math) + sp_val = (cutlass.Float32(1.0) / softplus_beta) * cute.log( + cutlass.Float32(1.0) + exp_bx, fastmath=fast_math + ) + use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) + sp_x = use_sp * sp_val + (cutlass.Float32(1.0) - use_sp) * x + sG[t_tok, k_start + c] = cute.exp( + -r_exp_A * sp_x, fastmath=fast_math + ) # g_t directly (exact prefix product in P2) + sKQ[t_tok, k_start + c] = r_kf[c] + sKQ[BT + t_tok, k_start + c] = r_qf[c] + if lane_id == 0: + sBeta[t_tok] = cutlass.Float32(1.0) / ( + cutlass.Float32(1.0) + cute.exp(-cutlass.Float32(b[i_n, t_tok, i_hv]), fastmath=fast_math) + ) + for rp in cutlass.range_constexpr(BT - T): + sKQ[T + rp, tidx] = cutlass.Float32(0.0) + sKQ[BT + T + rp, tidx] = cutlass.Float32(0.0) + sKsuf[T + rp, tidx] = cutlass.Float32(0.0) + if tidx >= T: + if tidx < BT: + sBeta[tidx] = cutlass.Float32(0.0) + cute.arch.barrier() + + # ---- P2a: T*T scores in plain fp32 with bounded decay-ratio chains. + # Runs BEFORE the prefix scaling so sKQ still holds raw kn/q_scaled: + # A[t,i] = sum_k kn_t kn_i * r(t,i), P[t,i] = sum_k qn_t kn_i * r(t,i), + # r(t,i) = prod_{i sKsuf; then forward prefix scaling + # kdec/qdec; scratch stores raw (k, g) for the bounded flush rebuild. ---- + kc = tidx # requires K == 128 == block size + suf_s = cutlass.Float32(1.0) + for tt in cutlass.range_constexpr(T): + i_t = T - 1 - tt + sKsuf[i_t, kc] = sKQ[i_t, kc] * suf_s + suf_s = suf_s * sG[i_t, kc] + bcum = cutlass.Float32(1.0) + for i_t in cutlass.range_constexpr(T): + g_t = sG[i_t, kc] + bcum = bcum * g_t + kn = sKQ[i_t, kc] + sKQ[i_t, kc] = kn * bcum + sKQ[BT + i_t, kc] = sKQ[BT + i_t, kc] * bcum + if cutlass.const_expr(write_ubuf): + if i_v == 0: + k_buf[i_n, i_t, i_hv, kc] = kn # raw key (was k/b_run) + g_buf[i_n, i_t, i_hv, kc] = g_t # per-step gate (was b_run) + sBlast[kc] = bcum + cute.arch.barrier() + if tidx < BT * BT: + ri = tidx // BT + ci = tidx % BT + one = cutlass.Float32(1.0) if ri == ci else cutlass.Float32(0.0) + sInv[ri, ci] = one # inv starts at I: each doubling step does inv += inv@Lp_old + # (with Lp_old = Ls^(2^step)), so I+Ls is produced by step 0 + sLp[ri, ci] = sL[ri, ci] + cute.arch.barrier() + + # ---- P4: doubling chain + Pinv on the 8x8 mats in PLAIN fp32 + ri = tidx // BT + ci = tidx % BT + for step in cutlass.range_constexpr(3): # 3 steps: (I+Ls)(I+Ls^2)(I+Ls^4), nilpotency 8 + if tidx < 2 * BT * BT: # rows 0..7 -> Lp@Lp, rows 8..15 -> inv@Lp + rr = ri % BT + acc = cutlass.Float32(0.0) + for l in cutlass.range_constexpr(BT): + if ri < BT: + acc += sLp[rr, l] * sLp[l, ci] + else: + acc += sInv[rr, l] * sLp[l, ci] + sPart[ri, ci] = acc + cute.arch.barrier() + if tidx < BT * BT: + sLp[ri, ci] = sPart[ri, ci] + sInv[ri, ci] = sInv[ri, ci] + sPart[BT + ri, ci] + cute.arch.barrier() + + # ---- P5 consumer. V tiled 3 ways (outer->inner): + # num_v_tiles : V split across CTAs (grid=N*HV*num_v_tiles) + # BV=32 : V rows/block = 4 warps x mma-N(8); 1 n-tile/warp, uniform barriers + # num_v_blocks : BV-blocks each CTA walks serially + num_v_blocks: cutlass.Constexpr[int] = V // BV // num_v_tiles + for vb in cutlass.range_constexpr(num_v_blocks): + v_base = (i_v * num_v_blocks + vb) * BV # global V-row start of this block + row_vecs = K // vec_size # float4s per V row + # stage S0[BV,K] -> sS0: 128 threads (blockDim), one float4 each; + # passes = BV*K / (128*vec_size) + for j in cutlass.range_constexpr(BV * K // (128 * vec_size)): + flat = j * 128 + tidx # float4-group id + s_row = flat // row_vecs # V row + s_col = flat % row_vecs # float4 within row + h_tile = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_base + s_row, s_col)) + cute.autovec_copy(h_tile, r_s) + for cc in cutlass.range_constexpr(vec_size): + sS0[s_row, s_col * vec_size + cc] = r_s[cc] + cute.arch.barrier() + + nb = warp_idx * 8 # current warp's n-tile = V rows [nb, nb+8) within the BV block + # the two adjacent V indices this lane owns (mma N-frag: 2*tig, 2*tig+1) + vc0 = nb + 2 * tig + vc1 = nb + 2 * tig + 1 + # GEMM1: [kdec; qdec] @ S0^T -> Skdec (rows 0..7) + Sqdec (rows 8..15) + e0 = cutlass.Float32(0.0) + e1 = cutlass.Float32(0.0) + e2 = cutlass.Float32(0.0) + e3 = cutlass.Float32(0.0) + for ks in cutlass.range_constexpr(K // 8): + kb = ks * 8 + a0 = sKQ[gid, kb + tig] + a1 = sKQ[gid + 8, kb + tig] + a2 = sKQ[gid, kb + tig + 4] + a3 = sKQ[gid + 8, kb + tig + 4] + b0 = sS0[nb + gid, kb + tig] + b1 = sS0[nb + gid, kb + tig + 4] + e0, e1, e2, e3 = _mma_m16n8k8_3xtf32(a0, a1, a2, a3, b0, b1, e0, e1, e2, e3) + # x = beta * (v - Skdec) from the top half; Sqdec (e2/e3) stays in registers + vmask = cutlass.Float32(1.0) if gid < T else cutlass.Float32(0.0) + vv0 = cutlass.Float32(v[i_n, gid % T, i_hv, v_base + vc0]) * vmask + vv1 = cutlass.Float32(v[i_n, gid % T, i_hv, v_base + vc1]) * vmask + sX[gid, vc0] = sBeta[gid] * (vv0 - e0) + sX[gid, vc1] = sBeta[gid] * (vv1 - e1) + cute.arch.barrier() + + # u = inv @ x in exact fp32 + f0 = cutlass.Float32(0.0) + f1 = cutlass.Float32(0.0) + for l in cutlass.range_constexpr(BT): + f0 += sInv[gid, l] * sX[l, vc0] + f1 += sInv[gid, l] * sX[l, vc1] + sU[gid, vc0] = f0 + sU[gid, vc1] = f1 + if cutlass.const_expr(write_ubuf): + if gid < T: + d_buf[i_n, gid, i_hv, v_base + vc0] = f0 + d_buf[i_n, gid, i_hv, v_base + vc1] = f1 + cute.arch.barrier() + # o = Sqdec + P@u combined in exact fp32 from sU (16 FMA/lane — removes the + # extra tf32 hop that the stacked [inv;Pinv]@x route put on the output path) + if cutlass.const_expr(emit_output): + if gid < T: + ov0 = e2 + ov1 = e3 + for l in cutlass.range_constexpr(BT): + ov0 += sP[gid, l] * sU[l, vc0] + ov1 += sP[gid, l] * sU[l, vc1] + o[(i_n, gid, i_hv, v_base + vc0)] = cutlass.BFloat16(ov0) + o[(i_n, gid, i_hv, v_base + vc1)] = cutlass.BFloat16(ov1) + + # state: S_T = b_last * S0 + u^T @ ksuf (ksuf bounded; b_last only + # rescales the S0 term), M = v rows, single k-slab + if cutlass.const_expr(not disable_state_update): + m_tiles: cutlass.Constexpr[int] = BV // 16 + pairs: cutlass.Constexpr[int] = m_tiles * (K // 8) + for pp in cutlass.range_constexpr((pairs + num_warps - 1) // num_warps): + pidx = pp * num_warps + warp_idx + if pidx < pairs: + m_t = pidx % m_tiles + n_t = pidx // m_tiles + mb = m_t * 16 + nb = n_t * 8 + g0 = cutlass.Float32(0.0) + g1 = cutlass.Float32(0.0) + g2 = cutlass.Float32(0.0) + g3 = cutlass.Float32(0.0) + a0 = sU[tig, mb + gid] + a1 = sU[tig, mb + gid + 8] + a2 = sU[tig + 4, mb + gid] + a3 = sU[tig + 4, mb + gid + 8] + b0 = sKsuf[tig, nb + gid] + b1 = sKsuf[tig + 4, nb + gid] + # 3xTF32 for near-fp32 state precision; only the dsu=0 + # path hits this GEMM (serving verify commits via flush). + g0, g1, g2, g3 = _mma_m16n8k8_3xtf32(a0, a1, a2, a3, b0, b1, g0, g1, g2, g3) + for fi in cutlass.range_constexpr(4): + vrow = mb + gid + (fi // 2) * 8 + kcol = nb + 2 * tig + (fi % 2) + gv = g0 + if cutlass.const_expr(fi == 1): + gv = g1 + if cutlass.const_expr(fi == 2): + gv = g2 + if cutlass.const_expr(fi == 3): + gv = g3 + h0_source[(flat_state_idx, v_base + vrow, kcol)] = sBlast[kcol] * sS0[vrow, kcol] + gv + cute.arch.barrier() + + +@cute.jit +def run_kda_mtp_tensor_core_kvbuffer_kernel( + h0_source: cute.Tensor, + A_log: cute.Tensor, + a: cute.Tensor, + dt_bias: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + d_buf: cute.Tensor, + k_buf: cute.Tensor, + g_buf: cute.Tensor, + vec_size: cutlass.Constexpr[int], + BV: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + softplus_beta: cutlass.Constexpr[float], + softplus_threshold: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + use_qk_l2norm: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + emit_output: cutlass.Constexpr[bool], + write_ubuf: cutlass.Constexpr[bool], + fast_math: cutlass.Constexpr[bool], + use_lower_bound: cutlass.Constexpr[bool], + lower_bound: cutlass.Constexpr[float], + stream: cuda.CUstream, +): + """BT=8 stacked tensor_core launcher: grid = N*HV*num_v_tiles, block = 128.""" + n_indices = h0_indices.layout.shape[0] + grid_size = n_indices * HV * num_v_tiles + smem_bytes = ( + 2 * 4 * BT * (K + 8) # sKQ (stacked) + + 2 * 4 * BT * (K + 8) # sKsuf + sG + + 4 * BT + + 4 * K # sBeta + sBlast + + 4 * 64 * 12 # sPart + + 4 * 4 * BT * (BT + 1) # sL/sP/sInv/sLp + + 2 * 4 * BT * (BV + 1) # sX/sU + + 4 * BV * (K + 8) # sS0 + + 512 + ) + kda_mtp_tensor_core_kvbuffer_kernel( + h0_source, + A_log, + a, + dt_bias, + q, + k, + v, + b, + o, + h0_indices, + d_buf, + k_buf, + g_buf, + vec_size, + BV, + num_v_tiles, + softplus_beta, + softplus_threshold, + scale, + HV, + T, + H, + K, + V, + use_qk_l2norm, + disable_state_update, + emit_output, + write_ubuf, + fast_math, + use_lower_bound, + lower_bound, + ).launch(grid=(grid_size, 1, 1), block=[128, 1, 1], smem=smem_bytes, stream=stream) + + +# --------------------------------------------------------------------------- +# KVBuffer verify dispatch: route between the two kvbuffer verify ops by T. +# --------------------------------------------------------------------------- +def _select_kvb_variant(N: int, HV: int, T: int) -> str: + """Pick "shuffle" or "tensor_core" kvbuffer variant; wu = N*HV.""" + wu = N * HV + if T <= 2: + return "shuffle" + if T == 3: + return "shuffle" if wu <= 64 else "tensor_core" + if T == 4: + return "shuffle" if wu <= 32 else "tensor_core" + return "tensor_core" + + +def _kvbuffer_prefer_tensor_core(N: int, HV: int, T: int) -> bool: + """True iff the kvbuffer dispatch picks tensor_core.""" + return _select_kvb_variant(N, HV, T) == "tensor_core" + + +def kda_decode_mtp_kvbuffer( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + out: torch.Tensor | None = None, + disable_state_update: bool = True, + emit_output: bool = True, + d_buffer: torch.Tensor | None = None, + k_buffer: torch.Tensor | None = None, + g_buffer: torch.Tensor | None = None, + t_crossover: int | None = None, + opt_level: int = 3, + fast_math: bool = True, + lower_bound: float | None = None, +) -> torch.Tensor: + """KDA MTP KVBuffer verify dispatch between shuffle-kvbuffer (token-parallel SIMT) and + tensor_core-kvbuffer (CuTe tensor-core GEMM, flat-in-T). With ``t_crossover=None`` + (default) the choice follows the kernel-level chain bench via + ``_kvbuffer_prefer_tensor_core`` (a function of the work size S = HV*N and T); pass an + int to force the legacy T-only rule (tensor_core iff T >= t_crossover). Routes only + among kvbuffer ops; the recurrent fallback is a higher-layer concern. + """ + T = q.shape[1] + N = q.shape[0] + HV = v.shape[2] + common = dict( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + out=out, + disable_state_update=disable_state_update, + emit_output=emit_output, + d_buffer=d_buffer, + k_buffer=k_buffer, + g_buffer=g_buffer, + opt_level=opt_level, + fast_math=fast_math, + lower_bound=lower_bound, + ) + if t_crossover is None: + use_tensor_core = _select_kvb_variant(N, HV, T) == "tensor_core" + else: + use_tensor_core = t_crossover <= T + if use_tensor_core: + return kda_decode_mtp_tensor_core_kvbuffer(**common) + return kda_decode_mtp_shuffle_kvbuffer(**common) diff --git a/tests/test_kda_decode_mtp.py b/tests/test_kda_decode_mtp.py new file mode 100644 index 00000000..b037a615 --- /dev/null +++ b/tests/test_kda_decode_mtp.py @@ -0,0 +1,837 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import pathlib +import sys + +import pytest +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) # for sibling test import + +from test_kda_decode import torch_kda_decode_ref # trusted single-token reference + +from cula.kda import kda_decode +from cula.ops.kda.decode.mtp import ( + _select_mtp_config, + _select_mtp_tile_v, + kda_decode_mtp_recurrent, + kda_decode_mtp_recurrent_ws, +) +from cula.ops.kda.decode.mtp_kvbuffer import ( + _kvbuffer_prefer_tensor_core, + _select_kvb_tile_v, + _select_shuffle_kvb_ilp_rows, + kda_decode_mtp_kvbuffer, + kda_decode_mtp_shuffle_kvbuffer, + kda_decode_mtp_tensor_core_kvbuffer, + kda_flush_kvbuffer, +) + + +def torch_kda_mtp_ref( + q, k, v, a, b, A_log, dt_bias, state, scale, use_l2norm=True, softplus_beta=1.0, softplus_threshold=20.0, lower_bound=None +): + """fp32 ground truth: the single-token KDA recurrence threaded over T. Returns (o, final_state).""" + N, T, HV, V = v.shape + H = q.shape[2] + heads_per_group = HV // H + A = torch.exp(A_log) + state_cur = state.clone() + o = torch.zeros(N, T, HV, V, dtype=torch.float32, device=q.device) + for t in range(T): + for n in range(N): + for hv in range(HV): + i_h = hv // heads_per_group + x = a[n, t, hv, :] + dt_bias[hv, :] + if lower_bound is not None: + # safe gate: g = lower_bound * sigmoid(exp(A_log) * x) + gate = torch.exp(lower_bound * torch.sigmoid(A[hv] * x)) + else: + sp = F.softplus(x, beta=softplus_beta, threshold=softplus_threshold) + gate = torch.exp(-A[hv] * sp) + if use_l2norm: + q_vec = F.normalize(q[n, t, i_h, :], dim=0) * scale + k_vec = F.normalize(k[n, t, i_h, :], dim=0) + else: + q_vec = q[n, t, i_h, :] * scale + k_vec = k[n, t, i_h, :] + Hk = state_cur[n, hv] @ (gate * k_vec) + beta_val = torch.sigmoid(b[n, t, hv]) + v_new = beta_val * (v[n, t, hv, :] - Hk) + state_cur[n, hv] = gate[None, :] * state_cur[n, hv] + v_new[:, None] * k_vec[None, :] + o[n, t, hv, :] = state_cur[n, hv] @ q_vec + return o, state_cur + + +def make_inputs_mtp(N, T, H, HV, K, V, device="cuda", seed=42): + """Random MTP inputs (q/k/v/a/b bf16, A_log/dt_bias/state fp32).""" + torch.manual_seed(seed) + q = torch.randn(N, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn(N, T, H, K, device=device, dtype=torch.bfloat16) + v = torch.randn(N, T, HV, V, device=device, dtype=torch.bfloat16) + a = (torch.randn(N, T, HV, K, device=device, dtype=torch.float32) * 0.1).to(torch.bfloat16) + b = torch.randn(N, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) * 2 # negative -> A < 1 + dt_bias = torch.randn(HV, K, device=device, dtype=torch.float32) * 0.1 + state = torch.randn(N, HV, V, K, device=device, dtype=torch.float32) * 0.01 + return q, k, v, a, b, A_log, dt_bias, state + + +def run_kda_decode_mtp_via_loop_dense(q, k, v, a, b, A_log, dt_bias, state, scale): + """The "loop" baseline: T sequential single-token kda_decode calls, state carried across tokens.""" + N, T, H, K = q.shape + HV, V = v.shape[2], v.shape[3] + state_source = state.clone().contiguous() + indices = torch.arange(N, device=q.device, dtype=torch.int32) + o_all = torch.empty(N, T, HV, V, device=q.device, dtype=torch.bfloat16) + for t in range(T): + q_t = q[:, t].unsqueeze(1).contiguous() + k_t = k[:, t].unsqueeze(1).contiguous() + v_t = v[:, t].unsqueeze(1).contiguous() + a_t = a[:, t].unsqueeze(1).contiguous() + b_t = b[:, t].unsqueeze(1).contiguous() + o_t = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q_t.to(torch.bfloat16), + k=k_t.to(torch.bfloat16), + v=v_t.to(torch.bfloat16), + a=a_t.to(torch.bfloat16), + b=b_t.to(torch.bfloat16), + initial_state_source=state_source, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + o_all[:, t] = o_t.squeeze(1) + return o_all, state_source + + +def _assert_close(name, ref, actual, atol=3e-2, rtol=2e-2): + """allclose, printing the observed max/mean margin (pytest -s).""" + diff = (ref.float() - actual.float()).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + print(f" [{name}] max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f} (atol={atol}, rtol={rtol})") + ok = torch.allclose(ref.float(), actual.float(), atol=atol, rtol=rtol) + assert ok, f"{name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}, atol={atol}, rtol={rtol}" + + +def oracle_intermediate_states(q, k, v, a, b, A_log, dt_bias, state, scale): + """fp32 per-token state snapshots [N,T,HV,V,K] from the trusted single-token reference.""" + N, T = q.shape[0], q.shape[1] + HV, V, K = v.shape[2], v.shape[3], q.shape[3] + state_cur = state.clone() + inter = torch.zeros(N, T, HV, V, K, dtype=torch.float32, device=q.device) + for t in range(T): + _, state_cur = torch_kda_decode_ref( + q[:, t].float(), + k[:, t].float(), + v[:, t].float(), + a[:, t], + b[:, t].float(), + A_log, + dt_bias, + state_cur, + scale, + ) + inter[:, t] = state_cur + return inter + + +def run_recurrent( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, + *, + variant, + bv=-1, + k_split=-1, + disable_state_update=False, + intermediate=False, + lower_bound=None, +): + """Run kda_decode_mtp_recurrent; state fed/returned in vk layout (kv transposed in and back).""" + N = q.shape[0] + indices = torch.arange(N, device=q.device, dtype=torch.int32) + T = q.shape[1] + HV, V, K = v.shape[2], v.shape[3], q.shape[3] + inter = torch.zeros(N, T, HV, V, K, device=q.device, dtype=torch.float32) if intermediate else None + st = state.clone().contiguous() + if variant == "kv": + st = st.transpose(-2, -1).contiguous() # vk -> kv + rec_kwargs = dict( + A_log=A_log, + dt_bias=dt_bias, + q=q.to(torch.bfloat16), + k=k.to(torch.bfloat16), + v=v.to(torch.bfloat16), + a=a.to(torch.bfloat16), + b=b.to(torch.bfloat16), + initial_state_source=st, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + variant=variant, + k_split=k_split, + disable_state_update=disable_state_update, + intermediate_states_buffer=inter, + lower_bound=lower_bound, + ) + if variant == "vk": + rec_kwargs["bv"] = bv # kv is fixed 1-warp; bv stays at the WARP_BV default + o = kda_decode_mtp_recurrent(**rec_kwargs) + state_vk = st.transpose(-2, -1).contiguous() if variant == "kv" else st + return (o, state_vk, inter) if intermediate else (o, state_vk) + + +@pytest.mark.parametrize("T", [1, 2, 4, 8]) +def test_mtp_ref_is_threaded_single_token(T): + """Pure-torch: the MTP oracle equals the trusted single-token ref threaded over T.""" + N, H, HV, K, V = 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + o_mtp, st_mtp = torch_kda_mtp_ref(q.float(), k.float(), v.float(), a, b.float(), A_log, dt_bias, state.clone(), scale) + st_cur = state.clone() + o_manual = torch.zeros(N, T, HV, V, dtype=torch.float32, device=q.device) + for t in range(T): + o_t, st_cur = torch_kda_decode_ref( + q[:, t].float(), k[:, t].float(), v[:, t].float(), a[:, t], b[:, t].float(), A_log, dt_bias, st_cur, scale + ) + o_manual[:, t] = o_t + torch.testing.assert_close(o_mtp, o_manual, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(st_mtp, st_cur, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("zero_state", [False, True], ids=["randstate", "zerostate"]) +@pytest.mark.parametrize( + "N,T,H,HV", + [ + pytest.param(*c, id="N{}-T{}-H{}-HV{}".format(*c)) + for c in [(1, 1, 8, 16), (4, 4, 8, 16), (16, 8, 8, 16), (64, 2, 16, 32), (4, 4, 16, 32)] + ], +) +def test_oracle_vs_loop(N, T, H, HV, zero_state): + """The looped single-token kernel matches the fp32 oracle (small N).""" + K, V = 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + if zero_state: + state = torch.zeros_like(state) + o_ref, st_ref = torch_kda_mtp_ref(q.float(), k.float(), v.float(), a, b.float(), A_log, dt_bias, state.clone(), scale) + o_loop, st_loop = run_kda_decode_mtp_via_loop_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + _assert_close("loop output", o_ref, o_loop.float()) + _assert_close("loop final state", st_ref, st_loop) + + +@pytest.mark.parametrize( + "N,T,H,HV,variant,bv,k_split", + [ + pytest.param(*c, id="N{}-T{}-H{}-HV{}-{}-bv{}-ks{}".format(*c)) + for c in [ + # vk: bv sweep + auto, incl T=1 and GQA + (1, 1, 8, 16, "vk", -1, 1), + (4, 4, 8, 16, "vk", -1, 1), + (8, 2, 8, 16, "vk", -1, 1), + (4, 4, 8, 16, "vk", 8, 1), + (4, 4, 8, 16, "vk", 16, 1), + (4, 2, 8, 16, "vk", 32, 1), + (16, 4, 16, 32, "vk", -1, 1), + # kv: k_split sweep + auto, incl T=1 and GQA + (1, 1, 8, 16, "kv", 32, -1), + (4, 4, 8, 16, "kv", 32, -1), + (8, 2, 8, 16, "kv", 32, -1), + (4, 4, 8, 16, "kv", 32, 1), + (4, 4, 8, 16, "kv", 32, 2), + (4, 4, 8, 16, "kv", 32, 4), + (16, 4, 16, 32, "kv", 32, -1), + ] + ], +) +def test_recurrent_decode(N, T, H, HV, variant, bv, k_split): + """recurrent vk + kv vs loop: bv / k_split / auto / GQA in one table.""" + K, V = 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + o_loop, st_loop = run_kda_decode_mtp_via_loop_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + o_sb, st_sb = run_recurrent(q, k, v, a, b, A_log, dt_bias, state, scale, variant=variant, bv=bv, k_split=k_split) + tag = f"recurrent {variant} bv={bv} ks={k_split}" + _assert_close(f"{tag} output", o_loop.float(), o_sb.float()) + _assert_close(f"{tag} final state", st_loop, st_sb) + + +@pytest.mark.parametrize( + "kernel", ["recurrent_ws", "recurrent_ws_ilp4", "recurrent_ws_smem_v", "recurrent_vk", "recurrent_kv"] +) +@pytest.mark.parametrize( + "N,T,H,HV", + [ + pytest.param(*c, id="N{}-T{}-H{}-HV{}".format(*c)) + for c in [ + (1, 1, 8, 16), + (4, 4, 8, 16), + (8, 4, 8, 16), + (16, 4, 16, 32), + ] + ], +) +def test_lower_bound_safe_gate(kernel, N, T, H, HV): + """Safe-gate path g = lower_bound * sigmoid(exp(A_log) * x): the MTP kernels must + match the fp32 oracle (the single-token loop kernel has no safe-gate path).""" + K, V = 128, 128 + scale = K**-0.5 + lower_bound = -4.0 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + o_ref, st_ref = torch_kda_mtp_ref( + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state.clone(), + scale, + lower_bound=lower_bound, + ) + if kernel == "recurrent_ws": + o, st = run_recurrent_ws(q, k, v, a, b, A_log, dt_bias, state, scale, lower_bound=lower_bound) + elif kernel == "recurrent_ws_ilp4": + o, st = run_recurrent_ws(q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=16, ilp_rows=4, lower_bound=lower_bound) + elif kernel == "recurrent_ws_smem_v": + o, st = run_recurrent_ws( + q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=32, ilp_rows=4, use_smem_v=True, lower_bound=lower_bound + ) + elif kernel == "recurrent_vk": + o, st = run_recurrent(q, k, v, a, b, A_log, dt_bias, state, scale, variant="vk", lower_bound=lower_bound) + else: # recurrent_kv + o, st = run_recurrent(q, k, v, a, b, A_log, dt_bias, state, scale, variant="kv", lower_bound=lower_bound) + tag = f"lb {kernel} N={N} T={T} HV={HV}" + _assert_close(f"{tag} output", o_ref, o.float()) + _assert_close(f"{tag} final state", st_ref, st) + + +@pytest.mark.parametrize("kernel", ["recurrent_ws", "recurrent_ws_ilp4", "recurrent_vk", "recurrent_kv"]) +def test_disable_state_update(kernel): + """disable_state_update leaves the state pool unchanged while output still matches the loop.""" + N, T, H, HV, K, V = 4, 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + o_loop, _ = run_kda_decode_mtp_via_loop_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + + if kernel == "recurrent_ws": + o, st = run_recurrent_ws(q, k, v, a, b, A_log, dt_bias, state, scale, disable_state_update=True) + elif kernel == "recurrent_ws_ilp4": + o, st = run_recurrent_ws(q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=32, ilp_rows=4, disable_state_update=True) + else: + variant = "vk" if kernel == "recurrent_vk" else "kv" + o, st = run_recurrent(q, k, v, a, b, A_log, dt_bias, state, scale, variant=variant, disable_state_update=True) + + assert torch.equal(st, state), f"{kernel}: state pool modified despite disable_state_update=True" + _assert_close(f"{kernel} dsu output", o_loop.float(), o.float()) + + +@pytest.mark.parametrize("kernel", ["recurrent_ws", "recurrent_ws_smem_v", "recurrent_vk", "recurrent_kv"]) +def test_determinism(kernel): + """Bit-exact determinism: repeat the state-writeback launch, assert identical output + state.""" + N, T, H, HV, K, V = 16, 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + + def launch(): + if kernel == "recurrent_ws": + return run_recurrent_ws(q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=64, ilp_rows=4, use_packed_fma=False) + if kernel == "recurrent_ws_smem_v": + return run_recurrent_ws( + q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=64, ilp_rows=4, use_packed_fma=False, use_smem_v=True + ) + variant = "vk" if kernel == "recurrent_vk" else "kv" + return run_recurrent(q, k, v, a, b, A_log, dt_bias, state, scale, variant=variant) + + o_ref, st_ref = launch() + o_ref = o_ref.clone() + n_iters = int(os.environ.get("KDA_MTP_DET_ITERS", "100000")) + for i in range(n_iters): + o_i, st_i = launch() + assert torch.equal(o_i, o_ref), f"{kernel} output non-deterministic at iter {i}" + assert torch.equal(st_i, st_ref), f"{kernel} state non-deterministic at iter {i}" + + +def test_intermediate_disable_state_update(): + """disable_state_update leaves the pool untouched; snapshots still fire and match the oracle.""" + N, T, H, HV, K, V = 4, 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + inter_ref = oracle_intermediate_states(q, k, v, a, b, A_log, dt_bias, state.clone(), scale) + + _o, st_vk, inter = run_recurrent( + q, k, v, a, b, A_log, dt_bias, state, scale, variant="vk", disable_state_update=True, intermediate=True + ) + assert torch.equal(st_vk, state), "pool modified despite disable_state_update=True" + for t in range(T): + _assert_close(f"inter+dsu snapshot[t={t}]", inter_ref[:, t], inter[:, t]) + + +def test_intermediate_buffer_validation(): + """Bad intermediate_states_buffer shape / dtype must raise.""" + N, T, H, HV, K, V = 4, 2, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + st = state.clone().contiguous() + indices = torch.arange(N, device=q.device, dtype=torch.int32) + + def _call(buf): + return kda_decode_mtp_recurrent( + A_log=A_log, + dt_bias=dt_bias, + q=q.to(torch.bfloat16), + k=k.to(torch.bfloat16), + v=v.to(torch.bfloat16), + a=a.to(torch.bfloat16), + b=b.to(torch.bfloat16), + initial_state_source=st, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + variant="vk", + intermediate_states_buffer=buf, + ) + + with pytest.raises((ValueError, AssertionError)): + _call(torch.zeros(N, T + 1, HV, V, K, device="cuda", dtype=torch.float32)) + with pytest.raises((ValueError, AssertionError)): + _call(torch.zeros(N, T, HV, V, K, device="cuda", dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("N,T", [(1, 2), (4, 4), (8, 8), (4, 2), (16, 6)]) +def test_intermediate_recurrent_vk(N, T): + """vk per-token snapshot == fp32 oracle; t=T-1 snapshot == final state pool.""" + H, HV, K, V = 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + inter_ref = oracle_intermediate_states(q, k, v, a, b, A_log, dt_bias, state.clone(), scale) + o, st_vk, inter = run_recurrent( + q, k, v, a, b, A_log, dt_bias, state.clone(), scale, variant="vk", disable_state_update=False, intermediate=True + ) + for t in range(T): + _assert_close(f"sbvk inter snapshot[t={t}]", inter_ref[:, t], inter[:, t]) + assert torch.equal(inter[:, T - 1], st_vk), "sbvk: t=T-1 snapshot != final state" + + +def run_recurrent_ws( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, + *, + tile_v=None, + ilp_rows=None, + use_packed_fma=None, + use_smem_v=None, + disable_state_update=False, + intermediate=False, + lower_bound=None, +): + """Run kda_decode_mtp_recurrent_ws (vk). Returns (o, state) or (o, state, inter).""" + N, T, _, K = q.shape + HV, V = v.shape[2], v.shape[3] + st = state.clone().contiguous() + indices = torch.arange(N, device=q.device, dtype=torch.int32) + inter = torch.zeros(N, T, HV, V, K, device=q.device, dtype=torch.float32) if intermediate else None + o = kda_decode_mtp_recurrent_ws( + A_log=A_log, + dt_bias=dt_bias, + q=q.to(torch.bfloat16), + k=k.to(torch.bfloat16), + v=v.to(torch.bfloat16), + a=a.to(torch.bfloat16), + b=b.to(torch.bfloat16), + initial_state_source=st, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + tile_v=tile_v, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + use_smem_v=use_smem_v, + disable_state_update=disable_state_update, + intermediate_states_buffer=inter, + lower_bound=lower_bound, + ) + return (o, st, inter) if intermediate else (o, st) + + +@pytest.mark.parametrize( + "N,T,H,HV,tile_v,ilp_rows,use_smem_v", + [ + pytest.param(*c, id="N{}-T{}-H{}-HV{}-tv{}-ilp{}-smem{}".format(*c)) + for c in [ + # auto (None) across N incl GQA and large batch + (1, 2, 8, 16, None, None, None), + (4, 4, 8, 16, None, None, None), + (16, 4, 16, 32, None, None, None), + (64, 8, 8, 16, None, None, None), + (1024, 2, 8, 16, None, None, None), + (2048, 2, 8, 16, None, None, None), + # explicit tile_v sweep, ilp=2 + (4, 4, 8, 16, 8, 2, False), + (4, 4, 8, 16, 16, 2, False), + (4, 4, 8, 16, 32, 2, False), + (4, 2, 8, 16, 64, 2, False), + # ilp=4 (tile_v % 16 == 0), fused steps + double-accumulator + (4, 4, 8, 16, 16, 4, False), + (4, 4, 8, 16, 32, 4, False), + (4, 2, 8, 16, 64, 4, False), + # use_smem_v on + (4, 4, 8, 16, 32, 4, True), + (16, 2, 16, 32, 64, 4, True), + ] + ], +) +def test_recurrent_ws_decode(N, T, H, HV, tile_v, ilp_rows, use_smem_v): + """ws warp-spec vs loop: auto / tile_v / ilp 2,4 / use_smem_v / large N in one table.""" + K, V = 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + o_loop, st_loop = run_kda_decode_mtp_via_loop_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + o_ws, st_ws = run_recurrent_ws( + q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=tile_v, ilp_rows=ilp_rows, use_smem_v=use_smem_v + ) + tag = f"ws tv={tile_v} ilp={ilp_rows} smem={use_smem_v}" + _assert_close(f"{tag} output", o_loop.float(), o_ws.float()) + _assert_close(f"{tag} final state", st_loop, st_ws) + + +@pytest.mark.parametrize("tile_v,ilp_rows", [(8, 2), (16, 2), (32, 2), (64, 2), (16, 4), (32, 4), (64, 4)]) +def test_recurrent_ws_smem_v_bit_identical(tile_v, ilp_rows): + """use_smem_v is pure data movement: byte-for-byte identical to the GMEM path.""" + N, T, H, HV, K, V = 4, 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + o_g, st_g = run_recurrent_ws( + q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=tile_v, ilp_rows=ilp_rows, use_packed_fma=False, use_smem_v=False + ) + o_s, st_s = run_recurrent_ws( + q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=tile_v, ilp_rows=ilp_rows, use_packed_fma=False, use_smem_v=True + ) + assert torch.equal(o_s, o_g), f"smem_v output != GMEM (tile_v={tile_v}, ilp={ilp_rows})" + assert torch.equal(st_s, st_g), f"smem_v state != GMEM (tile_v={tile_v}, ilp={ilp_rows})" + + +def test_recurrent_ws_ilp4_rejects_bad_tile_v(): + """ilp=4 requires tile_v % 16 == 0; tile_v=8 must raise.""" + N, T, H, HV, K, V = 4, 2, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + with pytest.raises(AssertionError): + run_recurrent_ws(q, k, v, a, b, A_log, dt_bias, state, scale, tile_v=8, ilp_rows=4, use_packed_fma=False) + + +@pytest.mark.parametrize( + "N,HV,V,T,expected", + [ + (1, 16, 128, 2, (8, 2, False)), + (4, 16, 128, 4, (8, 2, False)), + (1, 65, 128, 2, (16, 4, False)), + (8, 16, 128, 2, (16, 4, False)), + (16, 16, 128, 2, (16, 2, False)), + (16, 16, 128, 4, (32, 4, False)), + (7, 64, 128, 2, (16, 2, False)), + (7, 64, 128, 8, (32, 4, False)), + (16, 64, 128, 2, (32, 4, False)), + (64, 16, 128, 8, (32, 4, False)), + (17, 64, 128, 2, (64, 4, True)), + (256, 64, 128, 8, (64, 4, True)), + (8, 16, 8, 2, (8, 2, False)), + (8, 16, 16, 2, (16, 4, False)), + ], +) +def test_select_mtp_config(N, HV, V, T, expected): + """The joint (tile_v, ilp_rows, use_smem_v) heuristic returns the expected config.""" + assert _select_mtp_config(N, HV, V, T) == expected + assert _select_mtp_tile_v(N, HV, V, T) == expected[0] + + +def test_select_mtp_config_ilp_capped_at_4(): + """ilp is capped at 4 (no ilp=8 path) in every bucket.""" + for N in (1, 8, 16, 64, 256, 4096): + for HV in (16, 64): + for T in (1, 2, 4, 8): + for dsu in (False, True): + _, ilp, _ = _select_mtp_config(N, HV, 128, T, disable_state_update=dsu) + assert ilp in (2, 4), f"N={N},HV={HV},T={T},dsu={dsu} -> ilp={ilp}" + + +@pytest.mark.parametrize("use_smem_v", [False, True]) +@pytest.mark.parametrize("tile_v,ilp_rows", [(16, 2), (32, 4), (64, 4)]) +def test_intermediate_vs_oracle_and_final(use_smem_v, tile_v, ilp_rows): + """Each per-token snapshot == fp32 oracle state; the t=T-1 snapshot == final state pool.""" + N, T, H, HV, K, V = 4, 4, 8, 16, 128, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K, V) + inter_ref = oracle_intermediate_states(q, k, v, a, b, A_log, dt_bias, state.clone(), scale) + _o, st_final, inter = run_recurrent_ws( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state, + scale, + tile_v=tile_v, + ilp_rows=ilp_rows, + use_packed_fma=False, + use_smem_v=use_smem_v, + intermediate=True, + ) + tag = f"inter smem={use_smem_v} tv={tile_v} ilp={ilp_rows}" + for t in range(T): + _assert_close(f"{tag} snapshot[t={t}]", inter_ref[:, t], inter[:, t]) + assert torch.equal(inter[:, T - 1], st_final), f"{tag}: t=T-1 snapshot != final state pool" + + +K_DIM = 128 # kvbuffer ops hard-require K=128 + + +def _alloc_ubufs(N, T, HV, V, device="cuda"): + """d_buffer [N,T,HV,V], k/g_buffer [N,T,HV,K] — fp32, matching the kernel contract.""" + return ( + torch.zeros(N, T, HV, V, dtype=torch.float32, device=device), + torch.zeros(N, T, HV, K_DIM, dtype=torch.float32, device=device), + torch.zeros(N, T, HV, K_DIM, dtype=torch.float32, device=device), + ) + + +def _kvb_verify(which, q, k, v, a, b, A_log, dt_bias, state, scale, *, ubufs=None, lower_bound=None): + """Run a kvbuffer verify op (disable_state_update=True). Returns output o [N,T,HV,V].""" + N = q.shape[0] + indices = torch.arange(N, device=q.device, dtype=torch.int32) + d_b, k_b, g_b = ubufs if ubufs is not None else (None, None, None) + op = kda_decode_mtp_shuffle_kvbuffer if which == "shuffle" else kda_decode_mtp_tensor_core_kvbuffer + return op( + A_log=A_log, + dt_bias=dt_bias, + q=q.to(torch.bfloat16), + k=k.to(torch.bfloat16), + v=v.to(torch.bfloat16), + a=a.to(torch.bfloat16), + b=b.to(torch.bfloat16), + initial_state_source=state.clone().contiguous(), + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + emit_output=True, + d_buffer=d_b, + k_buffer=k_b, + g_buffer=g_b, + lower_bound=lower_bound, + ) + + +def _kvb_oracle_out(q, k, v, a, b, A_log, dt_bias, state, scale): + o_ref, _ = torch_kda_mtp_ref( + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state, + scale, + ) + return o_ref + + +def _check_kvb_verify_and_flush(which, N, T, H, HV): + """verify output == oracle, u-buffer populated; flush(m) == m-th oracle snapshot (m=full/half/one).""" + V = K_DIM + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K_DIM, V) + scale = K_DIM**-0.5 + o_ref = _kvb_oracle_out(q, k, v, a, b, A_log, dt_bias, state, scale) + inter_ref = oracle_intermediate_states(q, k, v, a, b, A_log, dt_bias, state.clone(), scale) + + indices = torch.arange(N, device=q.device, dtype=torch.int32) + ubufs = _alloc_ubufs(N, T, HV, V) + o = _kvb_verify(which, q, k, v, a, b, A_log, dt_bias, state, scale, ubufs=ubufs) + _assert_close(f"{which}_verify N{N}T{T}", o_ref, o) + assert ubufs[0].abs().sum() > 0, f"{which}: d_buffer was not written" + + # flush each accept length m -> rebuilt S_m == oracle state after m tokens (snapshot m-1) + for m in sorted({T, max(1, T // 2), 1}): + pool = state.clone().contiguous() + kda_flush_kvbuffer(pool, indices, ubufs[0], ubufs[1], ubufs[2], accept_len=m) + _assert_close(f"{which}_flush N{N}T{T}m{m}", inter_ref[:, m - 1], pool) + + +@pytest.mark.parametrize("N,T,H,HV", [(2, 2, 16, 16), (4, 4, 16, 16), (2, 4, 32, 32)]) +def test_shuffle_kvbuffer_verify_and_flush(N, T, H, HV): + """shuffle-kvbuffer (token-parallel SIMT) verify output + rank-m flush match the fp32 oracle.""" + _check_kvb_verify_and_flush("shuffle", N, T, H, HV) + + +@pytest.mark.parametrize("N,T,H,HV", [(2, 3, 16, 16), (4, 6, 16, 16), (1, 8, 32, 32)]) +def test_tensor_core_kvbuffer_verify_and_flush(N, T, H, HV): + """tensor_core-kvbuffer (CuTe tensor-core gemm) verify output + rank-m flush match the fp32 oracle.""" + _check_kvb_verify_and_flush("tensor_core", N, T, H, HV) + + +@pytest.mark.parametrize( + "which,N,T,H,HV", + [("shuffle", 2, 2, 16, 16), ("shuffle", 4, 2, 16, 16), ("tensor_core", 2, 4, 16, 16), ("tensor_core", 1, 8, 32, 32)], +) +def test_lower_bound_kvbuffer(which, N, T, H, HV): + """kvbuffer (shuffle/tensor_core) safe-gate path: verify output matches the fp32 oracle with lower_bound.""" + V = K_DIM + scale = K_DIM**-0.5 + lower_bound = -4.0 + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K_DIM, V) + o_ref, _ = torch_kda_mtp_ref( + q.float(), + k.float(), + v.float(), + a, + b.float(), + A_log, + dt_bias, + state.clone(), + scale, + lower_bound=lower_bound, + ) + ubufs = _alloc_ubufs(N, T, HV, V) + o = _kvb_verify(which, q, k, v, a, b, A_log, dt_bias, state, scale, ubufs=ubufs, lower_bound=lower_bound) + _assert_close(f"lb {which} N{N}T{T}HV{HV}", o_ref, o) + + +@pytest.mark.parametrize( + "N,HV,T,routed", + [(2, 16, 2, "shuffle"), (8, 16, 4, "tensor_core")], # S=HV*N: 32 -> shuffle @T2 ; 128 -> tensor_core @T4 +) +def test_kvbuffer_dispatch_output_matches_oracle(N, HV, T, routed): + """kda_decode_mtp_kvbuffer auto-dispatch (S=HV*N + T rule) routes as expected and the + output matches the oracle whichever kvbuffer kernel it picks.""" + H, V = HV, K_DIM + assert _kvbuffer_prefer_tensor_core(N, HV, T) is (routed == "tensor_core") + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K_DIM, V) + scale = K_DIM**-0.5 + o_ref = _kvb_oracle_out(q, k, v, a, b, A_log, dt_bias, state, scale) + indices = torch.arange(N, device=q.device, dtype=torch.int32) + o = kda_decode_mtp_kvbuffer( + A_log=A_log, + dt_bias=dt_bias, + q=q.to(torch.bfloat16), + k=k.to(torch.bfloat16), + v=v.to(torch.bfloat16), + a=a.to(torch.bfloat16), + b=b.to(torch.bfloat16), + initial_state_source=state.clone().contiguous(), + initial_state_indices=indices, + scale=scale, + ) + _assert_close(f"dispatch N{N} HV{HV} T{T}->{routed}", o_ref, o) + + +def test_kvbuffer_prefer_tensor_core_matches_bench(): + """_kvbuffer_prefer_tensor_core reproduces the kvbuffer-family winner from the kernel-level chain + bench at grid points spanning the S=HV*N collapse (tensor_core iff T >= t_tc(S)).""" + cases = [ + # (HV, N, T, expect_tensor_core) -- kvbuffer-family winner per the kernel_level speedup table + (8, 1, 6, True), + (8, 4, 4, False), + (8, 4, 6, True), + (8, 8, 3, False), + (8, 8, 4, True), + (8, 32, 2, False), + (8, 32, 3, True), + (16, 2, 6, True), + (16, 4, 4, True), + (16, 16, 3, True), + (32, 1, 6, True), + (32, 2, 4, True), + (64, 1, 3, False), + (64, 1, 4, True), + (64, 4, 3, True), + (64, 128, 2, False), + ] + for hv, n, t, exp in cases: + assert _kvbuffer_prefer_tensor_core(n, hv, t) is exp, f"HV={hv} N={n} T={t}: want tensor_core={exp}" + + +@pytest.mark.parametrize("which,N,T,H,HV", [("shuffle", 4, 4, 16, 16), ("tensor_core", 4, 6, 16, 16)]) +def test_kvbuffer_verify_determinism(which, N, T, H, HV): + """Repeated kvbuffer verify launches produce a bit-identical output (and u-buffer).""" + V = K_DIM + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K_DIM, V) + scale = K_DIM**-0.5 + ub_ref = _alloc_ubufs(N, T, HV, V) + o_ref = _kvb_verify(which, q, k, v, a, b, A_log, dt_bias, state, scale, ubufs=ub_ref) + for i in range(int(os.environ.get("KDA_MTP_DET_ITERS", "100000"))): + ub_i = _alloc_ubufs(N, T, HV, V) + o_i = _kvb_verify(which, q, k, v, a, b, A_log, dt_bias, state, scale, ubufs=ub_i) + assert torch.equal(o_i, o_ref), f"{which} verify output non-deterministic at iter {i}" + assert torch.equal(ub_i[0], ub_ref[0]), f"{which} u-buffer non-deterministic at iter {i}" + + +@pytest.mark.parametrize("which,N,T,H,HV", [("shuffle", 4, 4, 16, 16), ("tensor_core", 4, 6, 16, 16)]) +def test_kvbuffer_flush_determinism(which, N, T, H, HV): + """Repeated flush launches rebuild a bit-identical state.""" + V = K_DIM + q, k, v, a, b, A_log, dt_bias, state = make_inputs_mtp(N, T, H, HV, K_DIM, V) + scale = K_DIM**-0.5 + indices = torch.arange(N, device=q.device, dtype=torch.int32) + ubufs = _alloc_ubufs(N, T, HV, V) + _kvb_verify(which, q, k, v, a, b, A_log, dt_bias, state, scale, ubufs=ubufs) + pool_ref = state.clone().contiguous() + kda_flush_kvbuffer(pool_ref, indices, ubufs[0], ubufs[1], ubufs[2], accept_len=T) + for i in range(int(os.environ.get("KDA_MTP_DET_ITERS", "100000"))): + pool_i = state.clone().contiguous() + kda_flush_kvbuffer(pool_i, indices, ubufs[0], ubufs[1], ubufs[2], accept_len=T) + assert torch.equal(pool_i, pool_ref), f"{which} flush state non-deterministic at iter {i}" + + +@pytest.mark.parametrize("V,N,HV", [(128, 1, 16), (128, 4, 32), (128, 16, 64)]) +def test_select_kvb_tile_v_invariants(V, N, HV): + """The auto tile_v must divide V and be a multiple of 4 (4-warp consumer).""" + tile_v = _select_kvb_tile_v(V, N, HV) + assert V % tile_v == 0 and tile_v % 4 == 0, f"tile_v={tile_v} violates V%tile_v==0 & tile_v%4==0" + + +@pytest.mark.parametrize("tile_v,T", [(64, 2), (32, 4), (64, 8), (16, 6)]) +def test_select_shuffle_kvb_ilp_rows_invariants(tile_v, T): + """ilp_rows must divide rows_per_group = tile_v/4 (the wrapper asserts this).""" + ilp = _select_shuffle_kvb_ilp_rows(tile_v, T) + assert ilp >= 1 and (tile_v // 4) % ilp == 0, f"ilp_rows={ilp} must divide tile_v/4={tile_v // 4}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--tb=short"]) From 0be4e507efe7f720e380e66e451e0a739cafb036 Mon Sep 17 00:00:00 2001 From: Fan Kun <784819644@qq.com> Date: Thu, 9 Jul 2026 11:29:03 +0800 Subject: [PATCH 30/34] [LA] Lightning Attention MTP decode + KVBuffer parallel verify / commit (#97) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: LA decode MTP kernel + tests + benchmark Fused multi-token (MTP) Lightning Attention decode kernel for speculative decoding: a single launch processes T draft tokens, with ILP variants and a work-unit heuristic (get_mtp_config). Includes packed F32x2 FMA on SM100. - cula/lightning/la_decode_mtp.py: kernel + config + shared dot/update helpers - tests/test_la_decode_mtp.py + tests/_la_mtp_ref.py: correctness vs PyTorch ref - benchmarks/bench_la_decode_mtp.py: vs sequential decode and FLA, with SOL model * feat: LA KVBuffer verify + state-update kernels + tests + benchmark KVBuffer-backed Lightning Attention for speculative decode verify/commit. Verify computes each draft step's output in closed form (paper Eq. 7) with the two dot-product GEMMs on tensor cores via inline-PTX mma.sync.m16n8k8 (TF32); state-update commits the accepted prefix into the pooled state (paper Eq. 8), bit-equivalent to the baseline T-loop at L == T. - cula/lightning/la_verify_kvbuffer.py: TF32 MMA verify kernel (+ shuffle variant) - cula/lightning/la_update_kvbuffer.py: KV buffer state-update (commit) kernel - tests/test_la_kvbuffer.py: correctness vs PyTorch ref (verify + update) - benchmarks/bench_la_kvbuffer.py: vs SGLang verify+commit (optional), with SOL model * test: cover odd T for KVBuffer verify/update; fix shuffle SMEM size - test_la_kvbuffer.py: add odd-T cases (verify T=1,3,5,7; state-update T=3,7) to guard the BT=8 M/N padding path that handles non-even draft lengths. - la_verify_kvbuffer.py: the shuffle launcher's SMEM byte estimate omitted the 16B per-allocation alignment padding (4 SMEM tensors), so the declared launch size could fall ~12B short of actual usage and trip CUTLASS's size check. Add the 4*16 padding term, matching the main-kernel launcher. * refactor: rename la_update_kvbuffer.py -> la_state_update_kvbuffer.py Module name now matches the public symbol linear_attention_state_update_kvbuffer. Pure rename plus import-path updates; no behavior change. * refactor: clean up la_decode_mtp kernel structure Structural cleanup of the LA decode-MTP kernel (no semantic change), split out of the prior pre-commit chore commit for reviewability. * chore: fix pre-commit issues and inline _la_mtp_ref into test files Formatting/lint fixes plus inlining the shared _la_mtp_ref helper directly into the test files. Benchmark updates included. * chore: add defensive bounds checks for T and V in LA kernels * chore: remove vestigial use_smem_v and use_packed_fma from MMA verify kernel * refactor: collapse ilp_rows branches in la_decode_mtp into generic constexpr loop Replace three explicit ilp_rows==2/==4/==8 branches with a single range_constexpr(ilp_rows) path, mirroring the pattern already used in la_state_update_kvbuffer. Cuts ~550 LOC without changing semantics. * harden LA verify/state-update kernels - la_verify_kvbuffer: re-check V % ilp_rows == 0 AFTER the ilp_rows->8 promotion (the pre-promotion assert could let a partial row-block be silently skipped); zero the sH0 M-padding rows before GEMM1 so the MMA fragment is well-defined instead of consuming stale/NaN SMEM. - assert K == 128 in the verify (MMA + shuffle) and state-update entry points, documenting the hardcoded head-dim assumption. * fix: pass correct cache key to MMA verify kernel in benchmark The kernel-only timing path passed shuffle-only args (use_smem_v, use_packed_fma) to _get_compiled_verify_kvbuffer_kernel, causing a TypeError at T >= MMA_MIN_T. * Reuse benchmark_cuda_fn from utils in LA benchmarks. Drop duplicated local benchmark_fn helpers; IQR-mean aggregation is already the default in benchmarks.utils.benchmark_cuda_fn. * Simplify LA MTP/KVBuffer benchmarks with layered compile-cache helpers. Align bench_la_decode_mtp and bench_la_kvbuffer with bench_la_decode_vs_fla: wrapper for correctness+compile warmup, then get_compiled_*_handle for kernel-only timing. Centralize cache-key dispatch in kernel modules. * perf(la): grid-search tile configs per kernel; drop use_smem_v The shared GDN-derived `get_mtp_config` was suboptimal for LA on both kernels. A grid search (B200, H=HV=64, K=V=128, B in [1..128], T in [2,4,8]) showed the two kernels sit in opposite regimes and need opposite tile directions: - decode (memory-bound, 83% mem SOL @ B=128,T=8): wants SMALL tiles + high occupancy. New thresholds tile_v in {32,16,8} as work_units grows (was 64 at large WU). 3-9% faster at medium-large B; SOL 74->81%. - verify/state-update (L1/compute-bound, 86% L1 SOL): wants LARGE tiles to fill m16n8k8 MMA and amortize q/k SMEM staging. New thresholds (64,8)/(128,8). 5-15% faster on the MMA path at large B. Split the single shared function into two `get_mtp_config` (one per kernel module) since one config cannot serve both regimes. state-update reuses the verify config (both prefer large tiles). Drop `use_smem_v` from the decode kernel: v has no cross-row reuse in LA, so SMEM staging only added a barrier. Grid search confirmed the direct-global path wins for every tile config — the True branch was dead code (never compiled). Removes sVdata/sOutput SMEM, the cooperative v-load, and the cooperative output writeback. 50/50 la_decode_mtp + la_kvbuffer tests pass. * fix(lightning): enforce fp32 for LA MTP and KVBuffer paths * chore: apply Ruff formatting --------- Co-authored-by: fankun.fan Co-authored-by: 范坤 --- benchmarks/bench_la_decode_mtp.py | 384 +++++++ benchmarks/bench_la_kvbuffer.py | 443 ++++++++ benchmarks/utils.py | 9 +- cula/lightning/__init__.py | 10 + cula/lightning/la_decode_mtp.py | 596 +++++++++++ cula/lightning/la_state_update_kvbuffer.py | 641 +++++++++++ cula/lightning/la_verify_kvbuffer.py | 1126 ++++++++++++++++++++ tests/test_la_decode_mtp.py | 455 ++++++++ tests/test_la_kvbuffer.py | 749 +++++++++++++ 9 files changed, 4412 insertions(+), 1 deletion(-) create mode 100644 benchmarks/bench_la_decode_mtp.py create mode 100644 benchmarks/bench_la_kvbuffer.py create mode 100644 cula/lightning/la_decode_mtp.py create mode 100644 cula/lightning/la_state_update_kvbuffer.py create mode 100644 cula/lightning/la_verify_kvbuffer.py create mode 100644 tests/test_la_decode_mtp.py create mode 100644 tests/test_la_kvbuffer.py diff --git a/benchmarks/bench_la_decode_mtp.py b/benchmarks/bench_la_decode_mtp.py new file mode 100644 index 00000000..9cfa5b7c --- /dev/null +++ b/benchmarks/bench_la_decode_mtp.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Benchmark: la_decode_mtp (CuTe DSL) vs alternatives on Lightning Attention MTP. + +Compares three implementations of T > 1 Lightning Attention decode: + 1. cula `linear_attention_decode_mtp` (this work — fused single-launch) + 2. fla `fused_recurrent_fwd` (Triton, T-aware) + 3. cula `linear_attention_decode` × T (cula self-comparison; T sequential calls) + +Two timing modes (mirroring bench_la_decode_vs_fla.py): + - kernel-only: pre-allocated buffers, pre-compiled kernel handle, pre-built stream + - wrapper: full Python entry point per call (cache lookup, CUstream, ...) + +Bandwidth analysis (SOL% against B200 HBM3e peak ~8 TB/s) printed alongside. + +Usage: + python benchmarks/bench_la_decode_mtp.py + python benchmarks/bench_la_decode_mtp.py --heads 64 --head-dim 128 --T 4 + python benchmarks/bench_la_decode_mtp.py --batch-sizes 1 4 16 64 --T 2 +""" + +import argparse +import os +import sys + +os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) + +import cuda.bindings.driver as cuda_drv +import torch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +try: + from fla.ops.common.fused_recurrent import fused_recurrent_fwd + + HAS_FLA = True +except ImportError: + HAS_FLA = False + +from benchmarks.utils import benchmark_cuda_fn, relative_rms_error +from cula.lightning.la_decode_mtp import get_compiled_la_mtp_handle, linear_attention_decode_mtp +from cula.ops.lightning.decode import linear_attention_decode +from cula.utils import USE_FAST_MATH + + +# ───────────────────────────────────────────────────────────────────────────── +# Bandwidth model — see spec §9.3 +# ───────────────────────────────────────────────────────────────────────────── +def la_mtp_bytes(B, T, H, HV, K, V, cache_intermediate_states, disable_state_update): + fp32 = 4 + qkv = B * T * H * K * fp32 * 2 + B * T * HV * V * fp32 # q, k, v reads + out_w = B * T * HV * V * fp32 # o writes + h0_r = B * HV * V * K * fp32 # h0 reads + h0_w = 0 if disable_state_update else B * HV * V * K * fp32 # h0 writes + inter = B * T * HV * V * K * fp32 if cache_intermediate_states else 0 + return qkv + out_w + h0_r + h0_w + inter + + +def sol_pct(byte_count: int, kernel_ms: float, peak_bps: float) -> float: + """Speed-of-light percent of HBM peak.""" + return (byte_count / (kernel_ms * 1e-3)) / peak_bps * 100.0 + + +# ───────────────────────────────────────────────────────────────────────────── +# Core benchmark for one (B, T) configuration +# ───────────────────────────────────────────────────────────────────────────── +def run_config( + B, T, H, HV, K, V, layer_idx, num_layers, peak_bps, cache_intermediate_states=False, disable_state_update=False +): + device = "cuda" + dtype = torch.float32 + scale = K**-0.5 + pool_size = B + + # Per-head log decay (Lightning Attention formula) + g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * torch.arange(H, device=device, dtype=torch.float32) + decay_scales = -g_gamma # la_decode_mtp convention: exp(-decay_scales) + + # ========================================================================= + # Layer 1 — Inputs & buffers + # ========================================================================= + torch.manual_seed(42) + q_4d = torch.randn(B, T, H, K, device=device, dtype=dtype) + k_4d = torch.randn(B, T, H, K, device=device, dtype=dtype) + v_4d = torch.randn(B, T, HV, V, device=device, dtype=dtype) + state_init = torch.randn(B, HV, K, V, device=device, dtype=torch.float32) * 0.01 # K-major + + s_offsets = torch.arange(B, device=device, dtype=torch.int32) + cu_seqlens_dummy = torch.empty(1, device=device, dtype=torch.int32) + inter = ( + torch.zeros(B * T * HV, V, K, device=device, dtype=torch.float32) + if cache_intermediate_states + else torch.empty(1, 1, 1, device=device, dtype=torch.float32) + ) + + # ========================================================================= + # Layer 2 — Correctness + compile warmup (wrapper, same config as benchmark) + # ========================================================================= + # fla reference + o_fla = None + if HAS_FLA: + state_fla = state_init.clone() + with torch.no_grad(): + o_fla_fp32, _ht_fla = fused_recurrent_fwd( + q_4d, + k_4d, + v_4d, + g_gamma=g_gamma, + scale=scale, + initial_state=state_fla, + output_final_state=True, + ) + o_fla = o_fla_fp32.to(dtype) + + # cuLA MTP — also populates the compile cache for kernel-only timing below + s_cute = state_init.clone().permute(0, 1, 3, 2).contiguous() # [B, HV, V, K] + out_cute = torch.zeros(B, T, HV, V, device=device, dtype=dtype) + with torch.no_grad(): + linear_attention_decode_mtp( + q_4d, + k_4d, + v_4d, + s_cute, + inter, + out_cute, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens_dummy, + softmax_scale=scale, + T=T, + cache_intermediate_states=cache_intermediate_states, + disable_state_update=disable_state_update, + is_varlen=False, + ) + + rmse = rel_maxdiff = float("nan") + if o_fla is not None and HV == H: + rmse = relative_rms_error(o_fla.float(), out_cute.float()) + ref_cmp = o_fla.float() + out_cmp = out_cute.float() + max_ref = torch.abs(ref_cmp).max().item() + rel_maxdiff = torch.abs(out_cmp - ref_cmp).max().item() / (max_ref + 1e-8) + + # ========================================================================= + # Layer 3a — Kernel-only timing (compiled handle + pre-built stream) + # ========================================================================= + compiled_cute = get_compiled_la_mtp_handle( + B, + T, + H, + HV, + K, + V, + pool_size, + scale, + q_4d.device, + disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, + is_varlen=False, + ) + stream_handle = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + + state_kk = state_init.clone().permute(0, 1, 3, 2).contiguous().view(pool_size * HV, V, K) + out_kk = torch.empty(B, T, HV, V, device=device, dtype=dtype) + inter_kk = inter + + def kernel_cute_mtp(): + compiled_cute( + state_kk, + inter_kk, + decay_scales, + q_4d, + k_4d, + v_4d, + out_kk, + s_offsets, + cu_seqlens_dummy, + stream_handle, + ) + + # cula self-baseline: T sequential la_decode (T=1) wrapper calls + state_seq = state_init.clone().permute(0, 1, 3, 2).contiguous().view(B * HV, V, K) + out_seq_buf = torch.empty(B, HV, V, device=device, dtype=dtype) + q_slices = [q_4d[:, t].contiguous() for t in range(T)] + k_slices = [k_4d[:, t].contiguous() for t in range(T)] + v_slices = [v_4d[:, t].contiguous() for t in range(T)] + + def kernel_cute_seq(): + for t in range(T): + linear_attention_decode( + q_slices[t], + k_slices[t], + v_slices[t], + state_seq, + out_seq_buf, + softmax_scale=scale, + stride_q=0, + stride_k=0, + stride_v=0, + stride_s=0, + stride_o=0, + s_offsets=s_offsets, + decay_scales=decay_scales, + HEAD_DIM=K, + K_SPLIT_DIM=K, + V_SPLIT_DIM=V, + ) + + with torch.no_grad(): + cute_mtp_ms = benchmark_cuda_fn(kernel_cute_mtp) + cute_seq_ms = benchmark_cuda_fn(kernel_cute_seq) + + # ========================================================================= + # Layer 3b — Wrapper timing (full Python entry path per call) + # ========================================================================= + s_wrap = state_init.clone().permute(0, 1, 3, 2).contiguous() + out_wrap = torch.empty(B, T, HV, V, device=device, dtype=dtype) + inter_wrap = ( + torch.zeros(B * T * HV, V, K, device=device, dtype=torch.float32) + if cache_intermediate_states + else torch.empty(1, 1, 1, device=device, dtype=torch.float32) + ) + + def wrapper_cute_mtp(): + linear_attention_decode_mtp( + q_4d, + k_4d, + v_4d, + s_wrap, + inter_wrap, + out_wrap, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens_dummy, + softmax_scale=scale, + T=T, + cache_intermediate_states=cache_intermediate_states, + disable_state_update=disable_state_update, + is_varlen=False, + ) + + with torch.no_grad(): + wrap_cute_ms = benchmark_cuda_fn(wrapper_cute_mtp) + + fla_ms = float("nan") + if HAS_FLA: + state_fla_bench = state_init.clone() + + def wrapper_fla(): + fused_recurrent_fwd( + q_4d, + k_4d, + v_4d, + g_gamma=g_gamma, + scale=scale, + initial_state=state_fla_bench, + output_final_state=True, + ) + + with torch.no_grad(): + fla_ms = benchmark_cuda_fn(wrapper_fla) + + # ========================================================================= + # Layer 4 — Roofline & summary + # ========================================================================= + bytes_moved = la_mtp_bytes( + B, + T, + H, + HV, + K, + V, + cache_intermediate_states=cache_intermediate_states, + disable_state_update=disable_state_update, + ) + sol = sol_pct(bytes_moved, cute_mtp_ms, peak_bps) + + return { + "B": B, + "T": T, + "cute_mtp_ms": cute_mtp_ms, + "cute_seq_ms": cute_seq_ms, + "fla_ms": fla_ms, + "wrap_cute_ms": wrap_cute_ms, + "speedup_seq": cute_seq_ms / cute_mtp_ms, + "speedup_fla": fla_ms / cute_mtp_ms if HAS_FLA else float("nan"), + "rmse": rmse, + "rel_maxdiff": rel_maxdiff, + "sol_pct": sol, + "bytes_GB": bytes_moved / 1e9, + } + + +# ───────────────────────────────────────────────────────────────────────────── +# Main +# ───────────────────────────────────────────────────────────────────────────── +def main(): + parser = argparse.ArgumentParser(description="Benchmark la_decode_mtp") + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64, 128]) + parser.add_argument("--T", type=int, nargs="+", default=[2, 4, 8]) + parser.add_argument("--heads", type=int, default=32) + parser.add_argument("--num-v-heads", type=int, default=None, help="HV (defaults to --heads for MHA)") + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--layer-idx", type=int, default=12) + parser.add_argument("--num-layers", type=int, default=24) + parser.add_argument("--peak-bps", type=float, default=8e12, help="HBM peak bytes/sec for SOL%% (B200 HBM3e ≈ 8e12)") + parser.add_argument("--cache-intermediate", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--disable-state-update", action=argparse.BooleanOptionalAction, default=True) + args = parser.parse_args() + + H = args.heads + HV = args.num_v_heads if args.num_v_heads is not None else H + K = V = args.head_dim + + print("Lightning Attention MTP Decode Benchmark") + print(f" H={H}, HV={HV}, K={K}, V={V}, layer={args.layer_idx}/{args.num_layers}") + print(f" dtype=fp32, state=fp32, peak={args.peak_bps:.2e} B/s") + print(f" cache_intermediate_states={args.cache_intermediate}, disable_state_update={args.disable_state_update}") + print(f" USE_FAST_MATH={USE_FAST_MATH}, fla available={HAS_FLA}") + + fla_avail = HAS_FLA and HV == H + if HAS_FLA and HV != H: + print(f" [warning] GQA HV={HV} != H={H}; fla baseline disabled (fla assumes HV==H)") + + cols = ( + f"{'B':>4} | {'T':>3} | {'cute_mtp(ms)':>12} | {'cute×T(ms)':>10} | " + f"{'fla(ms)':>9} | {'spd_seq':>7} | {'spd_fla':>7} | " + f"{'wrap(ms)':>9} | {'SOL%':>5} | {'GB':>6} | {'RMSE':>9}" + ) + print(f"\n{cols}") + print("─" * len(cols)) + + for T in args.T: + for B in args.batch_sizes: + r = run_config( + B, + T, + H, + HV, + K, + V, + args.layer_idx, + args.num_layers, + args.peak_bps, + cache_intermediate_states=args.cache_intermediate, + disable_state_update=args.disable_state_update, + ) + print( + f"{r['B']:>4} | {r['T']:>3} | {r['cute_mtp_ms']:>12.4f} | " + f"{r['cute_seq_ms']:>10.4f} | " + f"{(r['fla_ms'] if fla_avail else float('nan')):>9.4f} | " + f"{r['speedup_seq']:>6.2f}x | " + f"{(r['speedup_fla'] if fla_avail else float('nan')):>6.2f}x | " + f"{r['wrap_cute_ms']:>9.4f} | {r['sol_pct']:>5.1f} | " + f"{r['bytes_GB']:>6.3f} | {r['rmse']:>9.6f}" + ) + print() + + print("Notes:") + print(" cute_mtp : linear_attention_decode_mtp (fused single launch, T tokens)") + print(" cute×T : T sequential linear_attention_decode (T=1) calls — cula self-baseline") + print(" fla : fused_recurrent_fwd (Triton); kernel still re-launched per T internally") + print(" spd_seq : cute×T / cute_mtp (fusion benefit within cula)") + print(" spd_fla : fla / cute_mtp (vs industry reference)") + print(" wrap(ms) : cute_mtp full Python entry (cache lookup + CUstream + kernel)") + print(f" SOL% : (bytes / kernel_ms) / peak_bps × 100 (peak = {args.peak_bps:.2e} B/s)") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_la_kvbuffer.py b/benchmarks/bench_la_kvbuffer.py new file mode 100644 index 00000000..f137084f --- /dev/null +++ b/benchmarks/bench_la_kvbuffer.py @@ -0,0 +1,443 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +""" +Benchmark: cuLA LA KVBuffer verify + state-update kernels. + +Times the KVBuffer path (verify writes k/v to a pool buffer; state-update advances +the pooled state from it) and validates against a PyTorch reference. + +An optional SGLang baseline (seg_la_mtp_kernel + fused_mamba_state_scatter_with_mask) +is compared when available. Set LA_SGLANG_PYTHON=/path/to/sglang/python for a custom +SGLang checkout. + +Timing follows bench_la_decode_vs_fla.py: + - Layer 2: wrapper call for correctness + compile warmup (same config as benchmark) + - Layer 3: kernel-only via pre-compiled handles + pre-built stream + +Usage: + python benchmarks/bench_la_kvbuffer.py + python benchmarks/bench_la_kvbuffer.py --batch-sizes 1 4 16 64 --T 2 4 8 + LA_SGLANG_PYTHON=~/sglang/python python benchmarks/bench_la_kvbuffer.py --T 4 +""" + +import argparse +import os +import sys + +import torch + +os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) + +import cuda.bindings.driver as cuda_drv + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +# ── Optional SGLang baseline ───────────────────────────────────────────────── +_HAVE_SGLANG, _SGLANG_ERR = True, "" +SegLaMeta = seg_la_mtp_kernel = seg_la_sum_kernel = None +fused_mamba_state_scatter_with_mask = None +try: + _sg_path = os.environ.get("LA_SGLANG_PYTHON", "") + if _sg_path and os.path.isdir(_sg_path): + sys.path.insert(0, _sg_path) + from sglang.srt.layers.attention.linear.seg_la import ( + SegLaMeta, + seg_la_mtp_kernel, + seg_la_sum_kernel, + ) + from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( + fused_mamba_state_scatter_with_mask, + ) +except Exception as e: # noqa: BLE001 + _HAVE_SGLANG, _SGLANG_ERR = False, repr(e) + +from benchmarks.utils import benchmark_cuda_fn, relative_rms_error # noqa: E402 +from cula.lightning.la_state_update_kvbuffer import ( # noqa: E402 + get_compiled_state_update_kvbuffer_handle, + linear_attention_state_update_kvbuffer, +) +from cula.lightning.la_verify_kvbuffer import ( # noqa: E402 + get_compiled_verify_kvbuffer_handle, + linear_attention_verify_kvbuffer, +) +from cula.utils import USE_FAST_MATH # noqa: E402 + + +# ───────────────────────────────────────────────────────────────────────────── +# Reference & SGLang helpers +# ───────────────────────────────────────────────────────────────────────────── +def torch_la_mtp_ref(q, k, v, state, decay_scales, softmax_scale): + """Pure PyTorch reference for MTP decode (output only).""" + B, T, H, K = q.shape + V = v.shape[-1] + state = state.clone().float() + out = torch.zeros(B, T, H, V, device=q.device, dtype=torch.float32) + decay = torch.exp(-decay_scales).float() + + for t in range(T): + qt = q[:, t].float() * softmax_scale + kt = k[:, t].float() + vt = v[:, t].float() + state = state * decay[None, :, None, None] + kt.unsqueeze(-1) * vt.unsqueeze(-2) + out[:, t] = torch.einsum("bhk,bhkv->bhv", qt, state) + + return out + + +def run_sglang_mtp( + q_3d, + k_3d, + v_3d, + s_sglang, + caches_sglang, + s_offsets, + cache_indices, + decay_scales, + meta, + softmax_scale, + HEAD_DIM, + step, + K_SPLIT_DIM=32, + V_SPLIT_DIM=64, +): + """Invoke seg_la_mtp_kernel the same way seg_la_fwd does for the MTP path.""" + length = q_3d.shape[0] + qo_heads = q_3d.shape[1] + bs = meta.batch_size + + k_dim_block = HEAD_DIM // K_SPLIT_DIM + v_dim_block = HEAD_DIM // V_SPLIT_DIM + tmp = torch.empty((k_dim_block, length, qo_heads, HEAD_DIM), device=q_3d.device, dtype=q_3d.dtype) + grid = (bs, qo_heads, k_dim_block * v_dim_block) + + seg_la_mtp_kernel[grid]( + q_3d, + k_3d, + v_3d, + s_sglang, + caches_sglang, + tmp, + softmax_scale, + q_3d.stride(0), + k_3d.stride(0), + v_3d.stride(0), + s_sglang.stride(0), + caches_sglang.stride(0), + tmp.stride(0), + s_offsets, + cache_indices, + decay_scales, + step, + HEAD_DIM=HEAD_DIM, + K_SPLIT_DIM=K_SPLIT_DIM, + V_SPLIT_DIM=V_SPLIT_DIM, + num_warps=2, + num_stages=3, + ) + + if k_dim_block > 1: + if length < 2048: + o = tmp.sum(0) + else: + o = torch.empty((length, qo_heads, HEAD_DIM), device=q_3d.device, dtype=q_3d.dtype) + seg_la_sum_kernel[(length,)]( + tmp, + o, + DIM=qo_heads * HEAD_DIM, + NUM_BLOCK=k_dim_block, + num_warps=2, + num_stages=3, + ) + else: + o = tmp[0] + return o + + +def run_sglang_commit(s_sglang, caches_sglang, s_offsets, step_indices, B, H, K, V, T): + """Invoke fused_mamba_state_scatter_with_mask (SGLang commit step).""" + elem_per_entry = H * K * V + dst = s_sglang.reshape(1, -1, elem_per_entry) + src = caches_sglang.reshape(1, B, T, elem_per_entry) + fused_mamba_state_scatter_with_mask(dst, src, s_offsets, step_indices) + + +# ───────────────────────────────────────────────────────────────────────────── +# Core benchmark for one (B, T) configuration +# ───────────────────────────────────────────────────────────────────────────── +def run_config(B, T, H, K, V, layer_idx, num_layers): + device = "cuda" + input_dtype = torch.float32 + out_dtype = torch.float32 + scale = K**-0.5 + HV = H # SGLang seg_la does not support GQA + pool_size = B + + g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * torch.arange(H, device=device, dtype=torch.float32) + decay_scales = -g_gamma + + # ========================================================================= + # Layer 1 — Inputs & buffers (benchmark config: verify writes k/v, commit reads k/v buffers) + # ========================================================================= + torch.manual_seed(42) + q_4d = torch.randn(B, T, H, K, device=device, dtype=input_dtype) + k_4d = torch.randn(B, T, H, K, device=device, dtype=input_dtype) + v_4d = torch.randn(B, T, HV, V, device=device, dtype=input_dtype) + state_init = torch.randn(B, H, K, V, device=device, dtype=torch.float32) * 0.01 # K-major + + # cuLA state pool [pool_size, HV, V, K] + s_kvbuf = state_init.permute(0, 1, 3, 2).contiguous() + s_kk_view = s_kvbuf.view(pool_size * HV, V, K) + + out_kvbuf = torch.zeros(B, T, HV, V, device=device, dtype=out_dtype) + out_kk = torch.empty(B, T, HV, V, device=device, dtype=out_dtype) + + h0_indices = torch.arange(B, device=device, dtype=torch.int32) + accepted_len = torch.full((B,), T, device=device, dtype=torch.int32) + + k_buf = torch.zeros(pool_size, T, H, K, device=device, dtype=torch.float32) + v_buf = torch.zeros(pool_size, T, HV, V, device=device, dtype=torch.float32) + + # SGLang 3D views (length = B*T) + q_3d = q_4d.reshape(B * T, H, K).contiguous() + k_3d = k_4d.reshape(B * T, H, K).contiguous() + v_3d = v_4d.reshape(B * T, HV, V).contiguous() + + # ========================================================================= + # Layer 2 — Correctness + compile warmup (wrapper, same config as benchmark) + # ========================================================================= + with torch.no_grad(): + o_ref = torch_la_mtp_ref(q_4d, k_4d, v_4d, state_init, decay_scales, scale) + + linear_attention_verify_kvbuffer( + q_4d, + k_4d, + v_4d, + s_kvbuf, + out_kvbuf, + decay_scales, + h0_indices, + scale, + T, + k_buf=k_buf, + v_buf=v_buf, + ) + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_kvbuf, + decay_scales, + h0_indices, + accepted_len, + T, + ) + + rmse_kv = relative_rms_error(o_ref, out_kvbuf.float()) + + # SGLang baseline (optional): correctness call also JIT-compiles Triton kernels + rmse_sg = float("nan") + s_sglang = caches_sglang = s_offsets_sg = cache_indices_sg = meta = None + K_SPLIT_DIM = 32 + V_SPLIT_DIM = 32 if B <= 2 else 64 + if _HAVE_SGLANG: + s_sglang = state_init.reshape(pool_size, H, K, V).contiguous() + caches_sglang = torch.zeros(pool_size * T, H, K, V, device=device, dtype=torch.float32) + s_offsets_sg = torch.arange(B, device=device, dtype=torch.int64) + cache_indices_sg = torch.arange(B, device=device, dtype=torch.int64) * T + meta = SegLaMeta( + batch_size=B, + max_q_length=T, + q_offsets=torch.arange(B + 1, device=device, dtype=torch.int64) * T, + s_offsets=s_offsets_sg, + q_lengths=torch.full((B,), T, device=device, dtype=torch.int64), + s_scales=torch.ones(B, device=device, dtype=torch.int64), + ) + with torch.no_grad(): + o_sg = run_sglang_mtp( + q_3d, + k_3d, + v_3d, + s_sglang.clone(), + caches_sglang.clone(), + s_offsets_sg, + cache_indices_sg, + decay_scales, + meta, + scale, + K, + T, + K_SPLIT_DIM, + V_SPLIT_DIM, + ) + rmse_sg = relative_rms_error(o_ref, o_sg.reshape(B, T, HV, V).float()) + + # ========================================================================= + # Layer 3 — Kernel-only timing (compiled handles + pre-built stream) + # ========================================================================= + stream_handle = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + + compiled_verify = get_compiled_verify_kvbuffer_handle( + B, T, H, HV, K, V, pool_size, scale, write_kv=True, device=q_4d.device + ) + compiled_update = get_compiled_state_update_kvbuffer_handle(B, T, H, HV, K, V, pool_size, device=q_4d.device) + + def kernel_kvbuf_verify(): + compiled_verify( + s_kk_view, + decay_scales, + q_4d, + k_4d, + v_4d, + out_kk, + h0_indices, + k_buf, + v_buf, + stream_handle, + ) + + def kernel_kvbuf_update(): + compiled_update( + s_kk_view, + decay_scales, + h0_indices, + accepted_len, + k_buf, + v_buf, + stream_handle, + ) + + step_indices_sg = torch.full((B,), T - 1, device=device, dtype=torch.int32) + + def kernel_sglang_verify(): + run_sglang_mtp( + q_3d, + k_3d, + v_3d, + s_sglang, + caches_sglang, + s_offsets_sg, + cache_indices_sg, + decay_scales, + meta, + scale, + K, + T, + K_SPLIT_DIM, + V_SPLIT_DIM, + ) + + def kernel_sglang_commit(): + run_sglang_commit( + s_sglang, + caches_sglang, + s_offsets_sg.int(), + step_indices_sg, + B, + H, + K, + V, + T, + ) + + with torch.no_grad(): + cu_vfy_ms = benchmark_cuda_fn(kernel_kvbuf_verify) + cu_cmt_ms = benchmark_cuda_fn(kernel_kvbuf_update) + if _HAVE_SGLANG: + sg_vfy_ms = benchmark_cuda_fn(kernel_sglang_verify) + sg_cmt_ms = benchmark_cuda_fn(kernel_sglang_commit) + else: + sg_vfy_ms = sg_cmt_ms = float("nan") + + # ========================================================================= + # Layer 4 — Summary + # ========================================================================= + sg_total_ms = sg_vfy_ms + sg_cmt_ms + cu_total_ms = cu_vfy_ms + cu_cmt_ms + + return { + "B": B, + "T": T, + "sg_vfy_ms": sg_vfy_ms, + "sg_cmt_ms": sg_cmt_ms, + "sg_total_ms": sg_total_ms, + "cu_vfy_ms": cu_vfy_ms, + "cu_cmt_ms": cu_cmt_ms, + "cu_total_ms": cu_total_ms, + "speedup": (sg_total_ms / cu_total_ms) if _HAVE_SGLANG else float("nan"), + "rmse_sg": rmse_sg, + "rmse_kv": rmse_kv, + } + + +# ───────────────────────────────────────────────────────────────────────────── +# Main +# ───────────────────────────────────────────────────────────────────────────── +def main(): + parser = argparse.ArgumentParser(description="Benchmark LA KVBuffer verify + state-update") + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64, 128]) + parser.add_argument("--T", type=int, nargs="+", default=[2, 4, 8]) + parser.add_argument("--heads", type=int, default=32) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--layer-idx", type=int, default=12) + parser.add_argument("--num-layers", type=int, default=24) + args = parser.parse_args() + + H = args.heads + K = V = args.head_dim + + print("LA KVBuffer verify + state-update benchmark (cuLA, optional SGLang baseline)") + print(f" H={H}, K={K}, V={V}, layer={args.layer_idx}/{args.num_layers}") + print(" q/k/v=fp32, out=fp32, state=fp32, kv_buffer=fp32") + print(f" USE_FAST_MATH={USE_FAST_MATH}") + print(" Timing: kernel-only (wrapper for compile warmup; compiled handle for measure)") + if _HAVE_SGLANG: + print(" SGLang baseline: AVAILABLE (sg_* columns active)") + else: + print(f" SGLang baseline: UNAVAILABLE — sg_* columns show nan. ({_SGLANG_ERR})") + print(" set LA_SGLANG_PYTHON=/path/to/sglang/python to enable the comparison.") + + hdr = ( + f"{'B':>4} | {'T':>3} | " + f"{'sg_vfy(ms)':>10} | {'sg_cmt(ms)':>10} | {'sg_total':>9} | " + f"{'cu_vfy(ms)':>10} | {'cu_cmt(ms)':>10} | {'cu_total':>9} | " + f"{'speedup':>7} | {'rmse_sg':>9} | {'rmse_kv':>9}" + ) + print(f"\n{hdr}") + print("─" * len(hdr)) + + for T_val in args.T: + for B in args.batch_sizes: + r = run_config(B, T_val, H, K, V, args.layer_idx, args.num_layers) + print( + f"{r['B']:>4} | {r['T']:>3} | " + f"{r['sg_vfy_ms']:>10.4f} | {r['sg_cmt_ms']:>10.4f} | {r['sg_total_ms']:>9.4f} | " + f"{r['cu_vfy_ms']:>10.4f} | {r['cu_cmt_ms']:>10.4f} | {r['cu_total_ms']:>9.4f} | " + f"{r['speedup']:>6.2f}x | " + f"{r['rmse_sg']:>9.6f} | {r['rmse_kv']:>9.6f}" + ) + print() + + sg_mem = B * T_val * H * K * V * 4 + cu_mem = B * T_val * (H * K + H * V) * 4 + print(f"Memory per-pool (B={args.batch_sizes[-1]}, T={args.T[-1]}):") + print(f" SGLang intermediate caches: {sg_mem / 1e6:.1f} MB") + print(f" cuLA KV buffer: {cu_mem / 1e6:.1f} MB") + print(f" Ratio: {sg_mem / cu_mem:.0f}×") + + print("\nColumns:") + print(" sg_vfy : seg_la_mtp_kernel (Triton, SGLang upstream)") + print(" sg_cmt : fused_mamba_state_scatter_with_mask (Triton, SGLang)") + print(" cu_vfy : verify_kvbuffer with KV buffer write (CuTe DSL)") + print(" cu_cmt : state_update_kvbuffer reading from buffer (CuTe DSL)") + print(" speedup : sg_total / cu_total") + print(" rmse_* : relative RMS error vs PyTorch reference") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 55602490..75786e1f 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -61,7 +61,14 @@ def set_seed(seed: int): def benchmark_cuda_fn(fn, *, setup_fn=None, warmup=30, rep=200, aggregate="iqr_mean"): - """Benchmark a CUDA callable with events and return milliseconds per call.""" + """Benchmark a CUDA callable with CUDA events; return milliseconds per call. + + Args: + aggregate: How to summarize ``rep`` timed iterations. + ``"iqr_mean"`` (default) — mean of the middle 50% after sorting + (robust to outliers; used by la_decode / MTP benchmarks). + ``"mean"`` — arithmetic mean of all iterations. + """ for _ in range(warmup): if setup_fn is not None: setup_fn() diff --git a/cula/lightning/__init__.py b/cula/lightning/__init__.py index af2e3dcf..caef73e4 100644 --- a/cula/lightning/__init__.py +++ b/cula/lightning/__init__.py @@ -12,6 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from cula.lightning.la_decode_mtp import linear_attention_decode_mtp +from cula.lightning.la_state_update_kvbuffer import ( + linear_attention_state_update_kvbuffer, + linear_attention_state_update_kvbuffer_fused, +) +from cula.lightning.la_verify_kvbuffer import linear_attention_verify_kvbuffer from cula.ops.lightning.decode import linear_attention_decode from cula.ops.lightning.prefill_sm100 import ( LinearAttentionChunkwiseDecay, @@ -24,4 +30,8 @@ "lightning_attn_fwd", "lightning_attn_fwd_varlen", "linear_attention_decode", + "linear_attention_decode_mtp", + "linear_attention_verify_kvbuffer", + "linear_attention_state_update_kvbuffer", + "linear_attention_state_update_kvbuffer_fused", ] diff --git a/cula/lightning/la_decode_mtp.py b/cula/lightning/la_decode_mtp.py new file mode 100644 index 00000000..17d38da6 --- /dev/null +++ b/cula/lightning/la_decode_mtp.py @@ -0,0 +1,596 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Lightning Attention MTP (Multi-Token Processing) Decode Kernel. + +Processes T > 1 tokens in one launch with h held in registers across the +whole T-loop. Targeted at speculative-decoding verify scenarios. + +Per timestep: + h_t = exp(-decay_scales[h]) * h_{t-1} + k_t ⊗ v_t + o_t = (h_t @ q_t) * softmax_scale + +`decay_scales` is per-head and time-invariant, so `r_decay` is computed ONCE +outside the T-loop. + +Grid: (B * HV * num_v_tiles, 1, 1). Each block handles one [tile_v] slice +across all T timesteps; h for that slice stays in registers. + +Reference: flashinfer/flashinfer/gdn_kernels/gdn_decode_mtp.py (inline variant). +""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from cula.utils import USE_FAST_MATH, get_device_sm_version + +# ============================================================================ +# Global configuration +# ============================================================================ +TILE_K_MTP = 128 +NUM_THREADS_MTP = 128 # 4 warps + + +# ============================================================================ +# FMA pair helpers (packed F32x2 on SM100; scalar fallback on SM90) +# ============================================================================ +@cute.jit +def la_update_pair(h_lo, h_hi, k_lo, k_hi, v_j, decay, use_packed_fma: cutlass.Constexpr[bool]): + """Inner LA recurrence on a (lo, hi) pair: h = h*decay + k*v_j.""" + if cutlass.const_expr(use_packed_fma): + # h *= decay (packed mul implemented as FMA with src_c=0) + h_lo, h_hi = cute.arch.fma_packed_f32x2( + src_a=(h_lo, h_hi), + src_b=(decay, decay), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + # h += k * v_j + h_lo, h_hi = cute.arch.fma_packed_f32x2( + src_a=(k_lo, k_hi), + src_b=(v_j, v_j), + src_c=(h_lo, h_hi), + ) + return h_lo, h_hi + else: + return h_lo * decay + k_lo * v_j, h_hi * decay + k_hi * v_j + + +@cute.jit +def hq_dot_pair(h_lo, h_hi, q_lo, q_hi, sum_lo, sum_hi, use_packed_fma: cutlass.Constexpr[bool]): + """Accumulate dot product over a (lo, hi) pair: sum += h * q.""" + if cutlass.const_expr(use_packed_fma): + return cute.arch.fma_packed_f32x2( + src_a=(h_lo, h_hi), + src_b=(q_lo, q_hi), + src_c=(sum_lo, sum_hi), + ) + else: + return h_lo * q_lo + sum_lo, h_hi * q_hi + sum_hi + + +# TODO (perf): for configs with row_iters > 1 (e.g. tile_v=64, ilp=4), q/k are +# reloaded from global on every row-loop iteration because the row-outer / T-inner +# structure is required to keep h register-resident across T (r_h budget is 8 rows). +# Stage q/k in SMEM per i_t (cooperative load + barrier) to avoid the (row_iters - 1) +# redundant reads; worst case (tile_v=64, ilp=4) wastes 3x the q/k bandwidth. +# With the LA-tuned thresholds (tile_v <= 32), row_iters <= 2, so this is less +# urgent, but still worth doing for ilp=2 with larger tile_v. +def get_mtp_config(B: int, T: int, HV: int, V: int, disable_state_update: bool) -> tuple: + """Pick (tile_v, vec_size, ilp_rows) for the decode kernel based on work units. + + LA grid search on B200 (H=HV=64, K=V=128) with B ∈ [1..128], T ∈ [2,4,8]. + LA's per-step compute is ~30% lighter than GDN (no delta rule), so the + compute/memory ratio is lower — favouring smaller tiles with more blocks + to improve occupancy and amortize per-block overhead. + + The old GDN-derived thresholds (tile_v=64, ilp=4 for work_units > 1024) are + suboptimal for LA by 3-12% at medium-to-large B. ``use_smem_v`` was dropped: + v has no cross-row reuse in LA, so SMEM staging only added a barrier (grid + search confirmed the direct-global path wins for every tile config). + + ``disable_state_update`` is kept in the signature for API stability but no + longer affects the tile choice (the old state-update branch collapsed). + """ + work_units = B * HV + vec_size = 4 + + if work_units <= 256: + tile_v, ilp_rows = 32, 8 + elif work_units <= 1024: + tile_v, ilp_rows = 16, 4 + else: + tile_v, ilp_rows = 8, 2 + + tile_v = min(tile_v, V) + rows_per_group = tile_v // 4 + assert rows_per_group % ilp_rows == 0, ( + f"tile_v={tile_v} / num_groups=4 / ilp_rows={ilp_rows} doesn't divide cleanly " + f"(rows_per_group={rows_per_group}); the ILP loop would run zero iterations." + ) + return tile_v, vec_size, ilp_rows + + +# ============================================================================ +# Kernel +# ============================================================================ +@cute.kernel +def la_verify_kernel_mtp( + h0_source: cute.Tensor, # [pool_size * HV, V, K] fp32 + intermediate_states: cute.Tensor, # [pool_size * T * HV, V, K] fp32 (or dummy) + decay_scales: cute.Tensor, # [H] fp32 + q: cute.Tensor, # [B, T, H, K] fp32 + k: cute.Tensor, # [B, T, H, K] fp32 + v: cute.Tensor, # [B, T, HV, V] fp32 + o: cute.Tensor, # [B, T, HV, V] fp32 + h0_indices: cute.Tensor, # [B] int32 + cu_seqlens: cute.Tensor, # [B+1] int32 (dummy when is_varlen=False) + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + scale: cutlass.Constexpr[float], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + disable_state_update: cutlass.Constexpr[bool], + cache_intermediate_states: cutlass.Constexpr[bool], + is_varlen: cutlass.Constexpr[bool], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + threads_per_group: cutlass.Constexpr[int] = K // vec_size # 32 + groups_per_warp: cutlass.Constexpr[int] = 32 // threads_per_group # 1 + num_groups: cutlass.Constexpr[int] = 4 * groups_per_warp # 4 + + lane_in_group = lane_id % threads_per_group + group_in_warp = lane_id // threads_per_group + group_idx = warp_idx * groups_per_warp + group_in_warp + + block_idx, _, _ = cute.arch.block_idx() + i_v = block_idx % num_v_tiles + tmp = block_idx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + + # ------------------------------------------------------------------ + # Register tensors (LA decode is memory-bound — no SMEM staging; v has no + # cross-row reuse so staging it would only add a barrier. Grid search + # confirmed the direct-global path wins for every tile config.) + # ------------------------------------------------------------------ + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + # r_h always declared with 8 rows; ilp_rows constexpr picks which are used. + r_h = cute.make_rmem_tensor(cute.make_layout((8, vec_size), stride=(vec_size, 1)), cutlass.Float32) + + if cache_idx >= 0: + # r_decay is a T-loop invariant — computed ONCE. + r_decay = cute.exp(-cutlass.Float32(decay_scales[i_h]), fastmath=USE_FAST_MATH) + + rows_per_group: cutlass.Constexpr[int] = tile_v // num_groups + flat_state_idx = cache_idx * HV + i_hv + + # Process `ilp_rows` V-rows per iteration. ilp_rows is a compile-time + # constant, so range_constexpr fully unrolls the slot loops below — the + # generated SASS is identical to hand-unrolling each ilp_rows value, but + # one loop covers ilp_rows ∈ {2, 4, 8}. + num_chunks: cutlass.Constexpr[int] = rows_per_group // ilp_rows + for chunk in cutlass.range_constexpr(num_chunks): + v_idx_0 = i_v * tile_v + group_idx * rows_per_group + chunk * ilp_rows + if v_idx_0 + (ilp_rows - 1) < V: + # Load ilp_rows h-state rows ONCE; they stay register-resident across T. + for slot in cutlass.range_constexpr(ilp_rows): + h_tile = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_0 + slot, lane_in_group), + ) + cute.autovec_copy(h_tile, cute.slice_(r_h, (slot, None))) + + for i_t in cutlass.range_constexpr(T): + # ---- inline q/k load for this t ---- + q_tile = cute.local_tile( + q, + (1, 1, 1, vec_size), + (i_n, i_t, i_h, lane_in_group), + ) + k_tile = cute.local_tile( + k, + (1, 1, 1, vec_size), + (i_n, i_t, i_h, lane_in_group), + ) + cute.autovec_copy(q_tile, r_q) + cute.autovec_copy(k_tile, r_k) + for i in cutlass.range_constexpr(vec_size): + r_q[i] = r_q[i] * scale + + # Per-row dot-product accumulators (lo, hi) — zeroed each t step. + r_dot_lo = cute.make_rmem_tensor(cute.make_layout((ilp_rows,), stride=(1,)), cutlass.Float32) + r_dot_hi = cute.make_rmem_tensor(cute.make_layout((ilp_rows,), stride=(1,)), cutlass.Float32) + for slot in cutlass.range_constexpr(ilp_rows): + r_dot_lo[slot] = cutlass.Float32(0.0) + r_dot_hi[slot] = cutlass.Float32(0.0) + + # ---- fused decay + rank-1 update (per V-row) ---- + for slot in cutlass.range_constexpr(ilp_rows): + r_v_s = cutlass.Float32(v[i_n, i_t, i_hv, v_idx_0 + slot]) + for j in cutlass.range_constexpr(0, vec_size, 2): + r_h[slot, j], r_h[slot, j + 1] = la_update_pair( + r_h[slot, j], + r_h[slot, j + 1], + r_k[j], + r_k[j + 1], + r_v_s, + r_decay, + use_packed_fma, + ) + + # ---- optional intermediate-state cache ---- + if cutlass.const_expr(cache_intermediate_states): + flat_idx = i_n * T * HV + i_t * HV + i_hv + for slot in cutlass.range_constexpr(ilp_rows): + inter_tile = cute.local_tile( + intermediate_states, + (1, 1, vec_size), + (flat_idx, v_idx_0 + slot, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (slot, None)), inter_tile) + + # ---- o_t = h_t @ q_t (per-row warp reduce) ---- + for slot in cutlass.range_constexpr(ilp_rows): + for j in cutlass.range_constexpr(0, vec_size, 2): + r_dot_lo[slot], r_dot_hi[slot] = hq_dot_pair( + r_h[slot, j], + r_h[slot, j + 1], + r_q[j], + r_q[j + 1], + r_dot_lo[slot], + r_dot_hi[slot], + use_packed_fma, + ) + r_acc = r_dot_lo[slot] + r_dot_hi[slot] + for offset in [16, 8, 4, 2, 1]: + r_acc += cute.arch.shuffle_sync_bfly(r_acc, offset=offset, mask=-1, mask_and_clamp=31) + r_dot_lo[slot] = r_acc # reuse slot for final result + + # ---- writeback ---- + if lane_in_group == 0: + for slot in cutlass.range_constexpr(ilp_rows): + o[(i_n, i_t, i_hv, v_idx_0 + slot)] = r_dot_lo[slot] + + # Final state writeback + if cutlass.const_expr(not disable_state_update): + for slot in cutlass.range_constexpr(ilp_rows): + h_tile_out = cute.local_tile( + h0_source, + (1, 1, vec_size), + (flat_state_idx, v_idx_0 + slot, lane_in_group), + ) + cute.autovec_copy(cute.slice_(r_h, (slot, None)), h_tile_out) + + +# ============================================================================ +# Launcher +# ============================================================================ +@cute.jit +def run_la_verify_kernel_mtp( + h0_source: cute.Tensor, + intermediate_states: cute.Tensor, + decay_scales: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + scale: cutlass.Constexpr[float], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + vec_size: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + disable_state_update: cutlass.Constexpr[bool], + cache_intermediate_states: cutlass.Constexpr[bool], + is_varlen: cutlass.Constexpr[bool], + stream: cuda.CUstream, +): + _, v_dim, _ = ( + h0_source.layout.shape[0], + h0_source.layout.shape[1], + h0_source.layout.shape[2], + ) + + num_v_tiles = cute.ceil_div(v_dim, tile_v) + grid_size = B * HV * num_v_tiles + + # LA decode uses no SMEM (v has no cross-row reuse; grid search confirmed the + # direct-global path wins). Reserve a small alignment slack only. + smem_bytes = 128 + + la_verify_kernel_mtp( + h0_source, + intermediate_states, + decay_scales, + q, + k, + v, + o, + h0_indices, + cu_seqlens, + vec_size, + num_v_tiles, + tile_v, + scale, + B, + T, + H, + HV, + K, + V, + disable_state_update, + cache_intermediate_states, + is_varlen, + ilp_rows, + use_packed_fma, + ).launch( + grid=(grid_size, 1, 1), + block=[NUM_THREADS_MTP, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +# ============================================================================ +# Compile cache +# ============================================================================ +@functools.cache +def _get_compiled_la_mtp_kernel( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + softmax_scale: float, + disable_state_update: bool, + cache_intermediate_states: bool, + is_varlen: bool, + tile_v: int, + vec_size: int, + ilp_rows: int, + use_packed_fma: bool, +): + return {} + + +def _la_mtp_compile_cache( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + softmax_scale: float, + *, + disable_state_update: bool, + cache_intermediate_states: bool, + is_varlen: bool, + device: torch.device, +): + """Return (cache dict, kernel config tuple) for the given launch parameters.""" + tile_v, vec_size, ilp_rows = get_mtp_config(B, T, HV, V, disable_state_update) + assert V % ilp_rows == 0, f"V={V} % ilp_rows={ilp_rows} ≠ 0: partial row-blocks would be silently skipped" + use_packed_fma = get_device_sm_version(device)[0] >= 10 + cache = _get_compiled_la_mtp_kernel( + B, + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + disable_state_update, + cache_intermediate_states, + is_varlen, + tile_v, + vec_size, + ilp_rows, + use_packed_fma, + ) + return cache, (tile_v, vec_size, ilp_rows, use_packed_fma) + + +def get_compiled_la_mtp_handle( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + softmax_scale: float, + device: torch.device, + *, + disable_state_update: bool, + cache_intermediate_states: bool, + is_varlen: bool = False, +): + """Return a pre-compiled MTP kernel handle (benchmark kernel-only path). + + Call ``linear_attention_decode_mtp`` once with the same config first so the + cache entry is populated. + """ + cache, _ = _la_mtp_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, + is_varlen=is_varlen, + device=device, + ) + compiled = cache.get("compiled") + if compiled is None: + raise RuntimeError("MTP kernel not compiled for this config; call linear_attention_decode_mtp once first.") + return compiled + + +# ============================================================================ +# Public Python entry point +# ============================================================================ +def linear_attention_decode_mtp( + q: torch.Tensor, # [B, T, H, K] fp32 + k: torch.Tensor, # [B, T, H, K] fp32 + v: torch.Tensor, # [B, T, HV, V] fp32 + s: torch.Tensor, # [pool_size, HV, V, K] fp32 + intermediate_states: torch.Tensor, # [pool_size*T*HV, V, K] fp32 (or dummy) + out: torch.Tensor, # [B, T, HV, V] fp32 + decay_scales: torch.Tensor, # [H] fp32 + s_offsets: torch.Tensor, # [B] int32 (-1 to skip) + cu_seqlens: torch.Tensor, # [B+1] int32 (reserved; see note below) + softmax_scale: float, + T: int, + cache_intermediate_states: bool, + disable_state_update: bool, + is_varlen: bool, +) -> None: + """ + Lightning Attention multi-token decode (T > 1). + + Writes to ``out``; updates ``s`` in place unless ``disable_state_update`` is True; + writes ``intermediate_states`` when ``cache_intermediate_states`` is True. + + NOTE: For any batch ``i`` where ``s_offsets[i] < 0`` the kernel skips that batch + entirely — ``out[i]`` is LEFT UNCHANGED, and neither ``s`` nor + ``intermediate_states`` is written for that slot. Callers must initialize ``out`` + to a known value (e.g. ``torch.zeros``) before the call if any downstream code + may read those slots. + + NOTE: ``is_varlen`` and ``cu_seqlens`` are reserved in the signature to keep the + public API stable, but the early-stop branch is NOT implemented yet — same as + upstream flashinfer GDN MTP, which also exposes the flag without consuming it. + Callers should pass ``is_varlen=False`` and any int32 tensor for ``cu_seqlens``. + The kernel descriptor is built with ``assumed_align=16``, so even the dummy + ``cu_seqlens`` must be 16-byte aligned; pass a fresh ``torch.empty(N, dtype=int32)`` + (CUDA allocator guarantees alignment) — do NOT pass a slice that may misalign. + """ + B, T_q, H, K = q.shape + assert T_q == T, f"q.shape[1]={T_q} doesn't match T={T}" + _, _, HV, V = v.shape + pool_size = s.shape[0] + if q.dtype != torch.float32 or k.dtype != torch.float32 or v.dtype != torch.float32: + raise ValueError(f"q/k/v must be torch.float32, got {q.dtype}/{k.dtype}/{v.dtype}") + if s.dtype != torch.float32: + raise ValueError(f"s must be torch.float32, got {s.dtype}") + if intermediate_states.dtype != torch.float32: + raise ValueError(f"intermediate_states must be torch.float32, got {intermediate_states.dtype}") + if out.dtype != torch.float32: + raise ValueError(f"out must be torch.float32, got {out.dtype}") + if decay_scales.dtype != torch.float32: + raise ValueError(f"decay_scales must be torch.float32, got {decay_scales.dtype}") + if s_offsets.dtype != torch.int32 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"s_offsets/cu_seqlens must be torch.int32, got {s_offsets.dtype}/{cu_seqlens.dtype}") + + cache, (tile_v, vec_size, ilp_rows, use_packed_fma) = _la_mtp_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, + is_varlen=is_varlen, + device=q.device, + ) + + h0_view = s.view(pool_size * HV, V, K) + + if "compiled" not in cache: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + compiled = cute.compile( + run_la_verify_kernel_mtp, + from_dlpack(h0_view, assumed_align=16), + from_dlpack(intermediate_states, assumed_align=16), + from_dlpack(decay_scales, assumed_align=16), + from_dlpack(q, assumed_align=16), + from_dlpack(k, assumed_align=16), + from_dlpack(v, assumed_align=16), + from_dlpack(out, assumed_align=16), + from_dlpack(s_offsets, assumed_align=16), + from_dlpack(cu_seqlens, assumed_align=16), + scale=softmax_scale, + B=B, + T=T, + H=H, + HV=HV, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, + is_varlen=is_varlen, + stream=stream, + options="--enable-tvm-ffi", + ) + cache["compiled"] = compiled + + compiled = cache["compiled"] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled( + h0_view, + intermediate_states, + decay_scales, + q, + k, + v, + out, + s_offsets, + cu_seqlens, + stream, + ) diff --git a/cula/lightning/la_state_update_kvbuffer.py b/cula/lightning/la_state_update_kvbuffer.py new file mode 100644 index 00000000..81d4959d --- /dev/null +++ b/cula/lightning/la_state_update_kvbuffer.py @@ -0,0 +1,641 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Lightning Attention KVBuffer state-update kernel (paper Eq. 8 for LA). + +After a parallel-verify cycle, advances the pooled state from h_init to +h_state_L for a per-batch accepted prefix length L = accepted_len[b]: + + h_running = h_init + for i in 0..L-1: + h_running = exp(-decay_scales[h]) * h_running + k_i ⊗ v_i + s[cache_idx] = h_running + +The loop body is bit-identical to the baseline T-loop body, so at L == T the +result is bit-equivalent to running the baseline with disable_state_update=False. + +Reads s and pool-indexed k_buf/v_buf; writes s. Never touches q or o. + +Grid: (B * HV * num_v_tiles, 1, 1), 128 threads/block — identical layout to the +baseline verify kernel, so the state write aligns with the verify kernel's h0 read. +""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import ( + from_dlpack, + make_fake_compact_tensor, + make_fake_stream, +) +from cutlass.cute.typing import Int32 + +from cula.lightning.la_decode_mtp import ( + NUM_THREADS_MTP, + la_update_pair, +) +from cula.lightning.la_verify_kvbuffer import get_mtp_config +from cula.utils import USE_FAST_MATH, get_device_sm_version + + +@cute.kernel +def la_state_update_kernel( + h0_source: cute.Tensor, # [pool_size * HV, V, K] fp32 (read + written in place) + decay_scales: cute.Tensor, # [H] fp32 + h0_indices: cute.Tensor, # [B] int32 + accepted_len: cute.Tensor, # [B] int32 + k_buf: cute.Tensor, # [pool_size, T, H, K] fp32 + v_buf: cute.Tensor, # [pool_size, T, HV, V] fp32 + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + threads_per_group: cutlass.Constexpr[int] = K // vec_size # 32 + groups_per_warp: cutlass.Constexpr[int] = 32 // threads_per_group # 1 + num_groups: cutlass.Constexpr[int] = 4 * groups_per_warp # 4 + + lane_in_group = lane_id % threads_per_group + group_in_warp = lane_id // threads_per_group + group_idx = warp_idx * groups_per_warp + group_in_warp + + block_idx, _, _ = cute.arch.block_idx() + i_v = block_idx % num_v_tiles + tmp = block_idx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + L = accepted_len[i_n] + + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_h = cute.make_rmem_tensor(cute.make_layout((8, vec_size), stride=(vec_size, 1)), cutlass.Float32) + + if cache_idx >= 0 and L > 0: + r_decay = cute.exp(-cutlass.Float32(decay_scales[i_h]), fastmath=USE_FAST_MATH) + rows_per_group: cutlass.Constexpr[int] = tile_v // num_groups + flat_state_idx = cache_idx * HV + i_hv + + # Process `ilp_rows` V-rows per iteration. ilp_rows is a compile-time + # constant, so range_constexpr fully unrolls the slot loops below — the + # generated SASS is identical to hand-unrolling each ilp_rows value, but + # one loop covers ilp_rows in {2, 4, 8}. + num_chunks: cutlass.Constexpr[int] = rows_per_group // ilp_rows + for chunk in cutlass.range_constexpr(num_chunks): + v_idx_0 = i_v * tile_v + group_idx * rows_per_group + chunk * ilp_rows + if v_idx_0 + (ilp_rows - 1) < V: + # Load the ilp_rows h-state rows this thread owns into registers. + for slot in cutlass.range_constexpr(ilp_rows): + h_tile = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_idx_0 + slot, lane_in_group)) + cute.autovec_copy(h_tile, cute.slice_(r_h, (slot, None))) + + # Recurrence: h = decay * h + k_i (x) v_i, for i in 0..L-1. + for i in cutlass.range(0, L, unroll=0): + k_tile = cute.local_tile(k_buf, (1, 1, 1, vec_size), (cache_idx, i, i_h, lane_in_group)) + cute.autovec_copy(k_tile, r_k) + for slot in cutlass.range_constexpr(ilp_rows): + r_v_s = cutlass.Float32(v_buf[cache_idx, i, i_hv, v_idx_0 + slot]) + for j in cutlass.range_constexpr(0, vec_size, 2): + r_h[slot, j], r_h[slot, j + 1] = la_update_pair( + r_h[slot, j], r_h[slot, j + 1], r_k[j], r_k[j + 1], r_v_s, r_decay, use_packed_fma + ) + + # Write the advanced state back in place. + for slot in cutlass.range_constexpr(ilp_rows): + h_out = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_idx_0 + slot, lane_in_group)) + cute.autovec_copy(cute.slice_(r_h, (slot, None)), h_out) + + +@cute.jit +def run_la_state_update_kernel( + h0_source: cute.Tensor, + decay_scales: cute.Tensor, + h0_indices: cute.Tensor, + accepted_len: cute.Tensor, + k_buf: cute.Tensor, + v_buf: cute.Tensor, + B: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + vec_size: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + stream: cuda.CUstream, +): + num_v_tiles: cutlass.Constexpr[int] = (V + tile_v - 1) // tile_v + grid_size = B * HV * num_v_tiles + + la_state_update_kernel( + h0_source, + decay_scales, + h0_indices, + accepted_len, + k_buf, + v_buf, + vec_size, + num_v_tiles, + tile_v, + B, + T, + H, + HV, + K, + V, + ilp_rows, + use_packed_fma, + ).launch( + grid=(grid_size, 1, 1), + block=[NUM_THREADS_MTP, 1, 1], + stream=stream, + ) + + +@functools.cache +def _get_compiled_state_update_kernel( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + tile_v: int, + vec_size: int, + ilp_rows: int, + use_packed_fma: bool, +): + return {} + + +def _state_update_compile_cache( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + *, + device: torch.device, +): + """Return (cache dict, tile config tuple) for the given launch parameters.""" + tile_v, vec_size, ilp_rows = get_mtp_config(B, T, HV, V) + assert V % ilp_rows == 0, f"V={V} % ilp_rows={ilp_rows} ≠ 0: partial row-blocks would be silently skipped" + use_packed_fma = get_device_sm_version(device)[0] >= 10 + cache = _get_compiled_state_update_kernel( + B, + T, + H, + HV, + K, + V, + pool_size, + tile_v, + vec_size, + ilp_rows, + use_packed_fma, + ) + return cache, (tile_v, vec_size, ilp_rows, use_packed_fma) + + +def get_compiled_state_update_kvbuffer_handle( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + *, + device: torch.device, +): + """Return a pre-compiled state-update kernel handle (benchmark kernel-only path). + + Call ``linear_attention_state_update_kvbuffer`` once with the same config first. + """ + cache, _ = _state_update_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + device=device, + ) + compiled = cache.get("compiled") + if compiled is None: + raise RuntimeError( + "State-update kernel not compiled for this config; call linear_attention_state_update_kvbuffer once first." + ) + return compiled + + +def linear_attention_state_update_kvbuffer( + k_buf: torch.Tensor, # [pool_size, T, H, K] fp32 + v_buf: torch.Tensor, # [pool_size, T, HV, V] fp32 + s: torch.Tensor, # [pool_size, HV, V, K] fp32, WRITTEN IN PLACE + decay_scales: torch.Tensor, # [H] fp32 + h0_indices: torch.Tensor, # [B] int32, -1 to skip + accepted_len: torch.Tensor, # [B] int32, in [0, T] + T: int, +) -> None: + """ + Advance pooled state from h_init to h_state_L per batch (KVBuffer Eq. 8). + + Reads k/v from fp32 pool-indexed buffers. This matches the SGLang/Ling + integration path: verify writes per-layer draft k/v into the request pool, + then commit advances the fp32 temporal state directly from those buffers. + """ + pool_size, T_k, H, K = k_buf.shape + assert T_k == T, f"k.shape[1]={T_k} doesn't match T={T}" + assert K == 128, f"K={K} != 128: kernel hardcodes K=128 (threads_per_group, lane K-coverage)" + if k_buf.dtype != torch.float32 or v_buf.dtype != torch.float32: + raise ValueError(f"k_buf/v_buf must be torch.float32, got {k_buf.dtype}/{v_buf.dtype}") + if s.dtype != torch.float32: + raise ValueError(f"s must be torch.float32, got {s.dtype}") + if decay_scales.dtype != torch.float32: + raise ValueError(f"decay_scales must be torch.float32, got {decay_scales.dtype}") + if h0_indices.dtype != torch.int32 or accepted_len.dtype != torch.int32: + raise ValueError(f"h0_indices/accepted_len must be torch.int32, got {h0_indices.dtype}/{accepted_len.dtype}") + if s.shape[0] != pool_size: + raise ValueError(f"s pool_size={s.shape[0]} doesn't match k_buf pool_size={pool_size}") + if v_buf.shape[:3] != (pool_size, T, s.shape[1]): + raise ValueError(f"v_buf shape {tuple(v_buf.shape)} doesn't match expected prefix {(pool_size, T, s.shape[1])}") + HV, V = s.shape[1], s.shape[2] + if v_buf.shape != (pool_size, T, HV, V): + raise ValueError(f"v_buf shape {tuple(v_buf.shape)} doesn't match expected {(pool_size, T, HV, V)}") + if s.shape[3] != K: + raise ValueError(f"s K={s.shape[3]} doesn't match k_buf K={K}") + if decay_scales.shape[0] != H: + raise ValueError(f"decay_scales length={decay_scales.shape[0]} doesn't match H={H}") + B = h0_indices.shape[0] + + cache, (tile_v, vec_size, ilp_rows, use_packed_fma) = _state_update_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + device=k_buf.device, + ) + + h0_view = s.view(pool_size * HV, V, K) + + if "compiled" not in cache: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled = cute.compile( + run_la_state_update_kernel, + from_dlpack(h0_view, assumed_align=16), + from_dlpack(decay_scales, assumed_align=16), + from_dlpack(h0_indices, assumed_align=16), + from_dlpack(accepted_len, assumed_align=16), + from_dlpack(k_buf, assumed_align=16), + from_dlpack(v_buf, assumed_align=16), + B=B, + T=T, + H=H, + HV=HV, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + stream=stream, + options="--enable-tvm-ffi", + ) + cache["compiled"] = compiled + + compiled = cache["compiled"] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled( + h0_view, + decay_scales, + h0_indices, + accepted_len, + k_buf, + v_buf, + stream, + ) + + +# --------------------------------------------------------------------------- +# Layer-fused state-update: one launch advances ALL mamba layers in parallel. +# Replaces the per-layer Python loop (28 FFI launches -> 1). Grid gains a layer +# dimension; k_buf/v_buf/h0_source/decay_scales all gain a leading num_layers +# dim and are indexed by i_layer. Pool-indexed k/v semantics, so no +# host-side gather is needed either. +# --------------------------------------------------------------------------- +@cute.kernel +def la_state_update_kernel_fused( + h0_source: cute.Tensor, + decay_scales: cute.Tensor, + k_buf: cute.Tensor, + v_buf: cute.Tensor, + h0_indices: cute.Tensor, + accepted_len: cute.Tensor, + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + num_layers: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + threads_per_group: cutlass.Constexpr[int] = K // vec_size + groups_per_warp: cutlass.Constexpr[int] = 32 // threads_per_group + num_groups: cutlass.Constexpr[int] = 4 * groups_per_warp + + lane_in_group = lane_id % threads_per_group + group_in_warp = lane_id // threads_per_group + group_idx = warp_idx * groups_per_warp + group_in_warp + + # 3D grid: (HV * num_v_tiles, B, num_layers) — B is a runtime grid dim. + block_idx_x, block_idx_y, block_idx_z = cute.arch.block_idx() + i_v = block_idx_x % num_v_tiles + i_hv = block_idx_x // num_v_tiles + i_n = block_idx_y + i_layer = block_idx_z + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + L = accepted_len[i_n] + + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_h = cute.make_rmem_tensor(cute.make_layout((8, vec_size), stride=(vec_size, 1)), cutlass.Float32) + + if cache_idx >= 0 and L > 0: + r_decay = cute.exp(-cutlass.Float32(decay_scales[i_layer, i_h]), fastmath=USE_FAST_MATH) + rows_per_group: cutlass.Constexpr[int] = tile_v // num_groups + flat_state_idx = cache_idx * HV + i_hv + + num_chunks: cutlass.Constexpr[int] = rows_per_group // ilp_rows + for chunk in cutlass.range_constexpr(num_chunks): + v_idx_0 = i_v * tile_v + group_idx * rows_per_group + chunk * ilp_rows + if v_idx_0 + (ilp_rows - 1) < V: + for slot in cutlass.range_constexpr(ilp_rows): + h_tile = cute.local_tile( + h0_source, (1, 1, 1, vec_size), (i_layer, flat_state_idx, v_idx_0 + slot, lane_in_group) + ) + cute.autovec_copy(h_tile, cute.slice_(r_h, (slot, None))) + + for i in cutlass.range(0, L, unroll=0): + k_tile = cute.local_tile(k_buf, (1, 1, 1, 1, vec_size), (i_layer, cache_idx, i, i_h, lane_in_group)) + cute.autovec_copy(k_tile, r_k) + for slot in cutlass.range_constexpr(ilp_rows): + r_v_s = cutlass.Float32(v_buf[i_layer, cache_idx, i, i_hv, v_idx_0 + slot]) + for j in cutlass.range_constexpr(0, vec_size, 2): + r_h[slot, j], r_h[slot, j + 1] = la_update_pair( + r_h[slot, j], r_h[slot, j + 1], r_k[j], r_k[j + 1], r_v_s, r_decay, use_packed_fma + ) + + for slot in cutlass.range_constexpr(ilp_rows): + h_out = cute.local_tile( + h0_source, (1, 1, 1, vec_size), (i_layer, flat_state_idx, v_idx_0 + slot, lane_in_group) + ) + cute.autovec_copy(cute.slice_(r_h, (slot, None)), h_out) + + +@cute.jit +def run_la_state_update_kernel_fused( + h0_source: cute.Tensor, + decay_scales: cute.Tensor, + k_buf: cute.Tensor, + v_buf: cute.Tensor, + h0_indices: cute.Tensor, + accepted_len: cute.Tensor, + grid_y: Int32, + num_layers: cutlass.Constexpr[int], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + vec_size: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + stream: cuda.CUstream, +): + num_v_tiles: cutlass.Constexpr[int] = (V + tile_v - 1) // tile_v + + la_state_update_kernel_fused( + h0_source, + decay_scales, + k_buf, + v_buf, + h0_indices, + accepted_len, + vec_size, + num_v_tiles, + tile_v, + num_layers, + T, + H, + HV, + K, + V, + ilp_rows, + use_packed_fma, + ).launch( + grid=(HV * num_v_tiles, grid_y, num_layers), + block=[NUM_THREADS_MTP, 1, 1], + stream=stream, + ) + + +@functools.cache +def _get_compiled_state_update_kernel_fused( + num_layers: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + tile_v: int, + vec_size: int, + ilp_rows: int, + use_packed_fma: bool, +): + return {} + + +def linear_attention_state_update_kvbuffer_fused( + k_buf: torch.Tensor, + v_buf: torch.Tensor, + s: torch.Tensor, + decay_scales: torch.Tensor, + h0_indices: torch.Tensor, + accepted_len: torch.Tensor, + T: int, +) -> None: + num_layers, pool_size_l, HV, V, K = s.shape + num_layers_k, pool_size_k, T_k, H, K_k = k_buf.shape + assert T_k == T, f"k_buf T={T_k} doesn't match T={T}" + assert K_k == 128, f"K={K_k} != 128" + if num_layers_k != num_layers: + raise ValueError(f"k_buf num_layers={num_layers_k} doesn't match state num_layers={num_layers}") + if k_buf.dtype != torch.float32 or v_buf.dtype != torch.float32: + raise ValueError(f"k_buf/v_buf must be torch.float32, got {k_buf.dtype}/{v_buf.dtype}") + if s.dtype != torch.float32: + raise ValueError(f"s must be torch.float32, got {s.dtype}") + if decay_scales.dtype != torch.float32: + raise ValueError(f"decay_scales must be torch.float32, got {decay_scales.dtype}") + if h0_indices.dtype != torch.int32 or accepted_len.dtype != torch.int32: + raise ValueError(f"h0_indices/accepted_len must be torch.int32, got {h0_indices.dtype}/{accepted_len.dtype}") + if pool_size_k != pool_size_l: + raise ValueError(f"k_buf pool_size={pool_size_k} doesn't match state pool_size={pool_size_l}") + if k_buf.shape != (num_layers, pool_size_l, T, H, K): + raise ValueError(f"k_buf shape {tuple(k_buf.shape)} doesn't match expected {(num_layers, pool_size_l, T, H, K)}") + if v_buf.shape[:4] != (num_layers, pool_size_l, T, HV): + raise ValueError(f"v_buf shape {tuple(v_buf.shape)} doesn't match expected prefix {(num_layers, pool_size_l, T, HV)}") + if v_buf.shape[-1] != V: + raise ValueError(f"v_buf V={v_buf.shape[-1]} doesn't match state V={V}") + if decay_scales.shape != (num_layers, H): + raise ValueError(f"decay_scales shape {tuple(decay_scales.shape)} doesn't match expected {(num_layers, H)}") + B = h0_indices.shape[0] + if accepted_len.shape[0] != B: + raise ValueError(f"accepted_len length={accepted_len.shape[0]} doesn't match h0_indices length={B}") + + tile_v, vec_size, ilp_rows = get_mtp_config(B, T, HV, V) + assert V % ilp_rows == 0, f"V={V} % ilp_rows={ilp_rows} != 0" + use_packed_fma = get_device_sm_version(k_buf.device)[0] >= 10 + + cache = _get_compiled_state_update_kernel_fused( + num_layers, + T, + H, + HV, + K, + V, + pool_size_l, + tile_v, + vec_size, + ilp_rows, + use_packed_fma, + ) + + h0_view = s.view(num_layers, pool_size_l * HV, V, K) + + if "compiled" not in cache: + sym_b = cute.sym_int() + pool_hv = pool_size_l * HV + h0_fake = make_fake_compact_tensor( + cutlass.Float32, + (num_layers, pool_hv, V, K), + stride_order=(3, 2, 1, 0), + assumed_align=16, + ) + decay_fake = make_fake_compact_tensor( + cutlass.Float32, + (num_layers, H), + stride_order=(1, 0), + assumed_align=16, + ) + k_buf_fake = make_fake_compact_tensor( + cutlass.Float32, + (num_layers, pool_size_l, T, H, K), + stride_order=(4, 3, 2, 1, 0), + assumed_align=16, + ) + v_buf_fake = make_fake_compact_tensor( + cutlass.Float32, + (num_layers, pool_size_l, T, HV, V), + stride_order=(4, 3, 2, 1, 0), + assumed_align=16, + ) + idx_fake = make_fake_compact_tensor( + cutlass.Int32, + (sym_b,), + stride_order=(0,), + assumed_align=16, + ) + acc_fake = make_fake_compact_tensor( + cutlass.Int32, + (sym_b,), + stride_order=(0,), + assumed_align=16, + ) + stream_fake = make_fake_stream() + + compiled = cute.compile( + run_la_state_update_kernel_fused, + h0_fake, + decay_fake, + k_buf_fake, + v_buf_fake, + idx_fake, + acc_fake, + Int32(1), # grid_y (dummy B) + num_layers=num_layers, + T=T, + H=H, + HV=HV, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + stream=stream_fake, + options="--enable-tvm-ffi", + ) + cache["compiled"] = compiled + + compiled = cache["compiled"] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled( + h0_view, + decay_scales, + k_buf, + v_buf, + h0_indices, + accepted_len, + Int32(B), + stream, + ) diff --git a/cula/lightning/la_verify_kvbuffer.py b/cula/lightning/la_verify_kvbuffer.py new file mode 100644 index 00000000..4e6d6535 --- /dev/null +++ b/cula/lightning/la_verify_kvbuffer.py @@ -0,0 +1,1126 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Lightning Attention KVBuffer verify kernel (paper Eq. 7 for LA). + +Closed-form parallel verification — computes each draft step's output directly +from (h0, k, v) without materializing the intermediate states: + + o_t = alpha^{t+1} * (h0 @ q_t * scale) <- "term1" (HQ) + + sum_{i=0..t} alpha^{t-i} * (q_t . k_i) * scale * v_i <- "term2" (QK·V) + +The two dot-product GEMMs run on tensor cores via inline-PTX mma.sync.m16n8k8 +(TF32). Operands are staged in fp32 SMEM (manual fragment addressing — no +LdMatrix/StMatrix). Everything downstream of the GEMMs is plain scalar math. + +PARALLELISM + Grid: (B * HV * num_v_tiles, 1, 1) — one block per (sequence, v-head, V-tile) + Block: 128 threads = 4 warps. Each warp owns `rows_per_group` output V-rows. + +PIPELINE (per block) + Stage 0 cooperative load q*scale, k -> SMEM (sQ, sK) + Stage 1 GEMM2: QK[t,i] = q_t . k_i (warp 0 only) -> s_qk_scaled + Stage 2 per V-row-block: load h0 -> SMEM, GEMM1: HQ = h0 @ q_t, + then scalar combine term1+term2 -> o + +MMA m16n8k8 FRAGMENT MAP (lane = gid*4 + tig, gid=lane//4 in 0..7, tig=lane%4 in 0..3) + A[16,8] row-major : a0=A[gid,tig] a1=A[gid+8,tig] a2=A[gid,tig+4] a3=A[gid+8,tig+4] + B[8,8] col-major : b0=B[tig,gid] b1=B[tig+4,gid] + C[16,8] : c0=C[gid,2tig] c1=C[gid,2tig+1] c2=C[gid+8,2tig] c3=C[gid+8,2tig+1] + We only have 8 valid rows (BT=8), so A rows 8..15 are fed as zeros and the + corresponding outputs c2,c3 / e2,e3 are unused padding. +""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm as _llvm +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) +from cutlass.cute.typing import Int32 +from cutlass.cutlass_dsl import T as _T +from cutlass.cutlass_dsl import dsl_user_op + +from cula.lightning.la_decode_mtp import ( + NUM_THREADS_MTP, + hq_dot_pair, +) +from cula.utils import USE_FAST_MATH, get_device_sm_version + +# Dispatch threshold between the two verify implementations. +# The MMA (tensor-core) kernel wins at T>=4 (matches at T=4, +45% at T=8 vs the +# shuffle kernel), but the shuffle kernel wins at small T (T<=2) where the MMA +# GEMMs are under-utilised and its larger SMEM footprint caps occupancy. +# See docs/la_verify_kvbuffer_dev_history.md §6 for the full benchmark. +MMA_MIN_T: int = 4 + + +def get_mtp_config(B: int, T: int, HV: int, V: int) -> tuple: + """Pick (tile_v, vec_size, ilp_rows) for the verify + state-update kernels. + + Grid-searched on the LA verify kernel (B200, H=HV=64, K=V=128, + B ∈ [1..128], T ∈ [2,4,8]). This is the *opposite* regime from the decode + kernel's ``get_mtp_config``: the verify kernel runs tensor-core MMA GEMMs + (T>=4) and is compute-bound, so it wants LARGE tiles with ilp_rows=8 — more + V-rows per block amortize the q/k SMEM staging and fill the m16n8k8 MMA + tiles. Small tiles starve the tensor cores (75% output padding at tile_v=8). + + vs the old shared GDN thresholds (tile_v=64, ilp=4 for work_units>1024): + up to 1.6x faster on the MMA path at large B (e.g. B=128,T=4: 0.170→0.125 ms; + B=1024,T=4: 0.031→0.019 ms). At small work_units the kernel is + latency-bound, so the tile choice is immaterial (<2% spread). + + Returns a 3-tuple — the verify/state-update kernels have no ``use_smem_v`` + knob (they always stage v in SMEM). + """ + work_units = B * HV + vec_size = 4 + if work_units <= 512: + tile_v, ilp_rows = 64, 8 + else: + tile_v, ilp_rows = 128, 8 + + tile_v = min(tile_v, V) + rows_per_group = tile_v // 4 + assert rows_per_group % ilp_rows == 0, ( + f"tile_v={tile_v} / num_groups=4 / ilp_rows={ilp_rows} doesn't divide cleanly " + f"(rows_per_group={rows_per_group}); the ILP loop would run zero iterations." + ) + return tile_v, vec_size, ilp_rows + + +# --------------------------------------------------------------------------- +# Inline PTX mma.sync.m16n8k8.tf32 — copied from kda_decode_mtp_kvbuffer.py +# --------------------------------------------------------------------------- + + +@dsl_user_op +def _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """One mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32; returns (d0,d1,d2,d3).""" + f32 = _T.f32() + i32 = _T.i32() + + def _bits(v): + vv = v.ir_value(loc=loc, ip=ip) if hasattr(v, "ir_value") else v + return _arith.bitcast(i32, vv, loc=loc, ip=ip) + + def _f(v): + return v.ir_value(loc=loc, ip=ip) if hasattr(v, "ir_value") else v + + res_ty = _llvm.StructType.get_literal([f32, f32, f32, f32]) + res = _llvm.inline_asm( + res_ty, + [_bits(a0), _bits(a1), _bits(a2), _bits(a3), _bits(b0), _bits(b1), _f(c0), _f(c1), _f(c2), _f(c3)], + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {$0,$1,$2,$3}, {$4,$5,$6,$7}, {$8,$9}, {$10,$11,$12,$13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(_llvm.extractvalue(f32, res, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(_llvm.extractvalue(f32, res, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(_llvm.extractvalue(f32, res, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(_llvm.extractvalue(f32, res, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +BT: int = 8 # pad M and N dimensions to 8 for mma fragment + + +@cute.kernel +def la_verify_kvbuffer_kernel( + h0_source: cute.Tensor, # [pool_size * HV, V, K] fp32 (READ ONLY) + decay_scales: cute.Tensor, # [H] fp32 + q: cute.Tensor, # [B, T, H, K] fp32 + k: cute.Tensor, # [B, T, H, K] fp32 + v: cute.Tensor, # [B, T, HV, V] fp32 + o: cute.Tensor, # [B, T, HV, V] fp32 (WRITTEN) + h0_indices: cute.Tensor, # [B] int32 + k_buf: cute.Tensor, # [pool_size, T, H, K] fp32 (WRITTEN when write_kv) + v_buf: cute.Tensor, # [pool_size, T, HV, V] fp32 (WRITTEN when write_kv) + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + scale: cutlass.Constexpr[float], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + write_kv: cutlass.Constexpr[bool], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # MMA lane decomposition (see fragment map in module docstring). + gid = lane_id // 4 # 0..7: row index within the MMA tile + tig = lane_id % 4 # 0..3: k-pair within the current 8-wide K-slab + + # 4 warps/block; each warp owns a disjoint set of output V-rows. All 32 lanes + # of a warp cooperate over the full K dimension (K=128, vec_size=4). + NUM_WARPS: cutlass.Constexpr[int] = 4 + + # Block -> (sequence n, v-head hv, V-tile i_v); i_h maps the v-head to its q/k head. + block_idx, _, _ = cute.arch.block_idx() + i_v = block_idx % num_v_tiles + tmp = block_idx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + + # ---- Per-lane registers ---- + r_decay_pow = cute.make_rmem_tensor(cute.make_layout((T + 1,), stride=(1,)), cutlass.Float32) + r_q_f32 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k_f32 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + + # ---- SMEM (all fp32; MMA bitcasts fp32->TF32, no separate conversion) ---- + # KP = K+4 pads the row stride so 132%32=4: the gid*4+tig access pattern then + # hits 32 distinct banks, giving conflict-free SMEM reads in both GEMMs. + KP: cutlass.Constexpr[int] = K + 4 + smem = cutlass.utils.SmemAllocator() + # GEMM operands. sQ holds q*scale, doubles as GEMM2-A and GEMM1-B. + sQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, KP), stride=(KP, 1)), 16) + sK = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BT, KP), stride=(KP, 1)), 16) + # h0, one [BT, K] region per warp (each warp does GEMM1 for its own V-rows). + sH0 = smem.allocate_tensor(cutlass.Float32, cute.make_layout((NUM_WARPS, BT, KP), stride=(BT * KP, KP, 1)), 16) + # Decay-masked QK coefficients [T, T], produced by GEMM2, consumed by every warp. + s_qk_scaled = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, T), stride=(T, 1)), 16) + # v is lane-invariant within a warp; stage it once in SMEM and broadcast-read. + sVbuf = smem.allocate_tensor(cutlass.Float32, cute.make_layout((NUM_WARPS, T, BT), stride=(T * BT, BT, 1)), 16) + + if cache_idx >= 0: + alpha = cute.exp(-cutlass.Float32(decay_scales[i_h]), fastmath=USE_FAST_MATH) + + r_decay_pow[0] = cutlass.Float32(1.0) + for t in cutlass.range_constexpr(1, T + 1): + r_decay_pow[t] = r_decay_pow[t - 1] * alpha + + rows_per_group: cutlass.Constexpr[int] = tile_v // NUM_WARPS + flat_state_idx = cache_idx * HV + i_hv + + # ================================================================ + # Stage 0: cooperative load q*scale, k -> SMEM (sQ, sK), fp32. + # Warp w loads tokens {w, w+4, ...}; within a token, lane_id covers the + # K dimension (vec_size contiguous elements each). Rows T..BT-1 are the + # MMA M-padding and are zeroed. + # ================================================================ + tokens_per_warp: cutlass.Constexpr[int] = (BT + NUM_WARPS - 1) // NUM_WARPS + for tt in cutlass.range_constexpr(tokens_per_warp): + t_tok = tt * NUM_WARPS + warp_idx + if t_tok < T: + q_tile = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, t_tok, i_h, lane_id)) + k_tile = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, t_tok, i_h, lane_id)) + cute.autovec_copy(q_tile, r_q_f32) + cute.autovec_copy(k_tile, r_k_f32) + for c in cutlass.range_constexpr(vec_size): + col = lane_id * vec_size + c + sQ[(t_tok, col)] = cutlass.Float32(r_q_f32[c]) * scale + sK[(t_tok, col)] = cutlass.Float32(r_k_f32[c]) + # Persist k to the pool buffer while it is already in registers. + if cutlass.const_expr(write_kv): + if i_v == 0 and i_hv % (HV // H) == 0: + kb_tile = cute.local_tile(k_buf, (1, 1, 1, vec_size), (cache_idx, t_tok, i_h, lane_id)) + for c in cutlass.range_constexpr(vec_size): + kb_tile[c] = r_k_f32[c] + if t_tok >= T and t_tok < BT: + for c in cutlass.range_constexpr(vec_size): + col = lane_id * vec_size + c + sQ[(t_tok, col)] = cutlass.Float32(0.0) + sK[(t_tok, col)] = cutlass.Float32(0.0) + + cute.arch.barrier() + + # ================================================================ + # Stage 1: GEMM2 — QK[t,i] = q_t . k_i, accumulated over the full K. + # A = Q[8,K] (rows = tokens), B = K[8,K] read col-major as K^T. Warp 0 + # alone has enough lanes (M=N=T<=8), so the other warps skip this. + # ================================================================ + if warp_idx == 0: + c0 = cutlass.Float32(0.0) + c1 = cutlass.Float32(0.0) + c2 = cutlass.Float32(0.0) # c2,c3 = padding rows 8..15, unused + c3 = cutlass.Float32(0.0) + for ks in cutlass.range_constexpr(K // 8): + kb = ks * 8 + a0 = sQ[(gid, kb + tig)] + a1 = cutlass.Float32(0.0) + a2 = sQ[(gid, kb + tig + 4)] + a3 = cutlass.Float32(0.0) + b0 = sK[(gid, kb + tig)] + b1 = sK[(gid, kb + tig + 4)] + c0, c1, c2, c3 = _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3) + + # c0,c1 hold QK[gid, 2tig], QK[gid, 2tig+1]. Keep the causal lower + # triangle, pre-multiply by the decay alpha^{t-i}, store coefficients. + for fi in cutlass.range_constexpr(2): + row = gid + col = 2 * tig + fi + cv = c1 if cutlass.const_expr(fi == 1) else c0 + if row < T and col < T: + if col <= row: + s_qk_scaled[(row, col)] = r_decay_pow[row - col] * cv + else: + s_qk_scaled[(row, col)] = cutlass.Float32(0.0) + + cute.arch.barrier() + + # ================================================================ + # Stage 2: for each block of `ilp_rows` V-rows owned by this warp, + # load h0 -> SMEM, run GEMM1 (HQ = h0 @ q_t), then combine the two terms. + # ================================================================ + num_row_blocks: cutlass.Constexpr[int] = rows_per_group // ilp_rows + for row_block in cutlass.range_constexpr(num_row_blocks): + v_base = i_v * tile_v + warp_idx * rows_per_group + row_block * ilp_rows + if v_base + (ilp_rows - 1) < V: + # (a) Coalesced h0 load: lane_id indexes vec_size contiguous K + # elements, so the 32 lanes read one full contiguous row per step + # (no over-fetch). Each warp fills its own sH0 region. + sH0_w = sH0[(warp_idx, None, None)] # [BT, KP] + gH0 = h0_source[(flat_state_idx, None, None)] # [V, K] + for row in cutlass.range_constexpr(ilp_rows): + h_g = cute.local_tile(gH0, (1, vec_size), (v_base + row, lane_id)) + h_s = cute.local_tile(sH0_w, (1, vec_size), (row, lane_id)) + cute.autovec_copy(h_g, h_s) + # Zero the M-padding rows (ilp_rows..BT-1). GEMM1 reads all BT rows; + # their outputs are unused, but leaving stale/NaN SMEM as MMA inputs + # is unclean — explicitly zero so the fragment is well-defined. + for row in cutlass.range_constexpr(ilp_rows, BT): + for c in cutlass.range_constexpr(vec_size): + sH0_w[(row, lane_id * vec_size + c)] = cutlass.Float32(0.0) + cute.arch.sync_warp() # make sH0 writes visible to this warp's GEMM1 + + # (b) GEMM1: HQ[row, t] = h0_row . q_t, over the full K. + # A = sH0 (this warp's V-rows), B = sQ read col-major as Q^T. + e0 = cutlass.Float32(0.0) + e1 = cutlass.Float32(0.0) + e2 = cutlass.Float32(0.0) # e2,e3 = padding rows 8..15, unused + e3 = cutlass.Float32(0.0) + for ks in cutlass.range_constexpr(K // 8): + kb = ks * 8 + a0 = sH0[(warp_idx, gid, kb + tig)] + a1 = cutlass.Float32(0.0) + a2 = sH0[(warp_idx, gid, kb + tig + 4)] + a3 = cutlass.Float32(0.0) + b0 = sQ[(gid, kb + tig)] + b1 = sQ[(gid, kb + tig + 4)] + e0, e1, e2, e3 = _mma_m16n8k8_tf32(a0, a1, a2, a3, b0, b1, e0, e1, e2, e3) + # e0,e1 now hold HQ[gid, 2tig], HQ[gid, 2tig+1] (gid = V-row index). + + # (c) Stage v in SMEM (lane-invariant within the warp) and persist it. + if lane_id < ilp_rows: + for t in cutlass.range_constexpr(T): + vv = v[i_n, t, i_hv, v_base + lane_id] + sVbuf[(warp_idx, t, lane_id)] = cutlass.Float32(vv) + if cutlass.const_expr(write_kv): + v_buf[(cache_idx, t, i_hv, v_base + lane_id)] = cutlass.Float32(vv) + + # (d) Combine: o[t, row] = alpha^{t+1}*HQ[row,t] + sum_i qk[t,i]*v[i,row]. + # The (t, row) output grid has T*ilp_rows entries. Distribute them + # across the 32 lanes in a grid-stride fashion: lane L handles outputs + # L, L+32, L+64, ... so each lane emits ceil(T*ilp_rows/32) of them. + # This keeps every lane doing useful work for ANY T (T=4 -> 1 each, + # T=8 -> 2 each, T=2 -> half the lanes), with no redundant compute and + # no SMEM reshuffle — HQ is fetched straight from its owner lane. + num_out: cutlass.Constexpr[int] = T * ilp_rows + outs_per_lane: cutlass.Constexpr[int] = (num_out + 31) // 32 + for oj in cutlass.range_constexpr(outs_per_lane): + out_idx = lane_id + oj * 32 + my_t = out_idx // ilp_rows + my_slot = out_idx % ilp_rows + # shuffle_sync must execute on ALL lanes (warp-collective), so it + # stays outside the my_t= 10 + cache = _get_compiled_verify_kvbuffer_kernel_shuffle( + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + tile_v, + vec_size, + ilp_rows, + use_packed_fma, + write_kv, + ) + return cache, (tile_v, vec_size, ilp_rows, use_packed_fma) + + tile_v, vec_size, ilp_rows = get_mtp_config(B, T, HV, V) + assert T <= 8, f"T={T} > 8: MMA kernel's BT=8 token staging only covers T ≤ 8" + assert V % ilp_rows == 0, f"V={V} % ilp_rows={ilp_rows} ≠ 0: partial row-blocks would be silently skipped" + cache = _get_compiled_verify_kvbuffer_kernel( + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + tile_v, + vec_size, + ilp_rows, + write_kv, + ) + return cache, (tile_v, vec_size, ilp_rows, None) + + +def get_compiled_verify_kvbuffer_handle( + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + softmax_scale: float, + *, + write_kv: bool, + device: torch.device, +): + """Return a pre-compiled verify kernel handle (benchmark kernel-only path). + + Call ``linear_attention_verify_kvbuffer`` once with the same config first. + The returned handle has one stable signature for both MMA and shuffle paths: + ``(..., k_buf, v_buf, stream)``. It computes the runtime grid size internally. + """ + cache, (tile_v, _, _, _) = _verify_kvbuffer_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + write_kv=write_kv, + device=device, + ) + compiled = cache.get("compiled") + if compiled is None: + raise RuntimeError("Verify kernel not compiled for this config; call linear_attention_verify_kvbuffer once first.") + + num_v_tiles = (V + tile_v - 1) // tile_v + + def run_compiled( + h0_source, + decay_scales, + q, + k, + v, + o, + h0_indices, + k_buf, + v_buf, + stream, + ): + grid_size = q.shape[0] * HV * num_v_tiles + compiled( + h0_source, + decay_scales, + q, + k, + v, + o, + h0_indices, + k_buf, + v_buf, + Int32(grid_size), + stream, + ) + + return run_compiled + + +def linear_attention_verify_kvbuffer( + q: torch.Tensor, # [B, T, H, K] fp32 + k: torch.Tensor, # [B, T, H, K] fp32 + v: torch.Tensor, # [B, T, HV, V] fp32 + s: torch.Tensor, # [pool_size, HV, V, K] fp32, READ ONLY + out: torch.Tensor, # [B, T, HV, V] fp32, WRITTEN + decay_scales: torch.Tensor, # [H] fp32 + h0_indices: torch.Tensor, # [B] int32, -1 to skip + softmax_scale: float, + T: int, + k_buf: torch.Tensor | None = None, + v_buf: torch.Tensor | None = None, +) -> None: + """ + Closed-form parallel verify (KVBuffer Eq. 7). Writes out; does not touch s. + + When k_buf and v_buf are provided, also writes k,v to fp32 pool-indexed + buffers so the caller can free the original k,v tensors after this call + returns. This matches Ling/SGLang, where q/k/v arrive after fp32 RoPE and + are committed into a fp32 temporal state. + + Dispatches between two equivalent implementations by draft depth T: the + tensor-core MMA kernel below for T >= MMA_MIN_T, and the warp-shuffle kernel + for smaller T (where MMA's GEMMs are under-utilised). Both share the same + interface, grid, and KVBuffer write semantics. + """ + if T < MMA_MIN_T: + return linear_attention_verify_kvbuffer_shuffle( + q, + k, + v, + s, + out, + decay_scales, + h0_indices, + softmax_scale, + T, + k_buf=k_buf, + v_buf=v_buf, + ) + + B, T_q, H, K = q.shape + assert T_q == T, f"q.shape[1]={T_q} doesn't match T={T}" + assert K == 128, f"K={K} != 128: kernel hardcodes K=128 (threads_per_group, KP=K+4, lane K-coverage)" + _, _, HV, V = v.shape + pool_size = s.shape[0] + if q.dtype != torch.float32 or k.dtype != torch.float32 or v.dtype != torch.float32: + raise ValueError(f"q/k/v must be torch.float32, got {q.dtype}/{k.dtype}/{v.dtype}") + if s.dtype != torch.float32: + raise ValueError(f"s must be torch.float32, got {s.dtype}") + if out.dtype != torch.float32: + raise ValueError(f"out must be torch.float32, got {out.dtype}") + + write_kv = k_buf is not None and v_buf is not None + if (k_buf is None) != (v_buf is None): + raise ValueError("k_buf and v_buf must both be None or both be provided") + if write_kv and (k_buf.dtype != torch.float32 or v_buf.dtype != torch.float32): + raise ValueError(f"k_buf/v_buf must be torch.float32, got {k_buf.dtype}/{v_buf.dtype}") + + cache, (tile_v, vec_size, ilp_rows, _) = _verify_kvbuffer_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + write_kv=write_kv, + device=q.device, + ) + + h0_view = s.view(pool_size * HV, V, K) + + if not write_kv: + k_buf_t = torch.empty(1, T, H, K, device=q.device, dtype=torch.float32) + v_buf_t = torch.empty(1, T, HV, V, device=q.device, dtype=torch.float32) + else: + k_buf_t = k_buf + v_buf_t = v_buf + + if "compiled" not in cache: + # Use sym_int() for B so one compiled kernel handles all batch sizes + # (no per-B cute.compile JIT). Pattern from prefill (lightning_attn_sm100). + sym_b = cute.sym_int() + q_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, H, K), stride_order=(3, 2, 1, 0), assumed_align=16) + k_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, H, K), stride_order=(3, 2, 1, 0), assumed_align=16) + v_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, HV, V), stride_order=(3, 2, 1, 0), assumed_align=16) + o_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, HV, V), stride_order=(3, 2, 1, 0), assumed_align=16) + h0_fake = make_fake_compact_tensor(cutlass.Float32, (cute.sym_int(), V, K), stride_order=(2, 1, 0), assumed_align=16) + decay_fake = make_fake_compact_tensor(cutlass.Float32, (H,), stride_order=(0,), assumed_align=16) + idx_fake = make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,), assumed_align=16) + k_buf_fake = make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), T, H, K), stride_order=(3, 2, 1, 0), assumed_align=16 + ) + v_buf_fake = make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), T, HV, V), stride_order=(3, 2, 1, 0), assumed_align=16 + ) + stream_fake = make_fake_stream() + + compiled = cute.compile( + run_la_verify_kvbuffer_kernel, + h0_fake, + decay_fake, + q_fake, + k_fake, + v_fake, + o_fake, + idx_fake, + k_buf_fake, + v_buf_fake, + Int32(1), # grid_size (positional, before Constexpr kwargs) + scale=softmax_scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + ilp_rows=ilp_rows, + write_kv=write_kv, + stream=stream_fake, + options="--enable-tvm-ffi", + ) + cache["compiled"] = compiled + + compiled = cache["compiled"] + num_v_tiles_rt = (V + tile_v - 1) // tile_v + grid_size = B * HV * num_v_tiles_rt + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled( + h0_view, + decay_scales, + q, + k, + v, + out, + h0_indices, + k_buf_t, + v_buf_t, + Int32(grid_size), + stream, + ) + + +# =========================================================================== +# Warp-shuffle verify kernel (baseline). Dispatched for small T (T < MMA_MIN_T) +# by linear_attention_verify_kvbuffer above. Uses butterfly shuffle reduce for +# the dot products instead of tensor-core MMA — h0 stays in registers (no SMEM +# fragment staging), giving higher occupancy that wins when T is small. +# =========================================================================== + + +@cute.kernel +def la_verify_kvbuffer_shuffle_kernel( + h0_source: cute.Tensor, # [pool_size * HV, V, K] fp32 (READ ONLY) + decay_scales: cute.Tensor, # [H] fp32 + q: cute.Tensor, # [B, T, H, K] fp32 + k: cute.Tensor, # [B, T, H, K] fp32 + v: cute.Tensor, # [B, T, HV, V] fp32 + o: cute.Tensor, # [B, T, HV, V] fp32 (WRITTEN) + h0_indices: cute.Tensor, # [B] int32 + k_buf: cute.Tensor, # [pool_size, T, H, K] fp32 (WRITTEN when write_kv) + v_buf: cute.Tensor, # [pool_size, T, HV, V] fp32 (WRITTEN when write_kv) + vec_size: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + scale: cutlass.Constexpr[float], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + write_kv: cutlass.Constexpr[bool], +): + tidx, _, _ = cute.arch.thread_idx() + lane_id = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + threads_per_group: cutlass.Constexpr[int] = K // vec_size # 32 + groups_per_warp: cutlass.Constexpr[int] = 32 // threads_per_group # 1 + num_groups: cutlass.Constexpr[int] = 4 * groups_per_warp # 4 + + lane_in_group = lane_id % threads_per_group + group_in_warp = lane_id // threads_per_group + group_idx = warp_idx * groups_per_warp + group_in_warp + + block_idx, _, _ = cute.arch.block_idx() + i_v = block_idx % num_v_tiles + tmp = block_idx // num_v_tiles + i_hv = tmp % HV + i_n = tmp // HV + i_h = i_hv // (HV // H) + + cache_idx = h0_indices[i_n] + + r_q_f32 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_k_f32 = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + r_h = cute.make_rmem_tensor(cute.make_layout((8, vec_size), stride=(vec_size, 1)), cutlass.Float32) + r_decay_pow = cute.make_rmem_tensor(cute.make_layout((T + 1,), stride=(1,)), cutlass.Float32) + o_partial = cute.make_rmem_tensor(cute.make_layout((8,), stride=(1,)), cutlass.Float32) + + smem = cutlass.utils.SmemAllocator() + s_qk_scaled = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, T), stride=(T, 1)), 16) + # v staged to SMEM (block-shared over the whole v-tile). v has no K dim, so + # keeping it in per-lane registers wasted 8*T regs/thread and capped occupancy; + # SMEM costs only T*tile_v*4 bytes and is read warp-uniformly (broadcast). + sVdata = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T, tile_v), stride=(tile_v, 1)), 16) + # q (scaled) and k staged to SMEM. They depend only on lane_in_group (NOT on + # warp/group), so a single copy of 32 K-slices is shared by all 4 warps — + # this also removes the redundant per-warp q/k loads. Lane-minor layout + # (T, vec_size, 32) keeps the 32 lanes of a warp on consecutive banks + # (conflict-free); cost is 2 * T*vec_size*32*4 bytes (~8KB at T=8). + s_q = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((T, vec_size, threads_per_group), stride=(vec_size * threads_per_group, threads_per_group, 1)), + 16, + ) + s_k = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((T, vec_size, threads_per_group), stride=(vec_size * threads_per_group, threads_per_group, 1)), + 16, + ) + + if cache_idx >= 0: + alpha = cute.exp(-cutlass.Float32(decay_scales[i_h]), fastmath=USE_FAST_MATH) + + # alpha^0 .. alpha^T (T+1 powers; term1 uses alpha^{t+1}) + r_decay_pow[0] = cutlass.Float32(1.0) + for t in cutlass.range_constexpr(1, T + 1): + r_decay_pow[t] = r_decay_pow[t - 1] * alpha + + rows_per_group: cutlass.Constexpr[int] = tile_v // num_groups + flat_state_idx = cache_idx * HV + i_hv + + # Stage all T q (scaled) and k (fp32) into SMEM. q/k are warp-independent, + # so only warp 0 (its 32 lanes cover the full K dim) loads them once. + # The k_buf write is fused here, replacing the old per-warp redundant store. + if warp_idx == 0: + for t in cutlass.range_constexpr(T): + q_tile = cute.local_tile(q, (1, 1, 1, vec_size), (i_n, t, i_h, lane_id)) + k_tile = cute.local_tile(k, (1, 1, 1, vec_size), (i_n, t, i_h, lane_id)) + cute.autovec_copy(q_tile, r_q_f32) + cute.autovec_copy(k_tile, r_k_f32) + for j in cutlass.range_constexpr(vec_size): + s_q[(t, j, lane_id)] = cutlass.Float32(r_q_f32[j]) * scale + s_k[(t, j, lane_id)] = cutlass.Float32(r_k_f32[j]) + + # Write k to buffer — gated: only one block per (b, h, t) writes + if cutlass.const_expr(write_kv): + if i_v == 0 and i_hv % (HV // H) == 0: + kb_tile = cute.local_tile(k_buf, (1, 1, 1, vec_size), (cache_idx, t, i_h, lane_id)) + cute.autovec_copy(r_k_f32, kb_tile) + + # Cooperative v load: first tile_v threads each stage one v-row for all T + # steps into SMEM. v_buf write (when enabled) is fused here — every + # (cache_idx, t, hv, v_row) is written exactly once by its owning thread. + v_tile_start = i_v * tile_v + for t in cutlass.range_constexpr(T): + if tidx < tile_v: + v_global_idx = v_tile_start + tidx + if v_global_idx < V: + vv = v[i_n, t, i_hv, v_global_idx] + sVdata[(t, tidx)] = cutlass.Float32(vv) + if cutlass.const_expr(write_kv): + v_buf[(cache_idx, t, i_hv, v_global_idx)] = vv + + cute.arch.barrier() # q/k/v staged → visible to all warps + + # Phase 1: cooperative QK matrix — 4 warps split T*(T+1)/2 qk dot products. + # Warp w handles rows where min(t, T-1-t) % 4 == w (head-tail pairing) so that + # each warp's total row-length is balanced: heavy tail rows are paired with light + # head rows, making per-warp work ≈ T*(T+1)/8 regardless of T. + for t_assign in cutlass.range_constexpr(T): + if min(t_assign, T - 1 - t_assign) % 4 == warp_idx: + for i in cutlass.range_constexpr(t_assign + 1): + qk_lo = cutlass.Float32(0.0) + qk_hi = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(0, vec_size, 2): + qk_lo, qk_hi = hq_dot_pair( + s_q[t_assign, j, lane_in_group], + s_q[t_assign, j + 1, lane_in_group], + s_k[i, j, lane_in_group], + s_k[i, j + 1, lane_in_group], + qk_lo, + qk_hi, + use_packed_fma, + ) + qk = qk_lo + qk_hi + for offset in [16, 8, 4, 2, 1]: + qk += cute.arch.shuffle_sync_bfly(qk, offset=offset, mask=-1, mask_and_clamp=31) + if lane_in_group == 0: + s_qk_scaled[(t_assign, i)] = r_decay_pow[t_assign - i] * qk + + cute.arch.barrier() # s_qk_scaled written by Phase 1 → read by Phase 2 + + num_row_blocks: cutlass.Constexpr[int] = rows_per_group // ilp_rows + for row_block in cutlass.range_constexpr(num_row_blocks): + v_base = i_v * tile_v + group_idx * rows_per_group + row_block * ilp_rows + v_local = group_idx * rows_per_group + row_block * ilp_rows # offset within sVdata's v-tile + if v_base + (ilp_rows - 1) < V: + # Load h_init rows (persistent across the T loop). + for slot in cutlass.range_constexpr(ilp_rows): + h_tile = cute.local_tile(h0_source, (1, 1, vec_size), (flat_state_idx, v_base + slot, lane_in_group)) + cute.autovec_copy(h_tile, cute.slice_(r_h, (slot, None))) + + for t in cutlass.range_constexpr(T): + # term1: alpha^{t+1} * (h_init @ q_t) (per-slot warp reduce) + for slot in cutlass.range_constexpr(ilp_rows): + hq_lo = cutlass.Float32(0.0) + hq_hi = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(0, vec_size, 2): + hq_lo, hq_hi = hq_dot_pair( + r_h[slot, j], + r_h[slot, j + 1], + s_q[t, j, lane_in_group], + s_q[t, j + 1, lane_in_group], + hq_lo, + hq_hi, + use_packed_fma, + ) + hq = hq_lo + hq_hi + for offset in [16, 8, 4, 2, 1]: + hq += cute.arch.shuffle_sync_bfly(hq, offset=offset, mask=-1, mask_and_clamp=31) + o_partial[slot] = r_decay_pow[t + 1] * hq + + # term2: read pre-computed decay-scaled qk + staged v from SMEM + for i in cutlass.range_constexpr(t + 1): + coeff = s_qk_scaled[(t, i)] + for slot in cutlass.range_constexpr(ilp_rows): + o_partial[slot] = o_partial[slot] + coeff * sVdata[(i, v_local + slot)] + + # writeback (all lanes hold the reduced value; lane 0 writes) + if lane_in_group == 0: + for slot in cutlass.range_constexpr(ilp_rows): + o[(i_n, t, i_hv, v_base + slot)] = o_partial[slot] + + +@cute.jit +def run_la_verify_kvbuffer_shuffle_kernel( + h0_source: cute.Tensor, + decay_scales: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + h0_indices: cute.Tensor, + k_buf: cute.Tensor, + v_buf: cute.Tensor, + grid_size: Int32, + scale: cutlass.Constexpr[float], + T: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + vec_size: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + write_kv: cutlass.Constexpr[bool], + stream: cuda.CUstream, +): + num_v_tiles: cutlass.Constexpr[int] = (V + tile_v - 1) // tile_v + + # s_qk_scaled[T][T] + sVdata[T][tile_v] + s_q/s_k[T][vec_size][32] + threads_per_group = 32 + smem_bytes = ( + T * T * 4 # s_qk_scaled + + T * tile_v * 4 # sVdata + + 2 * T * vec_size * threads_per_group * 4 # s_q + s_k + + 4 * 16 # per-allocation 16B alignment padding (4 tensors) + ) + + la_verify_kvbuffer_shuffle_kernel( + h0_source, + decay_scales, + q, + k, + v, + o, + h0_indices, + k_buf, + v_buf, + vec_size, + num_v_tiles, + tile_v, + scale, + T, + H, + HV, + K, + V, + ilp_rows, + use_packed_fma, + write_kv, + ).launch( + grid=(grid_size, 1, 1), + block=[NUM_THREADS_MTP, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +@functools.cache +def _get_compiled_verify_kvbuffer_kernel_shuffle( + T: int, + H: int, + HV: int, + K: int, + V: int, + pool_size: int, + softmax_scale: float, + tile_v: int, + vec_size: int, + ilp_rows: int, + use_packed_fma: bool, + write_kv: bool, +): + return {} + + +def linear_attention_verify_kvbuffer_shuffle( + q: torch.Tensor, # [B, T, H, K] fp32 + k: torch.Tensor, # [B, T, H, K] fp32 + v: torch.Tensor, # [B, T, HV, V] fp32 + s: torch.Tensor, # [pool_size, HV, V, K] fp32, READ ONLY + out: torch.Tensor, # [B, T, HV, V] fp32, WRITTEN + decay_scales: torch.Tensor, # [H] fp32 + h0_indices: torch.Tensor, # [B] int32, -1 to skip + softmax_scale: float, + T: int, + k_buf: torch.Tensor | None = None, # [pool_size, T, H, K] fp32, WRITTEN + v_buf: torch.Tensor | None = None, # [pool_size, T, HV, V] fp32, WRITTEN +) -> None: + """ + Closed-form parallel verify (KVBuffer Eq. 7). Writes out; does not touch s. + + When k_buf and v_buf are provided, also writes k,v to fp32 pool-indexed + buffers so the caller can free the original k,v tensors after this call + returns. + + For batch b with h0_indices[b] < 0, out[b] is LEFT UNCHANGED — callers must + pre-initialize out if downstream code reads those slots. + """ + B, T_q, H, K = q.shape + assert T_q == T, f"q.shape[1]={T_q} doesn't match T={T}" + assert K == 128, f"K={K} != 128: kernel hardcodes K=128 (threads_per_group, lane K-coverage)" + _, _, HV, V = v.shape + pool_size = s.shape[0] + if q.dtype != torch.float32 or k.dtype != torch.float32 or v.dtype != torch.float32: + raise ValueError(f"q/k/v must be torch.float32, got {q.dtype}/{k.dtype}/{v.dtype}") + if s.dtype != torch.float32: + raise ValueError(f"s must be torch.float32, got {s.dtype}") + if out.dtype != torch.float32: + raise ValueError(f"out must be torch.float32, got {out.dtype}") + + write_kv = k_buf is not None and v_buf is not None + if (k_buf is None) != (v_buf is None): + raise ValueError("k_buf and v_buf must both be None or both be provided") + if write_kv and (k_buf.dtype != torch.float32 or v_buf.dtype != torch.float32): + raise ValueError(f"k_buf/v_buf must be torch.float32, got {k_buf.dtype}/{v_buf.dtype}") + + cache, (tile_v, vec_size, ilp_rows, use_packed_fma) = _verify_kvbuffer_compile_cache( + B, + T, + H, + HV, + K, + V, + pool_size, + softmax_scale, + write_kv=write_kv, + device=q.device, + ) + + h0_view = s.view(pool_size * HV, V, K) + + # Dummy fp32 tensors when write_kv=False (never accessed by kernel) + if not write_kv: + k_buf_t = torch.empty(1, T, H, K, device=q.device, dtype=torch.float32) + v_buf_t = torch.empty(1, T, HV, V, device=q.device, dtype=torch.float32) + else: + k_buf_t = k_buf + v_buf_t = v_buf + + if "compiled" not in cache: + sym_b = cute.sym_int() + q_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, H, K), stride_order=(3, 2, 1, 0), assumed_align=16) + k_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, H, K), stride_order=(3, 2, 1, 0), assumed_align=16) + v_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, HV, V), stride_order=(3, 2, 1, 0), assumed_align=16) + o_fake = make_fake_compact_tensor(cutlass.Float32, (sym_b, T, HV, V), stride_order=(3, 2, 1, 0), assumed_align=16) + h0_fake = make_fake_compact_tensor(cutlass.Float32, (cute.sym_int(), V, K), stride_order=(2, 1, 0), assumed_align=16) + decay_fake = make_fake_compact_tensor(cutlass.Float32, (H,), stride_order=(0,), assumed_align=16) + idx_fake = make_fake_compact_tensor(cutlass.Int32, (sym_b,), stride_order=(0,), assumed_align=16) + k_buf_fake = make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), T, H, K), stride_order=(3, 2, 1, 0), assumed_align=16 + ) + v_buf_fake = make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), T, HV, V), stride_order=(3, 2, 1, 0), assumed_align=16 + ) + stream_fake = make_fake_stream() + compiled = cute.compile( + run_la_verify_kvbuffer_shuffle_kernel, + h0_fake, + decay_fake, + q_fake, + k_fake, + v_fake, + o_fake, + idx_fake, + k_buf_fake, + v_buf_fake, + Int32(1), + scale=softmax_scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + tile_v=tile_v, + vec_size=vec_size, + ilp_rows=ilp_rows, + use_packed_fma=use_packed_fma, + write_kv=write_kv, + stream=stream_fake, + options="--enable-tvm-ffi", + ) + cache["compiled"] = compiled + + compiled = cache["compiled"] + num_v_tiles_rt = (V + tile_v - 1) // tile_v + grid_size = B * HV * num_v_tiles_rt + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled( + h0_view, + decay_scales, + q, + k, + v, + out, + h0_indices, + k_buf_t, + v_buf_t, + Int32(grid_size), + stream, + ) diff --git a/tests/test_la_decode_mtp.py b/tests/test_la_decode_mtp.py new file mode 100644 index 00000000..c33da98f --- /dev/null +++ b/tests/test_la_decode_mtp.py @@ -0,0 +1,455 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unit tests for la_decode_mtp (CuTe DSL Lightning Attention MTP decode kernel). + +Compares against a PyTorch reference implementation of multi-token +Lightning Attention decode (T > 1). + +Layouts: + q, k: [B, T, H, K] fp32 + v: [B, T, HV, V] fp32 + s: [pool_size, HV, V, K] fp32 (V-major, K-last) + intermediate_states: [pool_size * T * HV, V, K] fp32, or 1-elem dummy + out: [B, T, HV, V] fp32 + decay_scales: [H] fp32 (positive; kernel does exp(-x)) + s_offsets: [B] int32 (pool index per batch; -1 to skip) +""" + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) + +from cula.lightning.la_decode_mtp import linear_attention_decode_mtp + + +# --------------------------------------------------------------------------- +# Pure PyTorch reference for multi-token Lightning Attention decode +# --------------------------------------------------------------------------- +def torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T, cache_intermediate_states=False, disable_state_update=False): + """Pure PyTorch reference. + + Args: + q, k: [B, T, H, D] fp32 + v: [B, T, HV, D] fp32 + state: [B, HV, D, D] fp32 (K-major, V-minor) + decay_scales: [H] fp32 (positive; kernel does exp(-x)) + scale: float + T: int + cache_intermediate_states: cache per-step state to inter + disable_state_update: do not update state_new at end + + Returns: + out: [B, T, HV, D] fp32 + state_new: [B, HV, D, D] fp32 + inter: [B*T*HV, D, D] fp32 or None + """ + B, _, H, D = q.shape + HV = v.shape[2] + q_f = q.float() * scale + k_f, v_f = k.float(), v.float() + decay_per_q_head = torch.exp(-decay_scales) + decay_per_hv = decay_per_q_head.repeat_interleave(HV // H).view(1, HV, 1, 1) + + state_running = state.clone() + out = torch.zeros(B, T, HV, D, dtype=torch.float32, device=q.device) + inter = torch.zeros(B * T * HV, D, D, dtype=torch.float32, device=q.device) if cache_intermediate_states else None + + for t in range(T): + q_hv = q_f[:, t].repeat_interleave(HV // H, dim=1) + k_hv = k_f[:, t].repeat_interleave(HV // H, dim=1) + v_t = v_f[:, t] + state_running = state_running * decay_per_hv + k_hv.unsqueeze(-1) * v_t.unsqueeze(-2) + out[:, t] = torch.einsum("bhk,bhkv->bhv", q_hv, state_running) + if cache_intermediate_states: + for b in range(B): + inter[b * T * HV + t * HV : b * T * HV + (t + 1) * HV] = state_running[b] + + state_final = state.clone() if disable_state_update else state_running + return out, state_final, inter + + +def _skip_if_no_sm90_or_later(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + cc = torch.cuda.get_device_capability("cuda") + if cc[0] < 9: + pytest.skip(f"requires SM90+, got SM{cc[0]}{cc[1]}") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def make_inputs(B, T, H, HV, D, device="cuda", seed=42): + """Returns q[B,T,H,D] fp32, k[B,T,H,D] fp32, v[B,T,HV,D] fp32, state[B,HV,D,D] fp32.""" + torch.manual_seed(seed) + q = torch.randn(B, T, H, D, device=device, dtype=torch.float32) + k = torch.randn(B, T, H, D, device=device, dtype=torch.float32) + v = torch.randn(B, T, HV, D, device=device, dtype=torch.float32) + state = torch.randn(B, HV, D, D, device=device, dtype=torch.float32) * 0.01 + return q, k, v, state + + +def run_la_mtp( + q, + k, + v, + state_4d, + decay_scales, + scale, + T, + cache_intermediate_states=False, + disable_state_update=False, +): + """ + Wraps linear_attention_decode_mtp with proper state-layout conversion. + + state_4d: [B, HV, K, V] fp32 (K-major) + Kernel expects s: [pool_size=B, HV, V, K]; we transpose K and V. + """ + B, HV, K, V = state_4d.shape + H = q.shape[2] + assert HV % H == 0, "HV must be a multiple of H" + + # pretranspose: [B, HV, V, K] + s_cute = state_4d.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, V, device=q.device, dtype=torch.float32) + s_offsets = torch.arange(B, device=q.device, dtype=torch.int32) + + if cache_intermediate_states: + inter = torch.zeros(B * T * HV, V, K, device=q.device, dtype=torch.float32) + else: + inter = torch.empty(1, 1, 1, device=q.device, dtype=torch.float32) # dummy + + cu_seqlens = torch.empty(1, device=q.device, dtype=torch.int32) # dummy when is_varlen=False + + linear_attention_decode_mtp( + q, + k, + v, + s_cute, + inter, + out, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens, + softmax_scale=scale, + T=T, + cache_intermediate_states=cache_intermediate_states, + disable_state_update=disable_state_update, + is_varlen=False, + ) + + # convert state back: [B, HV, V, K] -> [B, HV, K, V] + state_out = s_cute.permute(0, 1, 3, 2).contiguous() + + if cache_intermediate_states: + # inter (kernel): [B*T*HV, V, K] -> ref layout [B*T*HV, K, V] + inter_out = inter.permute(0, 2, 1).contiguous() + else: + inter_out = None + + return out, state_out, inter_out + + +# --------------------------------------------------------------------------- +# Tests vs PyTorch reference +# --------------------------------------------------------------------------- +# Each (B, T) below targets a distinct heuristic config (with H=HV=64): +# B=1, T=4: work_units=64 → tile_v=8, ilp=2, smem_v=False +# B=2, T=2: work_units=128 → tile_v=16, ilp=4, smem_v=False +# B=2, T=4: work_units=128 → tile_v=16, ilp=4, smem_v=False +# B=8, T=4: work_units=512 → tile_v=32, ilp=4, smem_v=False +# B=32, T=2: work_units=2048 → tile_v=64, ilp=8, smem_v=False (state_update ON) +# B=32, T=4: work_units=2048 → tile_v=64, ilp=4, smem_v=True +@pytest.mark.parametrize( + "B,T,expected_config", + [ + (1, 4, "tile_v=8_ilp=2"), + (2, 2, "tile_v=16_ilp=4"), + (2, 4, "tile_v=16_ilp=4"), + (8, 4, "tile_v=32_ilp=4"), + (32, 2, "tile_v=64_ilp=8"), + (32, 4, "tile_v=64_ilp=4_smem_v"), + ], +) +def test_output_vs_torch_ref(B, T, expected_config): + _skip_if_no_sm90_or_later() + H, HV, D = 64, 64, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + o_ref, state_ref, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + o_cute, state_cute, _ = run_la_mtp(q, k, v, state, decay_scales, scale, T) + + # Output check + rmse = torch.sqrt(torch.mean((o_cute.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + rel = rmse / (max_ref + 1e-8) + assert rel < 0.01, f"B={B} T={T} [{expected_config}]: output rel RMSE {rel:.6f} too large" + + # State check + state_rmse = torch.sqrt(torch.mean((state_cute - state_ref) ** 2)).item() + state_max = torch.abs(state_ref).max().item() + state_rel = state_rmse / (state_max + 1e-8) + assert state_rel < 0.001, f"B={B} T={T} [{expected_config}]: state rel RMSE {state_rel:.6f} too large" + + +@pytest.mark.parametrize("H,HV", [(16, 16), (8, 32), (16, 64)]) # MHA + GQA +def test_different_heads(H, HV): + """GQA support: HV is multiple of H; q/k indexed by i_h = i_hv // (HV//H).""" + _skip_if_no_sm90_or_later() + B, T, D = 4, 4, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + o_ref, state_ref, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + o_cute, state_cute, _ = run_la_mtp(q, k, v, state, decay_scales, scale, T) + + rmse = torch.sqrt(torch.mean((o_cute.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + assert rmse / (max_ref + 1e-8) < 0.01, f"H={H} HV={HV}: output mismatch" + + state_rmse = torch.sqrt(torch.mean((state_cute - state_ref) ** 2)).item() + state_max = torch.abs(state_ref).max().item() + assert state_rmse / (state_max + 1e-8) < 0.001, f"H={H} HV={HV}: state mismatch" + + +def test_disable_state_update(): + """h0_source remains bitwise-equal to the input snapshot.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + state_snapshot = state.clone() + + _, state_out, _ = run_la_mtp( + q, + k, + v, + state, + decay_scales, + scale, + T, + disable_state_update=True, + ) + assert torch.equal(state_out, state_snapshot), "state was mutated despite disable_state_update=True" + + +def test_cache_intermediate_states(): + """Each per-t slice of inter matches the reference state_running at that step.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + _, _, inter_ref = torch_la_mtp_ref( + q, + k, + v, + state, + decay_scales, + scale, + T, + cache_intermediate_states=True, + ) + _, _, inter_cute = run_la_mtp( + q, + k, + v, + state, + decay_scales, + scale, + T, + cache_intermediate_states=True, + ) + + rmse = torch.sqrt(torch.mean((inter_cute - inter_ref) ** 2)).item() + max_ref = torch.abs(inter_ref).max().item() + assert rmse / (max_ref + 1e-8) < 0.001, f"intermediate states mismatch, rel_rmse={rmse / (max_ref + 1e-8):.6f}" + + inter_cute_v = inter_cute.view(B, T, HV, D, D) + inter_ref_v = inter_ref.view(B, T, HV, D, D) + for b in range(B): + for t in range(T): + slot_c = inter_cute_v[b, t] + slot_r = inter_ref_v[b, t] + slot_rmse = torch.sqrt(torch.mean((slot_c - slot_r) ** 2)).item() + slot_max = torch.abs(slot_r).max().item() + assert slot_rmse / (slot_max + 1e-8) < 0.001, ( + f"(b={b}, t={t}) intermediate mismatch, rel_rmse={slot_rmse / (slot_max + 1e-8):.6f}" + ) + + assert not torch.allclose(inter_cute_v[0, 0], inter_cute_v[0, 1]) + + +def test_skip_with_negative_offset(): + """s_offsets[i]=-1: that batch's `out` slot stays at initial value.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + sentinel = 123.0 + out = torch.full((B, T, HV, D), sentinel, device=q.device, dtype=torch.float32) + s_offsets = torch.arange(B, device=q.device, dtype=torch.int32) + s_offsets[2] = -1 # skip batch index 2 + + inter = torch.empty(1, 1, 1, device=q.device, dtype=torch.float32) + cu_seqlens = torch.empty(1, device=q.device, dtype=torch.int32) + linear_attention_decode_mtp( + q, + k, + v, + s_cute, + inter, + out, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens, + softmax_scale=scale, + T=T, + cache_intermediate_states=False, + disable_state_update=False, + is_varlen=False, + ) + # batch 2 should be untouched (sentinel value) + assert torch.all(out[2] == torch.full_like(out[2], sentinel)), "skipped batch was modified" + # other batches should differ from sentinel + assert not torch.all(out[0] == torch.full_like(out[0], sentinel)), "non-skipped batch unchanged" + + +def test_skip_with_negative_offset_cache_intermediate(): + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device=q.device, dtype=torch.float32) + s_offsets = torch.arange(B, device=q.device, dtype=torch.int32) + s_offsets[2] = -1 + + inter_sentinel = 7.5 + inter = torch.full((B * T * HV, D, D), inter_sentinel, device=q.device, dtype=torch.float32) + cu_seqlens = torch.empty(1, device=q.device, dtype=torch.int32) + + linear_attention_decode_mtp( + q, + k, + v, + s_cute, + inter, + out, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens, + softmax_scale=scale, + T=T, + cache_intermediate_states=True, + disable_state_update=False, + is_varlen=False, + ) + + skipped = inter[2 * T * HV : 3 * T * HV] + assert torch.all(skipped == inter_sentinel), ( + f"intermediate_states for skipped batch was written (min={skipped.min().item()}, max={skipped.max().item()})" + ) + + others = torch.cat([inter[: 2 * T * HV], inter[3 * T * HV :]], dim=0) + assert not torch.all(others == inter_sentinel), "non-skipped intermediate slots were not written" + + +def test_zero_decay(): + """With decay=0: state_new = state_old + k⊗v (no decay applied).""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = torch.zeros(H, device="cuda", dtype=torch.float32) + + q, k, v, state = make_inputs(B, T, H, HV, D) + o_ref, _, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + o_cute, _, _ = run_la_mtp(q, k, v, state, decay_scales, scale, T) + + rmse = torch.sqrt(torch.mean((o_cute.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + assert rmse / (max_ref + 1e-8) < 0.01, "zero decay: output mismatch" + + +def test_zero_state(): + """With zero initial state.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.ones(H, device="cuda", dtype=torch.float32) + + q, k, v, _ = make_inputs(B, T, H, HV, D) + state = torch.zeros(B, HV, D, D, device="cuda", dtype=torch.float32) + o_ref, _, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + o_cute, _, _ = run_la_mtp(q, k, v, state, decay_scales, scale, T) + + rmse = torch.sqrt(torch.mean((o_cute.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + assert rmse / (max_ref + 1e-8) < 0.01, "zero state: output mismatch" + + +def test_rejects_non_fp32_inputs(): + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 2, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + + q, k, v, state = make_inputs(B, T, H, HV, D) + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + inter = torch.empty(1, 1, 1, device="cuda", dtype=torch.float32) + s_offsets = torch.arange(B, device="cuda", dtype=torch.int32) + cu_seqlens = torch.empty(1, device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="q/k/v must be torch.float32"): + linear_attention_decode_mtp( + q.to(torch.bfloat16), + k, + v, + s_cute, + inter, + out, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens, + softmax_scale=scale, + T=T, + cache_intermediate_states=False, + disable_state_update=False, + is_varlen=False, + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--tb=short"]) diff --git a/tests/test_la_kvbuffer.py b/tests/test_la_kvbuffer.py new file mode 100644 index 00000000..8f23b230 --- /dev/null +++ b/tests/test_la_kvbuffer.py @@ -0,0 +1,749 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the KVBuffer verify + state-update kernels.""" + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) + +from cula.lightning.la_decode_mtp import linear_attention_decode_mtp +from cula.lightning.la_state_update_kvbuffer import ( + linear_attention_state_update_kvbuffer, + linear_attention_state_update_kvbuffer_fused, +) +from cula.lightning.la_verify_kvbuffer import linear_attention_verify_kvbuffer + + +# --------------------------------------------------------------------------- +# Pure PyTorch reference for multi-token Lightning Attention decode +# --------------------------------------------------------------------------- +def torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T, cache_intermediate_states=False, disable_state_update=False): + """Pure PyTorch reference. + + Args: + q, k: [B, T, H, D] fp32 + v: [B, T, HV, D] fp32 + state: [B, HV, D, D] fp32 (K-major, V-minor) + decay_scales: [H] fp32 (positive; kernel does exp(-x)) + scale: float + T: int + cache_intermediate_states: cache per-step state to inter + disable_state_update: do not update state_new at end + + Returns: + out: [B, T, HV, D] fp32 + state_new: [B, HV, D, D] fp32 + inter: [B*T*HV, D, D] fp32 or None + """ + B, _, H, D = q.shape + HV = v.shape[2] + q_f = q.float() * scale + k_f, v_f = k.float(), v.float() + decay_per_q_head = torch.exp(-decay_scales) + decay_per_hv = decay_per_q_head.repeat_interleave(HV // H).view(1, HV, 1, 1) + + state_running = state.clone() + out = torch.zeros(B, T, HV, D, dtype=torch.float32, device=q.device) + inter = torch.zeros(B * T * HV, D, D, dtype=torch.float32, device=q.device) if cache_intermediate_states else None + + for t in range(T): + q_hv = q_f[:, t].repeat_interleave(HV // H, dim=1) + k_hv = k_f[:, t].repeat_interleave(HV // H, dim=1) + v_t = v_f[:, t] + state_running = state_running * decay_per_hv + k_hv.unsqueeze(-1) * v_t.unsqueeze(-2) + out[:, t] = torch.einsum("bhk,bhkv->bhv", q_hv, state_running) + if cache_intermediate_states: + for b in range(B): + inter[b * T * HV + t * HV : b * T * HV + (t + 1) * HV] = state_running[b] + + state_final = state.clone() if disable_state_update else state_running + return out, state_final, inter + + +def _skip_if_no_sm90_or_later(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + cc = torch.cuda.get_device_capability("cuda") + if cc[0] < 9: + pytest.skip(f"requires SM90+, got SM{cc[0]}{cc[1]}") + + +def _make_inputs(B, T, H, HV, D, device="cuda", seed=42): + torch.manual_seed(seed) + q = torch.randn(B, T, H, D, device=device, dtype=torch.float32) + k = torch.randn(B, T, H, D, device=device, dtype=torch.float32) + v = torch.randn(B, T, HV, D, device=device, dtype=torch.float32) + state = torch.randn(B, HV, D, D, device=device, dtype=torch.float32) * 0.01 + return q, k, v, state + + +def _make_kv_buffers(k, v, h0_indices, pool_size=None): + B, T, H, D = k.shape + _, _, HV, V = v.shape + if pool_size is None: + pool_size = B + k_buf = torch.zeros(pool_size, T, H, D, device=k.device, dtype=torch.float32) + v_buf = torch.zeros(pool_size, T, HV, V, device=v.device, dtype=torch.float32) + for b in range(B): + pool_idx = int(h0_indices[b].item()) + if pool_idx >= 0: + k_buf[pool_idx] = k[b] + v_buf[pool_idx] = v[b] + return k_buf, v_buf + + +@pytest.mark.parametrize("T", [4, 16]) +def test_verify_rejects_non_fp32_state(T): + _skip_if_no_sm90_or_later() + B, H, HV, D = 2, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + s_cute = state.permute(0, 1, 3, 2).contiguous().clone().to(torch.bfloat16) + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="s must be torch.float32"): + linear_attention_verify_kvbuffer(q, k, v, s_cute, out, decay_scales, h0_indices, scale, T) + + +def test_state_update_L0_no_op(): + """accepted_len=0 everywhere: s must be byte-for-byte unchanged.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + _, k, v, state = _make_inputs(B, T, H, HV, D) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() # [B, HV, V, K] + s_snapshot = s_cute.clone() + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + accepted_len = torch.zeros(B, device="cuda", dtype=torch.int32) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_cute, + decay_scales, + h0_indices, + accepted_len, + T, + ) + assert torch.equal(s_cute, s_snapshot), "L=0 must leave state unchanged" + + +def _ref_state_after_L(state, k, v, decay_scales, L_per_batch, T): + """state[B,HV,K,V] fp32; returns the per-batch state after L recurrent steps.""" + B, HV, K, V = state.shape + H = k.shape[2] + k_f, v_f = k.float(), v.float() + decay_per_q_head = torch.exp(-decay_scales) + decay_per_hv = decay_per_q_head.repeat_interleave(HV // H).view(HV, 1, 1) + out = state.clone() + for b in range(B): + L = int(L_per_batch[b].item()) + running = state[b].clone() + for i in range(L): + k_hv = k_f[b, i].repeat_interleave(HV // H, dim=0) # [HV, K] + v_i = v_f[b, i] # [HV, V] + running = running * decay_per_hv + k_hv.unsqueeze(-1) * v_i.unsqueeze(-2) + out[b] = running + return out + + +@pytest.mark.parametrize( + "B,T,H,HV,D", + [(4, 4, 16, 16, 128), (8, 4, 64, 64, 128), (4, 3, 16, 16, 128), (8, 7, 64, 64, 128)], +) +def test_state_update_full_accept(B, T, H, HV, D): + """accepted_len=T everywhere: bit-exact vs baseline recurrence reference.""" + _skip_if_no_sm90_or_later() + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + _, k, v, state = _make_inputs(B, T, H, HV, D) + + L_per_batch = torch.full((B,), T, device="cuda", dtype=torch.int32) + ref = _ref_state_after_L(state, k, v, decay_scales, L_per_batch, T) # [B,HV,K,V] + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() # [B,HV,V,K] + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_cute, + decay_scales, + h0_indices, + L_per_batch, + T, + ) + got = s_cute.permute(0, 1, 3, 2).contiguous() # back to [B,HV,K,V] + rmse = torch.sqrt(torch.mean((got - ref) ** 2)).item() + rel = rmse / (torch.abs(ref).max().item() + 1e-8) + assert rel < 1e-3, f"full-accept state rel RMSE {rel:.6f} too large" + + +@pytest.mark.parametrize("L", [0, 1, 3]) +def test_state_update_partial(L): + """Uniform accepted_len=L across all batches.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + _, k, v, state = _make_inputs(B, T, H, HV, D) + + L_per_batch = torch.full((B,), L, device="cuda", dtype=torch.int32) + ref = _ref_state_after_L(state, k, v, decay_scales, L_per_batch, T) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_cute, + decay_scales, + h0_indices, + L_per_batch, + T, + ) + got = s_cute.permute(0, 1, 3, 2).contiguous() + rel = torch.sqrt(torch.mean((got - ref) ** 2)).item() / (torch.abs(ref).max().item() + 1e-8) + assert rel < 1e-3, f"L={L} state rel RMSE {rel:.6f}" + + +def test_state_update_per_batch_L(): + """accepted_len varies per batch: [0, 1, T-1, T].""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + _, k, v, state = _make_inputs(B, T, H, HV, D) + + L_per_batch = torch.tensor([0, 1, T - 1, T], device="cuda", dtype=torch.int32) + ref = _ref_state_after_L(state, k, v, decay_scales, L_per_batch, T) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_cute, + decay_scales, + h0_indices, + L_per_batch, + T, + ) + got = s_cute.permute(0, 1, 3, 2).contiguous() + for b in range(B): + rel = torch.sqrt(torch.mean((got[b] - ref[b]) ** 2)).item() / (torch.abs(ref[b]).max().item() + 1e-8) + assert rel < 1e-3, f"batch {b} (L={int(L_per_batch[b])}) rel RMSE {rel:.6f}" + + +def test_state_update_skip_negative_h0_indices(): + """h0_indices[b]=-1: that pool slot is untouched even with accepted_len>0.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + _, k, v, state = _make_inputs(B, T, H, HV, D) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + snapshot_b2 = s_cute[2].clone() + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + h0_indices[2] = -1 + L_per_batch = torch.full((B,), T, device="cuda", dtype=torch.int32) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_cute, + decay_scales, + h0_indices, + L_per_batch, + T, + ) + assert torch.equal(s_cute[2], snapshot_b2), "skipped batch slot was modified" + + +def test_verify_skip_negative_h0_indices(): + """h0_indices[b]=-1: out[b] stays at its sentinel value.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + sentinel = 123.0 + out = torch.full((B, T, HV, D), sentinel, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + h0_indices[2] = -1 + + linear_attention_verify_kvbuffer( + q, + k, + v, + s_cute, + out, + decay_scales, + h0_indices, + scale, + T, + ) + assert torch.all(out[2] == sentinel), "skipped batch out slot was modified" + + +@pytest.mark.parametrize( + "B,T", + [(1, 4), (2, 2), (2, 4), (8, 4), (32, 2), (32, 4), (2, 1), (2, 3), (8, 5), (8, 7)], +) +def test_verify_outputs_match_ref(B, T): + """Verify kernel o matches torch_la_mtp_ref across the baseline configs.""" + _skip_if_no_sm90_or_later() + H, HV, D = 64, 64, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + o_ref, _, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + linear_attention_verify_kvbuffer( + q, + k, + v, + s_cute, + out, + decay_scales, + h0_indices, + scale, + T, + ) + rel = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() / (torch.abs(o_ref.float()).max().item() + 1e-8) + assert rel < 1e-2, f"B={B} T={T}: verify output rel RMSE {rel:.6f} too large" + + +@pytest.mark.parametrize("H,HV", [(16, 16), (8, 32), (16, 64)]) +def test_verify_different_heads(H, HV): + _skip_if_no_sm90_or_later() + B, T, D = 4, 4, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + o_ref, _, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + linear_attention_verify_kvbuffer( + q, + k, + v, + s_cute, + out, + decay_scales, + h0_indices, + scale, + T, + ) + rel = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() / (torch.abs(o_ref.float()).max().item() + 1e-8) + assert rel < 1e-2, f"H={H} HV={HV}: verify output mismatch {rel:.6f}" + + +def test_verify_zero_decay(): + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = torch.zeros(H, device="cuda", dtype=torch.float32) + q, k, v, state = _make_inputs(B, T, H, HV, D) + o_ref, _, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + linear_attention_verify_kvbuffer(q, k, v, s_cute, out, decay_scales, h0_indices, scale, T) + rel = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() / (torch.abs(o_ref.float()).max().item() + 1e-8) + assert rel < 1e-2, f"zero decay: {rel:.6f}" + + +def test_verify_zero_state(): + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.ones(H, device="cuda", dtype=torch.float32) + q, k, v, _ = _make_inputs(B, T, H, HV, D) + state = torch.zeros(B, HV, D, D, device="cuda", dtype=torch.float32) + o_ref, _, _ = torch_la_mtp_ref(q, k, v, state, decay_scales, scale, T) + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + linear_attention_verify_kvbuffer(q, k, v, s_cute, out, decay_scales, h0_indices, scale, T) + rel = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() / (torch.abs(o_ref.float()).max().item() + 1e-8) + assert rel < 1e-2, f"zero state: {rel:.6f}" + + +def test_end_to_end_equivalence_with_baseline(): + """KVBuffer (verify + state_update L=T) == baseline (cache_inter=T, disable=T).""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 8, 4, 64, 64, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + # ---- Baseline: capture out + all intermediate states ---- + s_base = state.permute(0, 1, 3, 2).contiguous().clone() # [B,HV,V,K] + out_base = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + s_offsets = torch.arange(B, device="cuda", dtype=torch.int32) + inter = torch.zeros(B * T * HV, D, D, device="cuda", dtype=torch.float32) # [.,V,K] + cu_seqlens = torch.empty(1, device="cuda", dtype=torch.int32) + linear_attention_decode_mtp( + q, + k, + v, + s_base, + inter, + out_base, + decay_scales=decay_scales, + s_offsets=s_offsets, + cu_seqlens=cu_seqlens, + softmax_scale=scale, + T=T, + cache_intermediate_states=True, + disable_state_update=True, + is_varlen=False, + ) + + # ---- KVBuffer: verify writes out; state-update (L=T) writes state ---- + s_kv = state.permute(0, 1, 3, 2).contiguous().clone() # [B,HV,V,K] + out_kv = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + linear_attention_verify_kvbuffer( + q, + k, + v, + s_kv, + out_kv, + decay_scales, + h0_indices, + scale, + T, + ) + accepted_len = torch.full((B,), T, device="cuda", dtype=torch.int32) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_kv, + decay_scales, + h0_indices, + accepted_len, + T, + ) + + # (a) outputs match + rel_o = torch.sqrt(torch.mean((out_kv - out_base) ** 2)).item() / (torch.abs(out_base).max().item() + 1e-8) + assert rel_o < 1e-2, f"output mismatch vs baseline: {rel_o:.6f}" + + # (b) updated state == baseline's last intermediate slice [B,HV,V,K] + inter_v = inter.view(B, T, HV, D, D) # [B,T,HV,V,K] + last_state = inter_v[:, T - 1] # [B,HV,V,K] + rel_s = torch.sqrt(torch.mean((s_kv - last_state) ** 2)).item() / (torch.abs(last_state).max().item() + 1e-8) + assert rel_s < 1e-3, f"state mismatch vs baseline last intermediate: {rel_s:.6f}" + + +@pytest.mark.parametrize("B,T", [(4, 4), (8, 2), (32, 4)]) +def test_verify_writes_kv_buffer(B, T): + """Verify kernel with k_buf/v_buf writes correct copies of k and v.""" + _skip_if_no_sm90_or_later() + H, HV, D = 64, 64, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + pool_size = B + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + k_buf = torch.zeros(pool_size, T, H, D, device="cuda", dtype=torch.float32) + v_buf = torch.zeros(pool_size, T, HV, D, device="cuda", dtype=torch.float32) + + linear_attention_verify_kvbuffer( + q, + k, + v, + s_cute, + out, + decay_scales, + h0_indices, + scale, + T, + k_buf=k_buf, + v_buf=v_buf, + ) + + for b in range(B): + pool_idx = h0_indices[b].item() + assert torch.equal(k_buf[pool_idx], k[b]), f"k_buf mismatch at batch {b}" + assert torch.equal(v_buf[pool_idx], v[b]), f"v_buf mismatch at batch {b}" + + +def test_verify_output_unchanged_with_kv_write(): + """Output o is identical whether k_buf/v_buf are provided or not.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 8, 4, 64, 64, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + pool_size = B + s1 = state.permute(0, 1, 3, 2).contiguous().clone() + s2 = s1.clone() + out_no_buf = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + out_with_buf = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + + linear_attention_verify_kvbuffer( + q, + k, + v, + s1, + out_no_buf, + decay_scales, + h0_indices, + scale, + T, + ) + + k_buf = torch.zeros(pool_size, T, H, D, device="cuda", dtype=torch.float32) + v_buf = torch.zeros(pool_size, T, HV, D, device="cuda", dtype=torch.float32) + linear_attention_verify_kvbuffer( + q, + k, + v, + s2, + out_with_buf, + decay_scales, + h0_indices, + scale, + T, + k_buf=k_buf, + v_buf=v_buf, + ) + + assert torch.equal(out_no_buf, out_with_buf), "kv write should not affect output" + + +@pytest.mark.parametrize("B,T,H,HV,D", [(4, 4, 16, 16, 128), (8, 4, 64, 64, 128)]) +def test_state_update_from_buffer(B, T, H, HV, D): + """State update from pool-indexed k_buf/v_buf is deterministic.""" + _skip_if_no_sm90_or_later() + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + _, k, v, state = _make_inputs(B, T, H, HV, D) + + pool_size = B + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + L_per_batch = torch.full((B,), T, device="cuda", dtype=torch.int32) + + # Path A: read from pool-indexed k_buf/v_buf filled from raw k, v + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices) + s_raw = state.permute(0, 1, 3, 2).contiguous().clone() + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_raw, + decay_scales, + h0_indices, + L_per_batch, + T, + ) + + # Path B: same data, independently materialized buffer. + k_buf_2, v_buf_2 = _make_kv_buffers(k, v, h0_indices, pool_size=pool_size) + s_buf = state.permute(0, 1, 3, 2).contiguous().clone() + linear_attention_state_update_kvbuffer( + k_buf_2, + v_buf_2, + s_buf, + decay_scales, + h0_indices, + L_per_batch, + T, + ) + + assert torch.equal(s_raw, s_buf), "buffer-read state must match raw-read state" + + +def test_verify_skip_negative_indices_no_buffer_write(): + """h0_indices[b]=-1: k_buf and v_buf slots are untouched.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 4, 4, 16, 16, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + pool_size = B + sentinel = 42.0 + k_buf = torch.full((pool_size, T, H, D), sentinel, device="cuda", dtype=torch.float32) + v_buf = torch.full((pool_size, T, HV, D), sentinel, device="cuda", dtype=torch.float32) + k_buf_snap = k_buf.clone() + v_buf_snap = v_buf.clone() + + s_cute = state.permute(0, 1, 3, 2).contiguous().clone() + out = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + h0_indices[2] = -1 + + linear_attention_verify_kvbuffer( + q, + k, + v, + s_cute, + out, + decay_scales, + h0_indices, + scale, + T, + k_buf=k_buf, + v_buf=v_buf, + ) + + assert torch.equal(k_buf[2], k_buf_snap[2]), "skipped batch k_buf slot was modified" + assert torch.equal(v_buf[2], v_buf_snap[2]), "skipped batch v_buf slot was modified" + + +def test_end_to_end_with_buffer(): + """Full pipeline: verify(+kv write) → state_update(from buffer) matches baseline.""" + _skip_if_no_sm90_or_later() + B, T, H, HV, D = 8, 4, 64, 64, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(H, device="cuda", dtype=torch.float32) / H + q, k, v, state = _make_inputs(B, T, H, HV, D) + + pool_size = B + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + + # Reference: existing end-to-end (no buffer) + s_ref = state.permute(0, 1, 3, 2).contiguous().clone() + out_ref = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + linear_attention_verify_kvbuffer( + q, + k, + v, + s_ref, + out_ref, + decay_scales, + h0_indices, + scale, + T, + ) + accepted_len = torch.full((B,), T, device="cuda", dtype=torch.int32) + k_buf_ref, v_buf_ref = _make_kv_buffers(k, v, h0_indices, pool_size=pool_size) + linear_attention_state_update_kvbuffer( + k_buf_ref, + v_buf_ref, + s_ref, + decay_scales, + h0_indices, + accepted_len, + T, + ) + + # Buffer path: verify writes buffer, state_update reads buffer + s_buf = state.permute(0, 1, 3, 2).contiguous().clone() + out_buf = torch.zeros(B, T, HV, D, device="cuda", dtype=torch.float32) + k_buf = torch.zeros(pool_size, T, H, D, device="cuda", dtype=torch.float32) + v_buf = torch.zeros(pool_size, T, HV, D, device="cuda", dtype=torch.float32) + + linear_attention_verify_kvbuffer( + q, + k, + v, + s_buf, + out_buf, + decay_scales, + h0_indices, + scale, + T, + k_buf=k_buf, + v_buf=v_buf, + ) + linear_attention_state_update_kvbuffer( + k_buf, + v_buf, + s_buf, + decay_scales, + h0_indices, + accepted_len, + T, + ) + + assert torch.equal(out_ref, out_buf), "output mismatch with buffer pipeline" + assert torch.equal(s_ref, s_buf), "state mismatch with buffer pipeline" + + +def test_state_update_fused_matches_per_layer(): + """Layer-fused state update matches independent per-layer launches.""" + _skip_if_no_sm90_or_later() + num_layers, B, T, H, HV, D = 3, 4, 4, 16, 16, 128 + pool_size = B + h0_indices = torch.arange(B, device="cuda", dtype=torch.int32) + accepted_len = torch.tensor([0, 1, T - 1, T], device="cuda", dtype=torch.int32) + + k_buf_layers = [] + v_buf_layers = [] + states = [] + decays = [] + for layer in range(num_layers): + _, k, v, state = _make_inputs(B, T, H, HV, D, seed=100 + layer) + k_buf, v_buf = _make_kv_buffers(k, v, h0_indices, pool_size=pool_size) + k_buf_layers.append(k_buf) + v_buf_layers.append(v_buf) + states.append(state.permute(0, 1, 3, 2).contiguous()) + decays.append(0.3 * (layer + 1) * torch.arange(H, device="cuda", dtype=torch.float32) / H) + + k_buf_fused = torch.stack(k_buf_layers, dim=0) + v_buf_fused = torch.stack(v_buf_layers, dim=0) + s_fused = torch.stack(states, dim=0) + decay_fused = torch.stack(decays, dim=0) + + s_expected = s_fused.clone() + for layer in range(num_layers): + linear_attention_state_update_kvbuffer( + k_buf_fused[layer], + v_buf_fused[layer], + s_expected[layer], + decay_fused[layer], + h0_indices, + accepted_len, + T, + ) + + linear_attention_state_update_kvbuffer_fused( + k_buf_fused, + v_buf_fused, + s_fused, + decay_fused, + h0_indices, + accepted_len, + T, + ) + + assert torch.equal(s_fused, s_expected), "fused state update must match per-layer state update" From abfd99f574c941b2395de12bd56601971a55aa4d Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 9 Jul 2026 12:28:19 +0800 Subject: [PATCH 31/34] feat: add packed kda decode (#102) * feat: add packed kda decode * fix: address packed decode review comments * fix: add missing kda lazy exports --------- Co-authored-by: zhouaihui --- benchmarks/bench_kda_packed_decode.py | 264 +++++++++ cula/kda/__init__.py | 4 + cula/ops/kda/decode/cute.py | 560 ++++++++++++++++++ tests/test_kda_packed_decode.py | 807 ++++++++++++++++++++++++++ 4 files changed, 1635 insertions(+) create mode 100644 benchmarks/bench_kda_packed_decode.py create mode 100644 tests/test_kda_packed_decode.py diff --git a/benchmarks/bench_kda_packed_decode.py b/benchmarks/bench_kda_packed_decode.py new file mode 100644 index 00000000..d9e13e7d --- /dev/null +++ b/benchmarks/bench_kda_packed_decode.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""bench_kda_packed_decode.py — micro-benchmark: packed vs non-packed KDA decode. + +Compares single-token (T=1) decode routes that share the same CuTe DSL kernel +body, so the delta isolates host-side orchestration + the q/k/v materialization +that a real caller pays. + +Two routes are timed, BOTH seeded from the same packed ``mixed_qkv`` (the form +a conv layer actually produces, e.g. in SGLang's KDA decode): + + 1. non-packed (realistic): the caller must turn ``mixed_qkv`` back into + separate contiguous q/k/v before calling ``kda_decode`` — i.e. a per-call + ``split + unsqueeze + contiguous``. The ``.contiguous()`` is a REAL copy + when ``N>1`` because the split leaves row stride = qkv_dim. This is the + cost the packed path is meant to remove. + 2. packed: ``cula.kda.kda_packed_decode`` feeds q/k/v as strided views + directly — no materialization, no ``.contiguous()`` copy. + +A third "non-packed (pre-split)" column is included as an oracle: it uses +q/k/v that were split once outside the timed loop (no per-call copy). This +shows the floor of the non-packed approach — i.e. how fast decode could be if +the caller already had separate contiguous q/k/v. packed matching or beating +the realistic column while staying near the oracle is the win. + +All routes are timed with CUDA events bracketing the full callable (host-side +view construction + cache lookup + kernel launch). The ``mixed_qkv`` is reused +across iterations; only the per-call split+contiguous is inside the timed +region for the realistic non-packed route. + +Usage: + python benchmarks/bench_kda_packed_decode.py + python benchmarks/bench_kda_packed_decode.py --batch-sizes 1 4 16 64 128 + python benchmarks/bench_kda_packed_decode.py --head-pairs 8:8 16:16 32:32 64:64 + python benchmarks/bench_kda_packed_decode.py --ncu +""" + +import argparse +import pathlib +import sys + +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from benchmarks.utils import benchmark_cuda_fn +from cula.kda import fused_sigmoid_gating_delta_rule_update as cula_fused +from cula.kda import kda_packed_decode + + +def make_inputs(N, H, HV, K, V, device="cuda", seed=42): + torch.manual_seed(seed) + q = torch.randn(N, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn(N, H, K, device=device, dtype=torch.bfloat16) + v = torch.randn(N, HV, V, device=device, dtype=torch.bfloat16) + a = (torch.randn(N, HV, K, device=device, dtype=torch.float32) * 0.1).to(torch.bfloat16) + b = torch.randn(N, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) * 2 + dt_bias = torch.randn(HV, K, device=device, dtype=torch.float32) * 0.1 + state = torch.randn(N, HV, V, K, device=device, dtype=torch.float32) * 0.01 + return q, k, v, a, b, A_log, dt_bias, state + + +def run_config(N, H, HV, K, V, warmup, rep): + device = "cuda" + scale = K**-0.5 + + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V, device) + + q_4d = q.unsqueeze(1).contiguous() + k_4d = k.unsqueeze(1).contiguous() + v_4d = v.unsqueeze(1).contiguous() + a_flat = a.reshape(N, 1, -1).contiguous() + b_3d = b.unsqueeze(1).contiguous() + mixed_qkv = torch.cat([q.view(N, -1), k.view(N, -1), v.view(N, -1)], dim=-1).contiguous() + qk_dim = H * K + v_dim = HV * V + indices = torch.arange(N, device=device, dtype=torch.int32) + + state_init = state.clone().contiguous() + + # non-packed (oracle): separate contiguous q/k/v split once outside the loop. + def call_oracle(state_buf): + return cula_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="vk", + ) + + # non-packed (realistic): start from mixed_qkv, split+unsqueeze+contiguous + # every call — the cost a real SGLang caller pays when q/k/v are not kept + # pre-split. + def call_nonpacked(state_buf): + qq, kk, vv = mixed_qkv.split([qk_dim, qk_dim, v_dim], dim=-1) + return cula_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=qq.view(N, 1, H, K).contiguous(), + k=kk.view(N, 1, H, K).contiguous(), + v=vv.view(N, 1, HV, V).contiguous(), + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="vk", + ) + + # packed: kda_packed_decode route (mixed_qkv reused every iteration) + def call_packed(state_buf): + return kda_packed_decode( + mixed_qkv, + a_flat, + b_3d, + A_log=A_log, + dt_bias=dt_bias, + state=state_buf, + state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + state_layout="vk", + ) + + # Correctness sanity: packed must match the oracle. + state_ora = state_init.clone() + state_pck = state_init.clone() + with torch.no_grad(): + o_ora = call_oracle(state_ora) + o_pck = call_packed(state_pck) + out_diff = (o_ora.float() - o_pck.float()).abs().max().item() + state_diff = (state_ora.float() - state_pck.float()).abs().max().item() + + state_bench_ora = state_init.clone() + state_bench_non = state_init.clone() + state_bench_pck = state_init.clone() + + def setup_ora(): + state_bench_ora.copy_(state_init) + + def setup_non(): + state_bench_non.copy_(state_init) + + def setup_pck(): + state_bench_pck.copy_(state_init) + + with torch.no_grad(): + t_ora = benchmark_cuda_fn(lambda: call_oracle(state_bench_ora), setup_fn=setup_ora, warmup=warmup, rep=rep) + t_non = benchmark_cuda_fn(lambda: call_nonpacked(state_bench_non), setup_fn=setup_non, warmup=warmup, rep=rep) + t_pck = benchmark_cuda_fn(lambda: call_packed(state_bench_pck), setup_fn=setup_pck, warmup=warmup, rep=rep) + + # q/k/v bytes the non-packed path copies per call (split+contiguous): bf16. + qkv_dim = mixed_qkv.shape[1] + copy_bytes = 2 * (2 * H * K + HV * V) * N + + return { + "N": N, + "H": H, + "HV": HV, + "K": K, + "V": V, + "qkv_dim": qkv_dim, + "t_oracle_ms": t_ora, + "t_non_ms": t_non, + "t_packed_ms": t_pck, + "saved_vs_non_us": (t_non - t_pck) * 1e3, + "saved_vs_ora_us": (t_ora - t_pck) * 1e3, + "speedup_vs_non": t_non / t_pck if t_pck > 0 else float("inf"), + "out_diff": out_diff, + "state_diff": state_diff, + "copy_bytes": copy_bytes, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 4, 16, 64, 128]) + parser.add_argument("--Hs", type=int, nargs="+", default=[8, 16]) + parser.add_argument("--HV", type=int, default=16) + parser.add_argument( + "--head-pairs", + type=str, + nargs="+", + default=None, + help="Explicit H:HV pairs, e.g. --head-pairs 8:8 16:16 32:32 64:64", + ) + parser.add_argument("--K", type=int, default=128) + parser.add_argument("--V", type=int, default=128) + parser.add_argument("--warmup", type=int, default=30) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--ncu", action="store_true") + args = parser.parse_args() + + if args.ncu: + args.warmup, args.rep = 1, 1 + + gpu = torch.cuda.get_device_name(0) + print(f"# cuLA Packed KDA Decode micro-bench [{gpu}]") + print(f"# K={args.K} V={args.V} warmup={args.warmup} rep={args.rep}") + print() + + if args.head_pairs is None: + head_pairs = [(H, args.HV if args.HV >= 2 * H else 2 * H) for H in args.Hs] + else: + head_pairs = [] + for item in args.head_pairs: + try: + h_str, hv_str = item.split(":", 1) + H, HV = int(h_str), int(hv_str) + except ValueError as exc: + raise ValueError(f"Invalid --head-pairs entry {item!r}; expected H:HV, e.g. 8:8") from exc + if H <= 0 or HV <= 0 or HV % H != 0: + raise ValueError(f"Invalid head pair H={H}, HV={HV}; expected positive values with HV % H == 0") + head_pairs.append((H, HV)) + + for H, HV in head_pairs: + print(f"## H={H} HV={HV} K={args.K} V={args.V}") + hdr = ( + f"{'N':>5} | {'qkv_dim':>7} | {'oracle (ms)':>12} | {'non-packed (ms)':>16} " + f"| {'cula_packed (ms)':>17} | {'save vs non (us)':>16} | {'speedup vs non':>15} " + f"| {'out_diff':>9} | {'state_diff':>10}" + ) + print(hdr) + print("-" * len(hdr)) + for N in args.batch_sizes: + r = run_config(N, H, HV, args.K, args.V, args.warmup, args.rep) + print( + f"{r['N']:>5} | {r['qkv_dim']:>7} | {r['t_oracle_ms']:>12.4f} | {r['t_non_ms']:>16.4f} " + f"| {r['t_packed_ms']:>17.4f} | {r['saved_vs_non_us']:>16.2f} | {r['speedup_vs_non']:>14.3f}x " + f"| {r['out_diff']:>9.2e} | {r['state_diff']:>10.2e}" + ) + print() + + +if __name__ == "__main__": + main() diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index 0b61b13e..73f33219 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -20,6 +20,7 @@ "kda_decode_mtp", "kda_decode_mtp_recurrent", "kda_decode_mtp_recurrent_ws", + "kda_packed_decode", "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", "kda_prefill_hopper_opt", @@ -29,6 +30,8 @@ _LAZY = { "chunk_kda": ("cula.kda.chunk", "chunk_kda"), "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"), + "kda_prefill_hopper_opt": ("cula.kda.hopper_fused_fwd_opt", "cula_kda_prefill_opt"), + "kda_prefill_hopper_auto": ("cula.kda.auto_route", "cula_kda_prefill_auto"), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), @@ -36,6 +39,7 @@ "cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent_ws", ), + "kda_packed_decode": ("cula.ops.kda.decode.cute", "kda_packed_decode"), "fused_sigmoid_gating_delta_rule_update": ( "cula.ops.kda.decode.cute", "fused_sigmoid_gating_delta_rule_update", diff --git a/cula/ops/kda/decode/cute.py b/cula/ops/kda/decode/cute.py index d84c77bf..37c6025d 100644 --- a/cula/ops/kda/decode/cute.py +++ b/cula/ops/kda/decode/cute.py @@ -285,6 +285,199 @@ def _try_fast_dense_decode( return o +def _try_fast_dense_packed_decode( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + N: int, + H: int, + HV: int, + K: int, + V: int, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + is_varlen_decode: bool, + cu_seqlens: torch.Tensor | None, + scale: float | None, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float, + softplus_threshold: float, + out: torch.Tensor | None, + state_layout: str | None, +): + """Fast path for packed-QKV decode. + + Mirrors ``_try_fast_dense_decode`` but takes a packed ``mixed_qkv`` of shape + ``[N, qkv_dim]`` (``= [Q(H·K) | K(H·K) | V(HV·V)]``, head-major, last dim + contiguous) and constructs the q/k/v as strided views instead of requiring + three separate contiguous tensors — so no q/k/v materialization or + ``.contiguous()`` copy is needed. + + Returns ``None`` (falling back to the general path) when the inputs are not + already in the exact kernel-ready layout/dtype; never raises on shape/dtype + mismatches. + """ + if K != TILE_K or V % TILE_V_SMALL != 0 or V % TILE_V != 0: + return None + + qkv_dim = 2 * H * K + HV * V + if ( + mixed_qkv.ndim != 2 + or mixed_qkv.shape != (N, qkv_dim) + or mixed_qkv.stride(-1) != 1 + or mixed_qkv.stride(0) < qkv_dim + or mixed_qkv.device.type != "cuda" + or mixed_qkv.dtype != torch.bfloat16 + ): + return None + + if ( + initial_state_source.dtype != torch.float32 + or initial_state_indices.dtype != torch.int32 + or initial_state_indices.shape != (N,) + or A_log.dtype != torch.float32 + or dt_bias.dtype != torch.float32 + ): + return None + + if not ( + initial_state_source.is_contiguous() + and initial_state_indices.is_contiguous() + and A_log.is_contiguous() + and dt_bias.is_contiguous() + ): + return None + if A_log.numel() != HV or dt_bias.shape != (HV, K): + return None + + normalized_layout = "vk" if state_layout is None else str(state_layout).strip().lower() + if normalized_layout == "vk": + if initial_state_source.ndim != 4 or initial_state_source.shape[1:] != (HV, V, K): + return None + state_layout_is_kv = False + elif normalized_layout == "kv": + if initial_state_source.ndim != 4 or initial_state_source.shape[1:] != (HV, K, V): + return None + state_layout_is_kv = True + else: + return None + + if not a.is_contiguous() or a.device != mixed_qkv.device or a.dtype != torch.bfloat16: + return None + if is_varlen_decode: + # varlen kernel compiled a shape: (N, HV, K) -- 3D + if a.dim() == 3 and a.shape == (N, HV, K): + a_kernel = a + else: + return None + else: + # dense kernel compiled a shape: (N, 1, HV, K) -- 4D + if a.dim() == 4 and a.shape == (N, 1, HV, K): + a_kernel = a + elif a.dim() == 3 and a.shape == (N, 1, HV * K): + a_kernel = a.view(N, 1, HV, K) + elif a.dim() == 3 and a.shape == (N, HV, K): + a_kernel = a.unsqueeze(1) + else: + return None + + if b.device != mixed_qkv.device or b.dtype != torch.bfloat16 or not b.is_contiguous(): + return None + if is_varlen_decode: + # varlen b compiled: (N, HV) -- 2D + if b.dim() == 2 and b.shape == (N, HV): + b_kernel = b + else: + return None + else: + # dense b compiled: (N, 1, HV) -- 3D + if b.dim() == 3 and b.shape == (N, 1, HV): + b_kernel = b + elif b.dim() == 2 and b.shape == (N, HV): + b_kernel = b.unsqueeze(1) + else: + return None + + if scale is None: + scale = K**-0.5 + elif scale <= 0: + return None + + # Construct strided q/k/v views (row stride may include padding, last dim contiguous) + # and the matching output. Shapes align with the compile-time mock in the + # packed compile paths: varlen -> [1,N,...], dense -> [N,1,...]. + if is_varlen_decode: + q_view = mixed_qkv.narrow(1, 0, H * K).view(1, N, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(1, N, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(1, N, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (1, N, HV, V)) + else: + q_view = mixed_qkv.narrow(1, 0, H * K).view(N, 1, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(N, 1, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(N, 1, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (N, 1, HV, V)) + + if cu_seqlens is not None: + if cu_seqlens.dtype != torch.int32 or cu_seqlens.numel() != N + 1 or not cu_seqlens.is_contiguous(): + return None + cu_seqlens_to_use = cu_seqlens + else: + cache_key = (N, str(mixed_qkv.device)) + if cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cache_key] = torch.arange(N + 1, dtype=torch.int32, device=mixed_qkv.device) + cu_seqlens_to_use = _cu_seqlens_cache[cache_key] + + use_small_batch = N < SMALL_BATCH_THRESHOLD + if is_varlen_decode: + dense_small_hv_parallel = False + else: + dense_small_hv_parallel_head_threshold = ( + N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD + ) + dense_small_hv_parallel = ( + use_small_batch and dense_small_hv_parallel_head_threshold >= H and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + ) + num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) + + compiled_kernel = _get_compiled_packed_kernel( + N, + H, + HV, + K, + V, + mixed_qkv.stride(0), + initial_state_source.shape[0], + use_small_batch, + is_varlen_decode, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=False, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + compiled_kernel( + cu_seqlens_to_use, + q_view, + k_view, + v_view, + a_kernel, + b_kernel, + A_log, + dt_bias, + initial_state_source, + initial_state_indices, + o, + _get_cached_stream(mixed_qkv.device), + ) + return o + + def _define_kernels(): """Define CuTe DSL kernels for KDA normal and varlen decode modes.""" @@ -1668,6 +1861,149 @@ def _get_compiled_kernel( return compiled_kernel +def _get_compiled_packed_kernel( + N, + H, + HV, + K, + V, + row_stride, + pool_size, + use_small_batch, + is_varlen_decode, + scale, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + num_blocks_per_state_small, + dense_small_hv_parallel, + softplus_beta, + softplus_threshold, +): + """Get or lazily compile a packed-QKV kernel with static packed strides. + + The q/k/v mock tensors are strided views sliced from a packed + ``mixed_qkv`` mock, so their row stride is the compile-time constant + ``row_stride``. Runtime packed views with the same shape/stride can then + use the static layout specialization while still avoiding q/k/v + materialization. + """ + global _compiled_kernels + + qkv_dim = 2 * H * K + HV * V + key = ( + N, + H, + HV, + K, + V, + qkv_dim, + row_stride, + pool_size, + use_small_batch, + is_varlen_decode, + scale, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + num_blocks_per_state_small, + dense_small_hv_parallel, + softplus_beta, + softplus_threshold, + "packed_qkv", + ) + if key in _compiled_kernels: + return _compiled_kernels[key] + + cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda") + if row_stride < qkv_dim: + raise ValueError(f"row_stride={row_stride} must be >= qkv_dim={qkv_dim}") + mixed_qkv = torch.zeros(N, row_stride, dtype=torch.bfloat16, device="cuda") + + if is_varlen_decode: + q = mixed_qkv.narrow(1, 0, H * K).view(1, N, H, K) + k = mixed_qkv.narrow(1, H * K, H * K).view(1, N, H, K) + v = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(1, N, HV, V) + a = torch.zeros(N, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") + else: + q = mixed_qkv.narrow(1, 0, H * K).view(N, 1, H, K) + k = mixed_qkv.narrow(1, H * K, H * K).view(N, 1, H, K) + v = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(N, 1, HV, V) + a = torch.zeros(N, 1, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, 1, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") + + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + if state_layout_is_kv: + h0_source = torch.zeros(pool_size, HV, K, V, dtype=torch.float32, device="cuda") + else: + h0_source = torch.zeros(pool_size, HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + + cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16) + q_tensor = from_dlpack(q, assumed_align=16) + k_tensor = from_dlpack(k, assumed_align=16) + v_tensor = from_dlpack(v, assumed_align=16) + a_tensor = from_dlpack(a, assumed_align=16) + b_tensor = from_dlpack(b, assumed_align=16) + A_log_tensor = from_dlpack(A_log, assumed_align=16) + dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) + h0_source_tensor = from_dlpack(h0_source, assumed_align=16) + h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16) + o_tensor = from_dlpack(o, assumed_align=16) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + run_small, run_small_varlen, run_large, run_large_varlen = _get_jit_functions() + if use_small_batch: + kernel_func = run_small_varlen if is_varlen_decode else run_small + else: + kernel_func = run_large_varlen if is_varlen_decode else run_large + + compiled_kernel = cute.compile( + kernel_func, + cu_seqlens_tensor, + q_tensor, + k_tensor, + v_tensor, + a_tensor, + b_tensor, + A_log_tensor, + dt_bias_tensor, + h0_source_tensor, + h0_indices_tensor, + o_tensor, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + B=1 if is_varlen_decode else N, + T=N if is_varlen_decode else 1, + H=H, + K=K, + V=V, + HV=HV, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=precomputed_decay_beta, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + stream=stream, + options="--enable-tvm-ffi --opt-level 1", + ) + + _compiled_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA packed static-stride kernel compiled: " + f"N={N}, H={H}, HV={HV}, K={K}, V={V}, qkv_dim={qkv_dim}, row_stride={row_stride}, pool_size={pool_size}, " + f"small_batch={use_small_batch}, varlen={is_varlen_decode}" + ) + return compiled_kernel + + def _normalize_A_log(A_log: torch.Tensor, HV: int) -> torch.Tensor: if A_log.numel() != HV: raise ValueError(f"Unexpected A_log shape: {A_log.shape}; expected numel={HV}") @@ -2094,3 +2430,227 @@ def kda_decode( ) return o + + +def kda_packed_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + state_indices: torch.Tensor, + out: torch.Tensor | None = None, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + cu_seqlens: torch.Tensor | None = None, + state_layout: str = "vk", +) -> torch.Tensor: + """Packed-QKV variant of :func:`kda_decode`. + + Takes a packed ``mixed_qkv`` of shape ``[N, qkv_dim]`` laid out as + ``[Q(H·K) | K(H·K) | V(HV·V)]`` (head-major, last dim contiguous, row stride + >= qkv_dim) and feeds the existing CuTe DSL KDA decode kernel q/k/v as + strided views directly — avoiding the q/k/v materialization and the + ``.contiguous()`` copy that ``kda_decode`` performs. + + Numerics match ``kda_decode`` on the same inputs repacked, since the kernel + body is unchanged. The packed static-stride compile path builds q/k/v mock + tensors with the same row stride as runtime packed views, so the kernel + accepts the non-contiguous views without dynamic layout. + + Args: + mixed_qkv: ``[N, qkv_dim]`` bf16, last dim contiguous. + a: gate, dense ``(N,1,HV,K)`` / ``(N,HV,K)`` / varlen-compatible; see + ``_normalize_kda_a``. + b: ``(N,1,HV)`` (dense) or ``(N,HV)`` (varlen), bf16. + A_log: ``(HV,)`` fp32. + dt_bias: ``(HV,K)`` fp32. + state: ``(num_slots, HV, V, K)`` for ``state_layout='vk'`` or + ``(num_slots, HV, K, V)`` for ``'kv'``, fp32. + state_indices: ``(N,)`` int32, ``-1`` marks dummy slots. + out: optional preallocated output, dense ``(N,1,HV,V)`` or varlen + ``(1,N,HV,V)``, bf16, contiguous. + cu_seqlens: when not None, varlen decode (otherwise dense). + state_layout: ``"vk"`` (default) or ``"kv"``. + + Returns: + Output tensor of shape ``(N,1,HV,V)`` (dense) or ``(1,N,HV,V)`` + (varlen). ``state`` is updated in place. + """ + state_layout_canonical = _canonicalize_state_layout(state_layout) + state_layout_is_kv = state_layout_canonical == "kv" + + if state.dim() != 4: + raise ValueError(f"Unexpected state shape: {state.shape}; expected a 4D state tensor") + + if state_layout_is_kv: + pool_size, HV, K, V = state.shape + else: + pool_size, HV, V, K = state.shape + + if K != TILE_K: + raise ValueError(f"Current CuTe DSL KDA kernel requires K={TILE_K}, got K={K}") + if V % TILE_V_SMALL != 0 or V % TILE_V != 0: + raise ValueError(f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0 and V % {TILE_V} == 0, got V={V}") + + if mixed_qkv.ndim != 2: + raise ValueError(f"mixed_qkv must be 2D [N, qkv_dim], got shape {mixed_qkv.shape}") + qkv_dim = mixed_qkv.shape[1] + qk_dim = qkv_dim - HV * V + if qk_dim % 2 != 0: + raise ValueError( + f"mixed_qkv q/k segment (qkv_dim - HV*V = {qk_dim}) must be even for equal Q|K, got qkv_dim={qkv_dim}" + ) + H = (qk_dim // 2) // K + if H <= 0 or qk_dim // 2 != H * K: + raise ValueError( + f"Inconsistent mixed_qkv layout: q/k segment {qk_dim // 2} not divisible by K={K}, got qkv_dim={qkv_dim}" + ) + if HV % H != 0: + raise ValueError(f"HV={HV} must be divisible by H={H} (grouped-head mapping)") + if qkv_dim != 2 * H * K + HV * V: + raise ValueError( + f"mixed_qkv qkv_dim={qkv_dim} != 2*H*K + HV*V = {2 * H * K + HV * V} for H={H}, HV={HV}, K={K}, V={V}" + ) + + N = state_indices.shape[0] + is_varlen_decode = cu_seqlens is not None + if mixed_qkv.shape[0] != N: + raise ValueError(f"mixed_qkv batch size {mixed_qkv.shape[0]} must match state_indices length {N}") + row_stride = mixed_qkv.stride(0) + if mixed_qkv.stride(-1) != 1 or row_stride < qkv_dim: + raise ValueError( + f"mixed_qkv must use packed-row layout with stride(-1)=1 and stride(0)>=qkv_dim={qkv_dim}; " + f"got stride={mixed_qkv.stride()}" + ) + + if scale is None: + scale = K**-0.5 + elif scale <= 0: + raise ValueError(f"scale must be positive, got {scale}") + + fast_out = _try_fast_dense_packed_decode( + A_log, + dt_bias, + mixed_qkv, + a, + b, + N=N, + H=H, + HV=HV, + K=K, + V=V, + initial_state_source=state, + initial_state_indices=state_indices, + is_varlen_decode=is_varlen_decode, + cu_seqlens=cu_seqlens, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + out=out, + state_layout=state_layout_canonical, + ) + if fast_out is not None: + return fast_out + + # --- General path: normalize then launch the packed static-stride kernel. --- + # We deliberately do NOT route through kda_decode's slow path, which would + # .contiguous()-copy the q/k/v (defeating packed) and requires separate + # q/k/v tensors. We normalize the meta tensors (a/b/A_log/dt_bias/state/ + # indices) but feed q/k/v as strided views into the packed kernel directly. + h0_source, pool_size, _ = _normalize_state_source( + state, + N=N, + HV=HV, + K=K, + V=V, + device=mixed_qkv.device, + state_layout=state_layout_canonical, + ) + a_kernel = _normalize_kda_a(a, is_varlen_decode=is_varlen_decode, N=N, HV=HV, K=K) + a_kernel = a_kernel if a_kernel.is_contiguous() else a_kernel.contiguous() + if is_varlen_decode: + if b.dim() == 3: + b_kernel = b.squeeze(0) + else: + b_kernel = b + else: + if b.dim() == 2: + b_kernel = b.unsqueeze(1) + else: + b_kernel = b + b_kernel = b_kernel if b_kernel.is_contiguous() else b_kernel.contiguous() + A_log_n = _normalize_A_log(A_log, HV) + dt_bias_n = _normalize_dt_bias(dt_bias, HV, K) + indices_n = _normalize_state_indices(state_indices, N=N, pool_size=pool_size, device=mixed_qkv.device) + + if is_varlen_decode: + q_view = mixed_qkv.narrow(1, 0, H * K).view(1, N, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(1, N, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(1, N, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (1, N, HV, V)) + else: + q_view = mixed_qkv.narrow(1, 0, H * K).view(N, 1, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(N, 1, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(N, 1, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (N, 1, HV, V)) + + if cu_seqlens is not None: + cu_seqlens_to_use = cu_seqlens.contiguous() + else: + cache_key = (N, str(mixed_qkv.device)) + if cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cache_key] = torch.arange(N + 1, dtype=torch.int32, device=mixed_qkv.device) + cu_seqlens_to_use = _cu_seqlens_cache[cache_key] + + use_small_batch = N < SMALL_BATCH_THRESHOLD + if is_varlen_decode: + dense_small_hv_parallel = False + else: + dense_small_hv_parallel_head_threshold = ( + N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD + ) + dense_small_hv_parallel = ( + use_small_batch and dense_small_hv_parallel_head_threshold >= H and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + ) + num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) + + compiled_kernel = _get_compiled_packed_kernel( + N, + H, + HV, + K, + V, + row_stride, + pool_size, + use_small_batch, + is_varlen_decode, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=False, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + compiled_kernel( + cu_seqlens_to_use, + q_view, + k_view, + v_view, + a_kernel, + b_kernel, + A_log_n, + dt_bias_n, + h0_source, + indices_n, + o, + _get_cached_stream(mixed_qkv.device), + ) + return o diff --git a/tests/test_kda_packed_decode.py b/tests/test_kda_packed_decode.py new file mode 100644 index 00000000..259baed9 --- /dev/null +++ b/tests/test_kda_packed_decode.py @@ -0,0 +1,807 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unit tests for kda_packed_decode (packed-QKV CuTe DSL KDA decode kernel). + +Numerical ground truth = the existing non-packed ``kda_decode``: the same base +inputs (q, k, v, a, b, A_log, dt_bias, state) are repacked into a mixed_qkv of +shape ``[N, qkv_dim] = [Q(H·K) | K(H·K) | V(HV·V)]`` (head-major, last dim +contiguous) and ``kda_packed_decode`` is checked against ``kda_decode``. +""" + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.kda import kda_decode, kda_packed_decode +from tests.test_kda_decode import _assert_close, make_inputs # noqa: F401 + +K = 128 + + +def _pack_mixed_qkv(q, k, v, N): + """q,k: (N,H,K) bf16 ; v: (N,HV,V) bf16 -> mixed_qkv (N, qkv_dim) bf16.""" + return torch.cat([q.view(N, -1), k.view(N, -1), v.view(N, -1)], dim=-1) + + +# --------------------------------------------------------------------------- +# Dense: packed vs non-packed kda_decode +# --------------------------------------------------------------------------- +def _run_nonpacked_dense(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + state_ref = state.clone() + o = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(1), state_ref # (N,HV,V), (N,HV,V,K) + + +def _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + mixed = _pack_mixed_qkv(q, k, v, N) + state_p = state.clone() + o = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(1), state_p # (N,HV,V), (N,HV,V,K) + + +@pytest.mark.parametrize("N", [1, 2, 8, 16, 32, 64, 128]) +@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)]) +@pytest.mark.parametrize("V", [128, 256]) +def test_packed_dense(N, H, HV, V): + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + + # Sanity: the packed should be ~identical to non-packed (same kernel, strided view). + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Varlen: packed vs non-packed kda_decode +# --------------------------------------------------------------------------- +def _run_nonpacked_varlen(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + state_ref = state.clone() + o = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(0), state_ref # (N,HV,V), (N,HV,V,K) + + +def _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + mixed = _pack_mixed_qkv(q, k, v, N) + state_p = state.clone() + o = kda_packed_decode( + mixed, + a.contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(0), state_p # (N,HV,V), (N,HV,V,K) + + +@pytest.mark.parametrize("N", [2, 8, 16, 32, 64, 128]) +@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)]) +@pytest.mark.parametrize("V", [128, 256]) +def test_packed_varlen(N, H, HV, V): + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Equal-head coverage: H == HV (group ratio = 1) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("N", [1, 8, 64]) +@pytest.mark.parametrize("H", [8, 16, 32, 64]) +def test_packed_dense_equal_heads(N, H): + HV, V = H, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +@pytest.mark.parametrize("N", [2, 8, 64]) +@pytest.mark.parametrize("H", [8, 16, 32, 64]) +def test_packed_varlen_equal_heads(N, H): + HV, V = H, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# -1 dummy slots: output 0, corresponding state untouched +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("N", [1, 2, 8]) +def test_packed_minus1_dummy(N): + H, HV, V = 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + # pool larger than N; mark the first token as dummy (-1). + pool = N + 2 + state_pool = torch.zeros(pool, HV, V, K, device="cuda", dtype=torch.float32) + real_idx = torch.arange(N, device="cuda", dtype=torch.int32) + 1 # use slots 1..N + state_pool[real_idx] = state + indices = real_idx.clone() + indices[0] = -1 # first batch token is a dummy + + mixed = _pack_mixed_qkv(q, k, v, N) + state_before = state_pool.clone() + o = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_pool, + state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + ).squeeze(1) # (N,HV,V) + + # Dummy tokens (pool_idx == -1) skip ALL kernel work — the kernel neither + # reads/writes their state slot nor writes their output row. So: + # - the dummy's pool slot is untouched, and + # - the dummy's output row is UNINITIALIZED (leave it out of checks). + dummy_slot = int(real_idx[0]) if indices[0] == -1 else None + assert dummy_slot is not None + assert torch.equal(state_pool[dummy_slot], state_before[dummy_slot]), "dummy state slot was modified" + + # Compare against non-packed using the SAME pool + SAME indices (no remap), + # so packed and non-packed write the same physical state slots. + state_ref_pool = state_before.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref_pool, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + ).squeeze(1) + + if N > 1: + _assert_close("real output", o_ref[1:].float(), o[1:].float()) + _assert_close("real state", state_ref_pool[real_idx[1:]], state_pool[real_idx[1:]]) + else: + # N == 1: the only token is the dummy, so there is no real row to check. + # Just confirm the dummy slot's state survived untouched (already asserted). + pass + + +# --------------------------------------------------------------------------- +# KV state layout +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_kv_state_layout(is_varlen): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state_vk = make_inputs(N, H, HV, K, V) + state_kv = state_vk.permute(0, 1, 3, 2).contiguous() # (N,HV,K,V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(N + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + if is_varlen: + state_p = state_kv.clone() + o_kv = kda_packed_decode( + mixed, + a.contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + state_layout="kv", + ).squeeze(0) + else: + state_p = state_kv.clone() + o_kv = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + state_layout="kv", + ).squeeze(1) + + # Compare against vk-layout non-packed reference (same numerics). + state_vk_ref = state_vk.clone() + if is_varlen: + o_vk = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_vk_ref, + initial_state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ).squeeze(0) + else: + o_vk = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_vk_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("kv output", o_kv.float(), o_vk.float()) + # state: kv layout transposed back to vk for comparison + _assert_close("kv state", state_vk_ref, state_p.permute(0, 1, 3, 2).contiguous()) + + +# --------------------------------------------------------------------------- +# No L2 norm path +# --------------------------------------------------------------------------- +def test_packed_no_l2norm(): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + + state_p = state.clone() + o_p = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=False, + ).squeeze(1) + + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=False, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Focused API compatibility and validation coverage +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_explicit_out_scale_and_softplus(is_varlen): + N, H, HV, V = 8, 8, 16, 128 + scale = 0.25 + softplus_beta = 0.75 + softplus_threshold = 10.0 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(N + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + state_p = state.clone() + if is_varlen: + out = torch.empty(1, N, HV, V, device="cuda", dtype=torch.bfloat16) + o_p = kda_packed_decode( + mixed, + a.contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + out=out, + cu_seqlens=cu_seqlens, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + assert o_p is out + o_p_cmp = o_p.squeeze(0) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ).squeeze(0) + else: + out = torch.empty(N, 1, HV, V, device="cuda", dtype=torch.bfloat16) + o_p = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + out=out, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + assert o_p is out + o_p_cmp = o_p.squeeze(1) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p_cmp.float()) + _assert_close("state", state_ref, state_p) + + +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_slow_path_compatible_a_b_shapes(is_varlen): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(N + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + state_p = state.clone() + if is_varlen: + # Fast path expects a=(N,HV,K), b=(N,HV); these compatible public + # shapes force the general normalization path. + o_p = kda_packed_decode( + mixed, + a.reshape(1, N, HV * K).contiguous(), + b.unsqueeze(0).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ).squeeze(0) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ).squeeze(0) + else: + # Fast path does not accept dense a=(N,HV*K), so this covers the + # packed general path while preserving the same numerics. + o_p = kda_packed_decode( + mixed, + a.reshape(N, HV * K).contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + ).squeeze(1) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +def test_packed_mixed_qkv_allows_padded_row_stride(): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + padded = torch.empty(N, mixed.shape[1] + 8, device="cuda", dtype=torch.bfloat16) + mixed_padded_view = padded[:, : mixed.shape[1]] + mixed_padded_view.copy_(mixed) + assert mixed_padded_view.stride(-1) == 1 + assert mixed_padded_view.stride(0) != mixed.shape[1] + + indices = torch.arange(N, device="cuda", dtype=torch.int32) + state_p = state.clone() + o_p = kda_packed_decode( + mixed_padded_view, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + ).squeeze(1) + + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +def test_packed_mixed_qkv_allows_padded_row_stride_n1(): + N, H, HV, V = 1, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + padded = torch.empty(N, mixed.shape[1] + 8, device="cuda", dtype=torch.bfloat16) + mixed_padded_view = padded[:, : mixed.shape[1]] + mixed_padded_view.copy_(mixed) + assert mixed_padded_view.stride(-1) == 1 + assert mixed_padded_view.stride(0) != mixed.shape[1] + + indices = torch.arange(N, device="cuda", dtype=torch.int32) + state_p = state.clone() + o_p = kda_packed_decode( + mixed_padded_view, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + ).squeeze(1) + + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_general_path_accepts_noncontiguous_a_b(is_varlen): + N, H, HV, V = 4, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + mixed = _pack_mixed_qkv(q, k, v, N) + + a_flat = a.reshape(N, HV * K) + a_padded = torch.empty(N, HV * K + 1, device="cuda", dtype=torch.bfloat16) + a_arg = a_padded[:, : HV * K] + a_arg.copy_(a_flat) + b_padded = torch.empty(N, HV + 1, device="cuda", dtype=torch.bfloat16) + b_arg = b_padded[:, :HV] + b_arg.copy_(b) + assert not a_arg.is_contiguous() + assert not b_arg.is_contiguous() + + state_p = state.clone() + if is_varlen: + o_p = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(0) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(0) + else: + o_p = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(1) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Padded batch + CUDA graph capture/replay +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_padded_batch(is_varlen): + N_real, N_pad, H, HV, V = 6, 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state_real = make_inputs(N_real, H, HV, K, V) + + # Pad to N_pad with extra rows; mark them as dummy (-1). + q_pad = torch.zeros(N_pad, H, K, device="cuda", dtype=torch.bfloat16) + k_pad = torch.zeros(N_pad, H, K, device="cuda", dtype=torch.bfloat16) + v_pad = torch.zeros(N_pad, HV, V, device="cuda", dtype=torch.bfloat16) + a_pad = torch.zeros(N_pad, HV, K, device="cuda", dtype=torch.bfloat16) + b_pad = torch.zeros(N_pad, HV, device="cuda", dtype=torch.bfloat16) + q_pad[:N_real], k_pad[:N_real], v_pad[:N_real] = q, k, v + a_pad[:N_real] = a + b_pad[:N_real] = b + + pool = N_pad + state_pool = torch.zeros(pool, HV, V, K, device="cuda", dtype=torch.float32) + indices = torch.arange(N_pad, device="cuda", dtype=torch.int32) + indices[N_real:] = -1 # padded slots are dummies + state_pool[:N_real] = state_real + + mixed = _pack_mixed_qkv(q_pad, k_pad, v_pad, N_pad) + cu_seqlens = torch.arange(N_pad + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + # Forward args shared by eager + graph path + if is_varlen: + a_arg, b_arg = a_pad.contiguous(), b_pad.contiguous() + else: + a_arg, b_arg = a_pad.unsqueeze(1).contiguous(), b_pad.unsqueeze(1).contiguous() + + state_p = state_pool.clone() + o_p = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ) + if is_varlen: + o_real = o_p[:, :N_real] # (1,N_real,HV,V) + else: + o_real = o_p[:N_real] # (N_real,1,HV,V) + + # Compare real rows against non-packed run on the real-only inputs. + state_ref = state_real.clone() + if is_varlen: + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N_real, device="cuda", dtype=torch.int32), + cu_seqlens=torch.arange(N_real + 1, device="cuda", dtype=torch.int32), + scale=scale, + ) # (1,N_real,HV,V) + else: + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N_real, device="cuda", dtype=torch.int32), + scale=scale, + ) # (N_real,1,HV,V) + + _assert_close("real output", o_ref.float(), o_real.float()) + # dummy (padded) state slots must be untouched; their output rows are + # UNINITIALIZED by the kernel (no write), so we don't assert on them. + assert torch.equal(state_p[N_real:], state_pool[N_real:]), "dummy state slots were modified" + + # CUDA graph capture + replay smoke (capture must not raise; replay must not raise). + # Run once eagerly first so the kernel is compiled and any warming cu_seqlens + # cache is populated before capture. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + state_g = state_pool.clone() + _ = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_g, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ) + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_g2 = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_g, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ) + g.replay() + assert out_g2 is not None From 9ff1edb1a0279dab2b4fb8045c8e699a6a63e266 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=97=A0=E8=A8=80=E7=8B=AC=E4=B8=8A=E6=9C=BA=E6=88=BF?= <88866917+sjmshsh@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:08:59 +0800 Subject: [PATCH 32/34] [Feature] Add GVA support for Lightning (#85) * support gva * support gva * style: format Lightning Attention files --------- Co-authored-by: sunnyxyli --- benchmarks/bench_la_decode_vs_fla.py | 64 ++++---- benchmarks/bench_lightning_attn.py | 134 +++++++++++------ cula/ops/lightning/decode.py | 90 +++++++---- cula/ops/lightning/prefill_sm100.py | 215 +++++++++++++++++---------- tests/test_la_decode.py | 113 +++++++++++--- tests/test_la_decode_pool.py | 53 ++++++- tests/test_lightning_attn.py | 162 ++++++++++++-------- 7 files changed, 555 insertions(+), 276 deletions(-) diff --git a/benchmarks/bench_la_decode_vs_fla.py b/benchmarks/bench_la_decode_vs_fla.py index 75f1dc79..da9f6e58 100644 --- a/benchmarks/bench_la_decode_vs_fla.py +++ b/benchmarks/bench_la_decode_vs_fla.py @@ -38,6 +38,7 @@ Usage: python benchmarks/bench_la_decode_vs_fla.py python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128 + python benchmarks/bench_la_decode_vs_fla.py --heads 32 --num-v-heads 64 python benchmarks/bench_la_decode_vs_fla.py --batch-sizes 1 8 64 256 """ @@ -63,28 +64,32 @@ # ───────────────────────────────────────────────────────────────────────────── # Core benchmark for one configuration # ───────────────────────────────────────────────────────────────────────────── -def run_config(B, H, K, V, layer_idx, num_layers): +def run_config(B, H, HV, K, V, layer_idx, num_layers): + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" device = "cuda" dtype = torch.bfloat16 scale = K**-0.5 + group_size = HV // H # Per-head log decay (Lightning Attention formula) - g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * torch.arange(H, device=device, dtype=torch.float32) + g_gamma = -(8 / HV * (1 - layer_idx / num_layers)) * torch.arange(HV, device=device, dtype=torch.float32) decay_scales = -g_gamma # la_decode convention # ── Random inputs ────────────────────────────────────────────────────── torch.manual_seed(42) q_4d = torch.randn(B, 1, H, K, device=device, dtype=dtype) k_4d = torch.randn(B, 1, H, K, device=device, dtype=dtype) - v_4d = torch.randn(B, 1, H, V, device=device, dtype=dtype) - state_init = torch.randn(B, H, K, V, device=device, dtype=torch.float32) * 0.01 + v_4d = torch.randn(B, 1, HV, V, device=device, dtype=dtype) + state_init = torch.randn(B, HV, K, V, device=device, dtype=torch.float32) * 0.01 + q_fla_4d = q_4d.repeat_interleave(group_size, dim=2) + k_fla_4d = k_4d.repeat_interleave(group_size, dim=2) # ── fla reference output ─────────────────────────────────────────────── state_fla = state_init.clone() with torch.no_grad(): o_fla_fp32, ht_fla = fused_recurrent_fwd( - q_4d, - k_4d, + q_fla_4d, + k_fla_4d, v_4d, g_gamma=g_gamma, scale=scale, @@ -94,11 +99,11 @@ def run_config(B, H, K, V, layer_idx, num_layers): o_fla = o_fla_fp32.to(dtype) # ── la_decode output ─────────────────────────────────────────────────── - state_cute = state_init.clone().permute(0, 1, 3, 2).reshape(B * H, V, K).contiguous() + state_cute = state_init.clone().permute(0, 1, 3, 2).reshape(B * HV, V, K).contiguous() q_3d = q_4d.squeeze(1) k_3d = k_4d.squeeze(1) v_3d = v_4d.squeeze(1) - out_cute = torch.zeros(B, H, V, device=device, dtype=dtype) + out_cute = torch.zeros(B, HV, V, device=device, dtype=dtype) s_offsets = torch.arange(B, device=device, dtype=torch.int32) with torch.no_grad(): @@ -128,7 +133,7 @@ def run_config(B, H, K, V, layer_idx, num_layers): max_ref = torch.abs(o_fla_cmp).max().item() rel_maxdiff = torch.abs(o_cute_cmp - o_fla_cmp).max().item() / (max_ref + 1e-8) - state_cute_back = state_cute.reshape(B, H, V, K).permute(0, 1, 3, 2).contiguous() + state_cute_back = state_cute.reshape(B, HV, V, K).permute(0, 1, 3, 2).contiguous() state_relative_rms_error = relative_rms_error(ht_fla, state_cute_back) # ================================================================== @@ -140,16 +145,16 @@ def run_config(B, H, K, V, layer_idx, num_layers): BV_fla = min(triton.next_power_of_2(V), 64) NK = triton.cdiv(K, BK_fla) NV = triton.cdiv(V, BV_fla) - fla_o_buf = torch.empty(NK, B, 1, H, V, device=device, dtype=torch.float32) - fla_ht_buf = torch.empty(B, H, K, V, device=device, dtype=torch.float32) - fla_o_sum = torch.empty(B, 1, H, V, device=device, dtype=torch.float32) + fla_o_buf = torch.empty(NK, B, 1, HV, V, device=device, dtype=torch.float32) + fla_ht_buf = torch.empty(B, HV, K, V, device=device, dtype=torch.float32) fla_state_k = state_init.clone() - grid_fla = (NV, NK, B * H) + fla_o_sum = torch.empty(B, 1, HV, V, device=device, dtype=torch.float32) + grid_fla = (NV, NK, B * HV) def kernel_fla(): fused_recurrent_fwd_kernel[grid_fla]( - q=q_4d, - k=k_4d, + q=q_fla_4d, + k=k_fla_4d, v=v_4d, g=None, g_gamma=g_gamma, @@ -162,7 +167,7 @@ def kernel_fla(): scale=scale, B=B, T=1, - H=H, + H=HV, K=K, V=V, BK=BK_fla, @@ -176,9 +181,9 @@ def kernel_fla(): torch.sum(fla_o_buf, dim=0, out=fla_o_sum) # cute kernel: pre-create compiled + stream handle - cute_state_k = state_init.clone().permute(0, 1, 3, 2).reshape(B * H, V, K).contiguous() - out_cute_k = torch.empty(B, H, V, device=device, dtype=dtype) - cache = _get_compiled_kernel(B, 1, H, K, V, cute_state_k.shape[0], scale, USE_FAST_MATH) + cute_state_k = state_init.clone().permute(0, 1, 3, 2).reshape(B * HV, V, K).contiguous() + out_cute_k = torch.empty(B, HV, V, device=device, dtype=dtype) + cache = _get_compiled_kernel(B, 1, H, HV, K, V, cute_state_k.shape[0], scale, USE_FAST_MATH) compiled_cute = cache["compiled"] stream_handle = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) @@ -193,13 +198,13 @@ def kernel_cute(): # Mode 2: WRAPPER (full call path as used in production) # ================================================================== wrap_fla_state = state_init.clone() - wrap_cute_state = state_init.clone().permute(0, 1, 3, 2).reshape(B * H, V, K).contiguous() - wrap_cute_out = torch.empty(B, H, V, device=device, dtype=dtype) + wrap_cute_state = state_init.clone().permute(0, 1, 3, 2).reshape(B * HV, V, K).contiguous() + wrap_cute_out = torch.empty(B, HV, V, device=device, dtype=dtype) def wrapper_fla(): fused_recurrent_fwd( - q_4d, - k_4d, + q_fla_4d, + k_fla_4d, v_4d, g_gamma=g_gamma, scale=scale, @@ -233,6 +238,8 @@ def wrapper_cute(): return { "B": B, + "H": H, + "HV": HV, "kernel_fla_ms": kernel_fla_ms, "kernel_cute_ms": kernel_cute_ms, "kernel_speedup": kernel_fla_ms / kernel_cute_ms, @@ -257,16 +264,19 @@ def main(): default=[1, 2, 4, 8, 16, 32, 64, 128, 256], ) parser.add_argument("--heads", type=int, default=32) + parser.add_argument("--num-v-heads", type=int, default=None, help="Number of value heads (default: same as --heads)") parser.add_argument("--head-dim", type=int, default=128) parser.add_argument("--layer-idx", type=int, default=12) parser.add_argument("--num-layers", type=int, default=24) args = parser.parse_args() - H, K, V = args.heads, args.head_dim, args.head_dim + H, HV, K, V = args.heads, args.num_v_heads or args.heads, args.head_dim, args.head_dim + if HV < H or HV % H != 0: + raise ValueError(f"num_v_heads ({HV}) must be >= heads ({H}) and divisible by heads") print("Lightning Attention Decode Benchmark") print(" la_decode (CuTe DSL) vs fla fused_recurrent_fwd (Triton)") - print(f" H={H}, K={K}, V={V}, layer={args.layer_idx}/{args.num_layers}") + print(f" H={H}, HV={HV}, K={K}, V={V}, layer={args.layer_idx}/{args.num_layers}") print(" dtype=bf16, state=fp32, T=1") # ── Kernel-only comparison ────────────────────────────────────────── @@ -282,7 +292,7 @@ def main(): results = [] for B in args.batch_sizes: - r = run_config(B, H, K, V, args.layer_idx, args.num_layers) + r = run_config(B, H, HV, K, V, args.layer_idx, args.num_layers) results.append(r) print( f"{r['B']:>5} | {r['kernel_fla_ms']:>10.4f} | {r['kernel_cute_ms']:>10.4f} | " @@ -293,7 +303,7 @@ def main(): # ── Wrapper comparison ────────────────────────────────────────────── print(f"\n{'=' * 100}") print(" Mode 2: WRAPPER (fused_recurrent_fwd vs linear_attention_decode, full call path)") - print(" fla: alloc o[NK,B,1,H,V]+ht[B,H,K,V] + kernel + sum(0); cute: cache lookup + CUstream + kernel") + print(" fla: alloc o[NK,B,1,HV,V]+ht[B,HV,K,V] + kernel + sum(0); cute: cache lookup + CUstream + kernel") print(f"{'=' * 100}") print(f"{'B':>5} | {'fla (ms)':>10} | {'cute (ms)':>10} | {'speedup':>8}") print("─" * 50) diff --git a/benchmarks/bench_lightning_attn.py b/benchmarks/bench_lightning_attn.py index bf8db7a3..62f77386 100644 --- a/benchmarks/bench_lightning_attn.py +++ b/benchmarks/bench_lightning_attn.py @@ -33,6 +33,9 @@ # Custom varlen workloads python benchmarks/bench_lightning_attn.py --modes varlen --num-heads 32 64 --iterations 50 + + # GVA (value heads > Q/K heads) + python benchmarks/bench_lightning_attn.py --modes no_state h0_ht varlen --num-heads 32 --num-v-heads 64 """ import argparse @@ -80,6 +83,17 @@ def compute_decay(H, layer_idx=12, num_layers=24): return (8 / H * (1 - layer_idx / num_layers)) * torch.arange(H, dtype=torch.float32, device=DEVICE) +def expand_qk_to_value_heads(Q, K, V): + """Expand Q/K from H heads to HV heads for FLA and naive references.""" + H = Q.shape[2] + HV = V.shape[2] + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + if HV == H: + return Q, K + group_size = HV // H + return Q.repeat_interleave(group_size, dim=2), K.repeat_interleave(group_size, dim=2) + + @torch.no_grad() def torch_naive_lightning_attn(Q, K, V, decay, scale=1.0, initial_state=None, output_final_state=False): """Recurrent FP32 reference for lightning attention (simple_gla). @@ -87,15 +101,17 @@ def torch_naive_lightning_attn(Q, K, V, decay, scale=1.0, initial_state=None, ou O(B*T*H*D^2) — exact ground truth, all computation in FP32. """ B, T, H, D = Q.shape + HV = V.shape[2] + Q, K = expand_qk_to_value_heads(Q, K, V) q, k, v = Q.float(), K.float(), V.float() - decay_factor = torch.exp(-decay.float()) # [H] + decay_factor = torch.exp(-decay.float()) # [HV] S = ( initial_state.float().clone() if initial_state is not None - else torch.zeros(B, H, D, D, dtype=torch.float32, device=Q.device) + else torch.zeros(B, HV, D, D, dtype=torch.float32, device=Q.device) ) - O = torch.zeros(B, T, H, D, dtype=torch.float32, device=Q.device) + O = torch.zeros(B, T, HV, D, dtype=torch.float32, device=Q.device) for t in range(T): S = S * decay_factor[None, :, None, None] @@ -111,13 +127,14 @@ def torch_naive_lightning_attn(Q, K, V, decay, scale=1.0, initial_state=None, ou # ============================================================================= def run_fla(Q, K, V, decay, initial_state, output_final_state, warmup, iters): """Run FLA chunk_simple_gla_fwd (standard, non-varlen).""" + Q_fla, K_fla = expand_qk_to_value_heads(Q, K, V) g_gamma = -decay scale = 1.0 def fn(): return chunk_simple_gla_fwd( - q=Q, - k=K, + q=Q_fla, + k=K_fla, v=V, g_gamma=g_gamma, scale=scale, @@ -183,13 +200,14 @@ def fn(): def run_fla_varlen(Q, K, V, decay, cu_seqlens, warmup, iters): """Run FLA native varlen (single launch via cu_seqlens). FAIR baseline.""" + Q_fla, K_fla = expand_qk_to_value_heads(Q, K, V) g_gamma = -decay cu_long = cu_seqlens.to(torch.long) def fn(): return chunk_simple_gla_fwd( - q=Q, - k=K, + q=Q_fla, + k=K_fla, v=V, g_gamma=g_gamma, scale=1.0, @@ -207,25 +225,26 @@ def fn(): # ============================================================================= # Standard (non-varlen) benchmark # ============================================================================= -def benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, iters): +def benchmark_standard_config(B, T, H, HV, D, layer_idx, num_layers, mode, warmup, iters): """Benchmark a single standard (non-varlen) config. mode: "no_state" — no initial/final state "h0_ht" — provide random h0 and output ht """ + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" torch.manual_seed(42) Q = torch.randn(B, T, H, D, dtype=DTYPE, device=DEVICE) K = torch.randn(B, T, H, D, dtype=DTYPE, device=DEVICE) - V = torch.randn(B, T, H, D, dtype=DTYPE, device=DEVICE) - decay = compute_decay(H, layer_idx, num_layers) + V = torch.randn(B, T, HV, D, dtype=DTYPE, device=DEVICE) + decay = compute_decay(HV, layer_idx, num_layers) has_h0 = mode == "h0_ht" output_ht = mode == "h0_ht" - h0 = torch.randn(B, H, D, D, dtype=torch.float32, device=DEVICE) * 0.01 if has_h0 else None + h0 = torch.randn(B, HV, D, D, dtype=torch.float32, device=DEVICE) * 0.01 if has_h0 else None h0_fla = h0.clone() if h0 is not None else None h0_cute = h0.transpose(-1, -2).contiguous() if h0 is not None else None # BHVK for CuTe - result = {"B": B, "T": T, "H": H, "D": D, "mode": mode} + result = {"B": B, "T": T, "H": H, "HV": HV, "D": D, "mode": mode} ht_fla = None ht_cute = None @@ -289,20 +308,22 @@ def benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, i # ============================================================================= # Varlen benchmark # ============================================================================= -def benchmark_varlen_config(N, seq_lens, H, D, warmup, iters, dist=""): +def benchmark_varlen_config(N, seq_lens, H, HV, D, warmup, iters, dist=""): """Benchmark a varlen config: persistent vs non-persistent vs FLA varlen.""" + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" T = sum(seq_lens) torch.manual_seed(42) Q = torch.randn(1, T, H, D, dtype=DTYPE, device=DEVICE) K = torch.randn(1, T, H, D, dtype=DTYPE, device=DEVICE) - V = torch.randn(1, T, H, D, dtype=DTYPE, device=DEVICE) + V = torch.randn(1, T, HV, D, dtype=DTYPE, device=DEVICE) cu = torch.tensor([0] + list(np.cumsum(seq_lens)), dtype=torch.int32, device=DEVICE) - decay = compute_decay(H) + decay = compute_decay(HV) result = { "B": N, "T": T, "H": H, + "HV": HV, "D": D, "mode": "varlen", "seq_lens": seq_lens, @@ -394,7 +415,8 @@ def print_standard_header(): def print_standard_result(r): - cfg = f"B={r['B']},T={r['T']},H={r['H']}" + hv = r.get("HV", r["H"]) + cfg = f"B={r['B']},T={r['T']},H={r['H']},HV={hv}" fla = f"{r['fla_ms']:.3f}" if _valid(r.get("fla_ms", float("nan"))) else "ERR" dsl = f"{r['cutedsl_ms']:.3f}" if _valid(r.get("cutedsl_ms", float("nan"))) else "ERR" @@ -437,7 +459,8 @@ def print_varlen_header(): def print_varlen_result(r): - cfg = f"N={r['B']},T={r['T']},H={r['H']}" + hv = r.get("HV", r["H"]) + cfg = f"N={r['B']},T={r['T']},H={r['H']},HV={hv}" dist = r.get("dist", "") p_ms = f"{r['persistent_ms']:.3f}" if _valid(r.get("persistent_ms", float("nan"))) else "ERR" @@ -482,6 +505,7 @@ def run_benchmark_suite(args): warmup = args.warmup iters = args.iterations modes = args.modes + num_v_heads = getattr(args, "num_v_heads", None) print("\n" + "=" * 100) print("Lightning Attention Benchmark: CuteDSL vs FLA") @@ -490,6 +514,7 @@ def run_benchmark_suite(args): print(f" Batch sizes: {args.batch_sizes}") print(f" Seq lengths: {args.seq_lens}") print(f" Num heads: {args.num_heads}") + print(f" Num V heads: {num_v_heads or args.num_heads}") print(f" Head dim: {D}") print(f" Layer: {layer_idx}/{num_layers}") print(f" Warmup/Iters: {warmup}/{iters}") @@ -505,14 +530,18 @@ def run_benchmark_suite(args): for B in args.batch_sizes: for T in args.seq_lens: for H in args.num_heads: - total = B * T * H * D - if total > 2_147_483_648: - continue - if T > 4096 and B > 2: - continue - r = benchmark_standard_config(B, T, H, D, layer_idx, num_layers, mode, warmup, iters) - all_results.append(r) - print_standard_result(r) + for HV in num_v_heads or [H]: + if HV < H or HV % H != 0: + print(f"Skipping invalid GVA config H={H}, HV={HV}") + continue + total = B * T * HV * D + if total > 2_147_483_648: + continue + if T > 4096 and B > 2: + continue + r = benchmark_standard_config(B, T, H, HV, D, layer_idx, num_layers, mode, warmup, iters) + all_results.append(r) + print_standard_result(r) # ===================== Varlen mode ===================== if "varlen" in modes: @@ -542,22 +571,26 @@ def run_benchmark_suite(args): workloads = unique for H in args.num_heads: - print(f"\n --- H={H}, D={D} ---") - print_varlen_header() - - for N, T_total, dist in workloads: - if dist == "uniform": - seq_lens = gen_uniform(N, T_total) - elif dist == "skewed": - seq_lens = gen_skewed(N, T_total) - elif dist == "random": - seq_lens = gen_random(N, T_total) - else: - raise ValueError(f"Unknown dist: {dist}") + for HV in num_v_heads or [H]: + if HV < H or HV % H != 0: + print(f"Skipping invalid GVA config H={H}, HV={HV}") + continue + print(f"\n --- H={H}, HV={HV}, D={D} ---") + print_varlen_header() + + for N, T_total, dist in workloads: + if dist == "uniform": + seq_lens = gen_uniform(N, T_total) + elif dist == "skewed": + seq_lens = gen_skewed(N, T_total) + elif dist == "random": + seq_lens = gen_random(N, T_total) + else: + raise ValueError(f"Unknown dist: {dist}") - r = benchmark_varlen_config(N, seq_lens, H, D, warmup, iters, dist=dist) - all_results.append(r) - print_varlen_result(r) + r = benchmark_varlen_config(N, seq_lens, H, HV, D, warmup, iters, dist=dist) + all_results.append(r) + print_varlen_result(r) # ===================== Summary ===================== print(f"\n{'=' * 100}") @@ -662,7 +695,7 @@ def plot_results(all_results, modes): if not hr: ax.set_title("varlen (no data)") continue - labels = [f"N{r['B']}T{r['T']}\n{r.get('dist', '')[:3]}" for r in hr] + labels = [f"N{r['B']}T{r['T']}H{r['H']}HV{r.get('HV', r['H'])}\n{r.get('dist', '')[:3]}" for r in hr] p_ms = [r["persistent_ms"] for r in hr] np_ms = [r["nonpersistent_ms"] for r in hr] fla_ms = [r["fla_varlen_ms"] if _valid(r.get("fla_varlen_ms", float("nan"))) else 0 for r in hr] @@ -676,7 +709,7 @@ def plot_results(all_results, modes): if not hr: ax.set_title(f"{mode} (no data)") continue - labels = [f"B{r['B']}T{r['T']}H{r['H']}" for r in hr] + labels = [f"B{r['B']}T{r['T']}H{r['H']}HV{r.get('HV', r['H'])}" for r in hr] fla = [r["fla_ms"] for r in hr] dsl = [r["cutedsl_ms"] for r in hr] x = np.arange(len(labels)) @@ -703,6 +736,7 @@ def plot_results(all_results, modes): def generate_report(all_results, modes, args): from datetime import datetime + num_v_heads = getattr(args, "num_v_heads", None) path = os.path.join(os.path.dirname(__file__), "benchmark_report.md") with open(path, "w") as f: f.write("# Lightning Attention Benchmark Report\n\n") @@ -712,6 +746,7 @@ def generate_report(all_results, modes, args): f.write(f"- Batch sizes: {args.batch_sizes}\n") f.write(f"- Seq lengths: {args.seq_lens}\n") f.write(f"- Num heads: {args.num_heads}\n") + f.write(f"- Num V heads: {num_v_heads or args.num_heads}\n") f.write(f"- Head dim: {args.head_dim}\n") f.write(f"- Layer: {args.layer_idx}/{args.num_layers}\n") f.write(f"- Warmup/Iters: {args.warmup}/{args.iterations}\n\n") @@ -723,8 +758,12 @@ def generate_report(all_results, modes, args): f.write(f"## Mode: {mode}\n\n") if mode == "varlen": - f.write("| N | T | Dist | Persist(ms) | NonPer(ms) | FLA_vl(ms) | P/NP | P/FLAvl | O diff | ht diff |\n") - f.write("|---|---|------|-------------|------------|------------|------|---------|--------|--------|\n") + f.write( + "| N | T | H | HV | Dist | Persist(ms) | NonPer(ms) | FLA_vl(ms) | P/NP | P/FLAvl | O diff | ht diff |\n" + ) + f.write( + "|---|---|---|----|------|-------------|------------|------------|------|---------|--------|--------|\n" + ) for r in mr: p = f"{r['persistent_ms']:.3f}" if _valid(r.get("persistent_ms", float("nan"))) else "-" np_ = f"{r['nonpersistent_ms']:.3f}" if _valid(r.get("nonpersistent_ms", float("nan"))) else "-" @@ -736,7 +775,7 @@ def generate_report(all_results, modes, args): od = f"{r['p_vs_np_O_diff']:.1e}" if not np.isnan(r.get("p_vs_np_O_diff", float("nan"))) else "-" hd = f"{r['p_vs_np_ht_diff']:.1e}" if not np.isnan(r.get("p_vs_np_ht_diff", float("nan"))) else "-" f.write( - f"| {r['B']} | {r['T']} | {r.get('dist', '')} | {p} | {np_} | {fla_vl} | {pvnp} | {pvfla_vl} | {od} | {hd} |\n" + f"| {r['B']} | {r['T']} | {r['H']} | {r.get('HV', r['H'])} | {r.get('dist', '')} | {p} | {np_} | {fla_vl} | {pvnp} | {pvfla_vl} | {od} | {hd} |\n" ) else: has_ht = mode == "h0_ht" @@ -753,7 +792,7 @@ def generate_report(all_results, modes, args): "|--------|---------|-------------|---------|---------------------------|----------------------------|\n" ) for r in mr: - cfg = f"B={r['B']},T={r['T']},H={r['H']}" + cfg = f"B={r['B']},T={r['T']},H={r['H']},HV={r.get('HV', r['H'])}" sp = f"{r['speedup']:.2f}x" if _valid(r.get("speedup", float("nan"))) else "-" fla = f"{r['fla_ms']:.3f}" if _valid(r.get("fla_ms", float("nan"))) else "-" dsl = f"{r['cutedsl_ms']:.3f}" if _valid(r.get("cutedsl_ms", float("nan"))) else "-" @@ -838,6 +877,9 @@ def parse_args(): "--seq-lens", type=int, nargs="+", default=[256, 1024, 4096, 8192, 32768], help="Sequence lengths for standard modes" ) p.add_argument("--num-heads", type=int, nargs="+", default=[64], help="Number of heads to test") + p.add_argument( + "--num-v-heads", type=int, nargs="+", default=None, help="Number of value heads to test (default: same as H)" + ) p.add_argument("--head-dim", type=int, default=128) p.add_argument("--layer-idx", type=int, default=12) p.add_argument("--num-layers", type=int, default=24) diff --git a/cula/ops/lightning/decode.py b/cula/ops/lightning/decode.py index 08831dfb..6ab1b467 100644 --- a/cula/ops/lightning/decode.py +++ b/cula/ops/lightning/decode.py @@ -76,7 +76,7 @@ def la_decode_kernel_small_batch_pretranspose( smem_layout_staged: cute.Layout, vec_size: cutlass.Constexpr[int], num_v_tiles: cutlass.Constexpr[int], - decay_scales: cute.Tensor, # [H] + decay_scales: cute.Tensor, # [HV] q: cute.Tensor, # [B, T, H, K] k: cute.Tensor, # [B, T, H, K] v: cute.Tensor, # [B, T, HV, V] @@ -86,6 +86,7 @@ def la_decode_kernel_small_batch_pretranspose( B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], NUM_WARPS: cutlass.Constexpr[int] = 4, @@ -94,7 +95,6 @@ def la_decode_kernel_small_batch_pretranspose( ): """Each block uses pipeline to load one batch and vectorized writeback""" - HV = H tidx, _, _ = cute.arch.thread_idx() lane_id = tidx % 32 warp_idx = cute.arch.warp_idx() @@ -121,7 +121,7 @@ def la_decode_kernel_small_batch_pretranspose( r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) r_v = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) r_h = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) - r_decay_scale = -cutlass.Float32(decay_scales[i_h]) + r_decay_scale = -cutlass.Float32(decay_scales[i_hv]) r_decay = cute.exp(r_decay_scale, fastmath=USE_FAST_MATH) cute.arch.barrier() @@ -244,7 +244,7 @@ def la_decode_kernel_big_batch_pretranspose( smem_layout_staged: cute.Layout, vec_size: cutlass.Constexpr[int], num_v_tiles: cutlass.Constexpr[int], - decay_scales: cute.Tensor, # [H] + decay_scales: cute.Tensor, # [HV] q: cute.Tensor, # [B, T, H, K] k: cute.Tensor, # [B, T, H, K] v: cute.Tensor, # [B, T, HV, V] @@ -254,6 +254,7 @@ def la_decode_kernel_big_batch_pretranspose( B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], NUM_WARPS: cutlass.Constexpr[int] = 4, @@ -262,7 +263,6 @@ def la_decode_kernel_big_batch_pretranspose( ): """Each block uses pipeline to load one batch and vectorized writeback""" - HV = H tidx, _, _ = cute.arch.thread_idx() lane_id = tidx % 32 warp_idx = cute.arch.warp_idx() @@ -330,7 +330,7 @@ def la_decode_kernel_big_batch_pretranspose( for i in cutlass.range_constexpr(vec_size): r_q[i] = r_q[i] * scale - r_g = cute.exp(-cutlass.Float32(decay_scales[i_h]), fastmath=USE_FAST_MATH) + r_g = cute.exp(-cutlass.Float32(decay_scales[i_hv]), fastmath=USE_FAST_MATH) # =================================================================== # Mainloop: All threads participate @@ -404,8 +404,8 @@ def la_decode_kernel_big_batch_pretranspose( @cute.jit def run_la_decode_kernel_big_batch_pretranspose( - h0_source: cute.Tensor, # [B*H, V, K] - decay_scales: cute.Tensor, # [H] + h0_source: cute.Tensor, # [pool_size*HV, V, K] + decay_scales: cute.Tensor, # [HV] q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -413,13 +413,14 @@ def run_la_decode_kernel_big_batch_pretranspose( h0_indices: cute.Tensor, softmax_scale: cutlass.Constexpr[float], H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], stream: cuda.CUstream, ): - # h0_source: (B*HV, V, K) + # h0_source: (pool_size*HV, V, K) _pool_dim0, v_dim, _k_dim = ( h0_source.layout.shape[0], h0_source.layout.shape[1], @@ -473,13 +474,14 @@ def run_la_decode_kernel_big_batch_pretranspose( B, T, H, + HV, K, V, NUM_WARPS_BIG, TILE_V_BIG, NUM_STAGES_BIG, ).launch( - grid=(B * H, 1, 1), + grid=(B * HV, 1, 1), block=[NUM_THREADS_BIG, 1, 1], smem=smem_bytes, stream=stream, @@ -488,8 +490,8 @@ def run_la_decode_kernel_big_batch_pretranspose( @cute.jit def run_la_decode_kernel_small_batch_pretranspose( - h0_source: cute.Tensor, # [B*H, V, K] - decay_scales: cute.Tensor, # [H] + h0_source: cute.Tensor, # [pool_size*HV, V, K] + decay_scales: cute.Tensor, # [HV] q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -497,13 +499,14 @@ def run_la_decode_kernel_small_batch_pretranspose( h0_indices: cute.Tensor, softmax_scale: cutlass.Constexpr[float], H: cutlass.Constexpr[int], + HV: cutlass.Constexpr[int], B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], stream: cuda.CUstream, ): - # h0_source: (B*H, V, K) + # h0_source: (pool_size*HV, V, K) _pool_dim0, v_dim, _k_dim = ( h0_source.layout.shape[0], h0_source.layout.shape[1], @@ -557,13 +560,14 @@ def run_la_decode_kernel_small_batch_pretranspose( B, T, H, + HV, K, V, NUM_WARPS_SMALL, TILE_V_SMALL, NUM_STAGES_SMALL, ).launch( - grid=(B * H * NUM_BLOCKS_PER_STATE, 1, 1), + grid=(B * HV * NUM_BLOCKS_PER_STATE, 1, 1), block=[NUM_THREADS_SMALL, 1, 1], smem=smem_bytes, stream=stream, @@ -572,18 +576,18 @@ def run_la_decode_kernel_small_batch_pretranspose( @functools.cache def _get_compiled_kernel( - B: int, T: int, H: int, K: int, V: int, pool_dim0: int, softmax_scale: float, use_fast_math: bool = True + B: int, T: int, H: int, HV: int, K: int, V: int, pool_dim0: int, softmax_scale: float, use_fast_math: bool = True ): """Get or create compiled kernel cache.""" return {} def linear_attention_decode( - q: torch.Tensor, # [B, 1, H, HEAD_DIM], same as [B, 1, H, K] - k: torch.Tensor, # [B, 1, H, HEAD_DIM], same as [B, 1, H, K] - v: torch.Tensor, # [B, 1, H, HEAD_DIM], same as [B, 1, H, V] - s: torch.Tensor, # [pool_size, heads, V, K] - out: torch.Tensor, # [B, 1, H, HEAD_DIM] + q: torch.Tensor, # [B, H, HEAD_DIM], same as [B, H, K] + k: torch.Tensor, # [B, H, HEAD_DIM], same as [B, H, K] + v: torch.Tensor, # [B, HV, HEAD_DIM], same as [B, HV, V] + s: torch.Tensor, # [pool_size * HV, V, K] + out: torch.Tensor, # [B, HV, HEAD_DIM] softmax_scale: float, stride_q: int, stride_k: int, @@ -591,7 +595,7 @@ def linear_attention_decode( stride_s: int, stride_o: int, s_offsets: torch.Tensor, # [B] - state pool indices - decay_scales: torch.Tensor, # [H] + decay_scales: torch.Tensor, # [HV] or [H] HEAD_DIM: int, K_SPLIT_DIM: int, V_SPLIT_DIM: int, @@ -603,9 +607,9 @@ def linear_attention_decode( Args: q: Query tensor [B, H, HEAD_DIM] k: Key tensor [B, H, HEAD_DIM] - v: Value tensor [B, H, HEAD_DIM] - s: State pool tensor [pool_size, heads, K*V] - out: Output tensor [k_dim_block, length, heads, HEAD_DIM] + v: Value tensor [B, HV, HEAD_DIM] + s: State pool tensor [pool_size * HV, V, K] in BHVK layout + out: Output tensor [B, HV, HEAD_DIM] softmax_scale: Softmax scale factor stride_q: Stride of q tensor stride_k: Stride of k tensor @@ -613,7 +617,7 @@ def linear_attention_decode( stride_s: Stride of s tensor stride_o: Stride of out tensor s_offsets: State pool indices [B] - decay_scales: Decay scales per head [H] + decay_scales: Decay scales per value head [HV]. A [H] tensor is accepted and expanded. HEAD_DIM: Head dimension K_SPLIT_DIM: K split dimension (must be HEAD_DIM for no split) V_SPLIT_DIM: V split dimension (must be HEAD_DIM for no split) @@ -621,8 +625,27 @@ def linear_attention_decode( Returns: None (modifies out and s in-place) """ + if q.ndim != 3 or q.shape[2] != HEAD_DIM: + raise ValueError(f"q must have shape (B, H, HEAD_DIM), got {tuple(q.shape)}") + if k.shape != q.shape: + raise ValueError(f"k must have the same shape as q, got k={tuple(k.shape)}, q={tuple(q.shape)}") B = q.shape[0] H = q.shape[1] + if v.ndim != 3 or v.shape[0] != B or v.shape[2] != HEAD_DIM: + raise ValueError(f"v must have shape (B, HV, HEAD_DIM), got {tuple(v.shape)}") + HV = v.shape[1] + if out.shape != (B, HV, HEAD_DIM): + raise ValueError(f"out must have shape {(B, HV, HEAD_DIM)}, got {tuple(out.shape)}") + if HV < H or HV % H != 0: + raise ValueError(f"HV ({HV}) must be >= H ({H}) and divisible by H") + if decay_scales.ndim != 1: + raise ValueError(f"decay_scales must be 1D, got {tuple(decay_scales.shape)}") + if decay_scales.shape[0] == H and HV != H: + decay_scales = decay_scales.repeat_interleave(HV // H).contiguous() + elif decay_scales.shape[0] == HV: + decay_scales = decay_scales.contiguous() + else: + raise ValueError(f"decay_scales must have shape ({HV},) or ({H},), got {tuple(decay_scales.shape)}") k_dim_block = HEAD_DIM // K_SPLIT_DIM if k_dim_block > 1: @@ -630,13 +653,13 @@ def linear_attention_decode( # Get compiled kernel (cached) pool_dim0 = s.shape[0] - cache_key = (B, 1, H, HEAD_DIM, HEAD_DIM, pool_dim0, softmax_scale, USE_FAST_MATH) + cache_key = (B, 1, H, HV, HEAD_DIM, HEAD_DIM, pool_dim0, softmax_scale, USE_FAST_MATH) cache = _get_compiled_kernel(*cache_key) h0_source = s # Validate state pool dimensions - assert s.shape[0] % H == 0, f"s.shape[0] must be divisible by H={H}, got {s.shape[0]}" + assert s.shape[0] % HV == 0, f"s.shape[0] must be divisible by HV={HV}, got {s.shape[0]}" # First-time compilation if "compiled" not in cache: stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -674,6 +697,7 @@ def linear_attention_decode( h0_idx_tensor, softmax_scale=softmax_scale, H=H, + HV=HV, B=B, T=1, K=HEAD_DIM, @@ -690,11 +714,11 @@ def linear_attention_decode( def seg_la_d_kernel_cute( - q: torch.Tensor, # [B, 1, heads, HEAD_DIM] - k: torch.Tensor, # [B, 1, heads, HEAD_DIM] - v: torch.Tensor, # [B, 1, heads, HEAD_DIM] - s: torch.Tensor, # [pool_size, heads, K*V] - out: torch.Tensor, # [B, 1, heads, HEAD_DIM] + q: torch.Tensor, # [B, H, HEAD_DIM] + k: torch.Tensor, # [B, H, HEAD_DIM] + v: torch.Tensor, # [B, HV, HEAD_DIM] + s: torch.Tensor, # [pool_size * HV, V, K] + out: torch.Tensor, # [B, HV, HEAD_DIM] softmax_scale: float, stride_q: int, stride_k: int, @@ -702,7 +726,7 @@ def seg_la_d_kernel_cute( stride_s: int, stride_o: int, s_offsets: torch.Tensor, # [B] - state pool indices - decay_scales: torch.Tensor, # [H] + decay_scales: torch.Tensor, # [HV] or [H] HEAD_DIM: int, K_SPLIT_DIM: int, V_SPLIT_DIM: int, diff --git a/cula/ops/lightning/prefill_sm100.py b/cula/ops/lightning/prefill_sm100.py index 8a6b204e..3e516f90 100644 --- a/cula/ops/lightning/prefill_sm100.py +++ b/cula/ops/lightning/prefill_sm100.py @@ -107,7 +107,8 @@ class LinearAttentionChunkwiseDecay: chunk_size: Size of each attention chunk (default: 64) acc_dtype: Accumulator data type for all MMA computations (default: Float32) io_dtype: Input/output data type (default: BFloat16) - H: Number of attention heads + H: Number of Q/K heads + HV: Number of V/O heads. HV > H enables GVA. K: Key head dimension (must be 128) V: Value head dimension (must be 128) scale: Scaling factor for queries @@ -121,6 +122,7 @@ def __init__( has_initial_state: bool = False, output_final_state: bool = False, H: int = 64, + HV: int | None = None, K: int = 128, V: int = 128, scale: float = 1.0, @@ -129,6 +131,8 @@ def __init__( use_fast_math: bool = True, ): assert K == 128 and V == 128, f"K and V must both be 128, got K={K}, V={V}" + HV = H if HV is None else HV + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" assert_blackwell() self.use_fast_math = use_fast_math self.chunk_size = chunk_size @@ -145,6 +149,7 @@ def __init__( self.has_initial_state = has_initial_state self.output_final_state = output_final_state self.H = H + self.HV = HV self.K = K self.V = V self.D = K # Internal shorthand: K == V == D @@ -349,16 +354,16 @@ def __call__( (zero-copy C-level dlpack). Pass None for initial_state_in / final_state_in when has_initial_state / output_final_state is False. - scale, H, D are compile-time constants stored in self.__init__. + scale, H, HV, D are compile-time constants stored in self.__init__. Args: q_in: Query tensor [B, S, H, D] or [1, T, H, D] for varlen k_in: Key tensor [B, S, H, D] or [1, T, H, D] for varlen - v_in: Value tensor [B, S, H, D] or [1, T, H, D] for varlen - o_in: Output tensor [B, S, H, D] or [1, T, H, D] for varlen - decay_in: Per-head decay tensor [H] (FP32) - initial_state_in: Initial state [B, H, D, D] or state pool [pool, H, D, D] (FP32) - final_state_in: Final state [B, H, D, D] (FP32) or None (varlen uses INPLACE_UPDATE) + v_in: Value tensor [B, S, HV, D] or [1, T, HV, D] for varlen + o_in: Output tensor [B, S, HV, D] or [1, T, HV, D] for varlen + decay_in: Per-value-head decay tensor [HV] (FP32) + initial_state_in: Initial state [B, HV, D, D] or state pool [pool, HV, D, D] (FP32) + final_state_in: Final state [B, HV, D, D] (FP32) or None (varlen uses INPLACE_UPDATE) cu_seqlens_in: [N+1] int32 cumulative sequence lengths (varlen only) initial_state_indices_in: [N] int32 indices into state pool (varlen only) problem_size: (N, T) for varlen or (B, S) dynamic problem dimensions @@ -366,6 +371,7 @@ def __call__( """ B, S = problem_size H = self.H + HV = self.HV D = self.D # Setup attributes @@ -373,16 +379,16 @@ def __call__( self.cta_group = tcgen05.CtaGroup.ONE - # It's ok since torch tensor is row major, hence we've layout=(B,S,H,D):(DHS, DH, D, 1). + # It's ok since torch tensor is row major, hence we've layout=(B,S,H/HV,D):(D*heads, DH, D, 1). # Below are just permutation tricks to ease the later processing. - # For varlen: input is [1, T, H, D] → view as (T, D, H) with stride (D*H, 1, D) - # For non-varlen: input is [B, S, H, D] → view as (S, D, (H,B)) + # For varlen: Q/K are [1,T,H,D] -> (T,D,H); V/O are [1,T,HV,D] -> (D,T,HV) + # For non-varlen: Q/K use (S,D,(H,B)); V/O use (D,S,(HV,B)). if cutlass.const_expr(self.is_varlen): # Varlen: B=N (num_seqs), S=T (total_tokens), no batch stride q_layout = cute.make_layout((S, D, H), stride=(D * H, 1, D)) k_layout = cute.make_layout((S, D, H), stride=(D * H, 1, D)) - v_layout = cute.make_layout((D, S, H), stride=(1, D * H, D)) - o_layout = cute.make_layout((D, S, H), stride=(1, D * H, D)) + v_layout = cute.make_layout((D, S, HV), stride=(1, D * HV, D)) + o_layout = cute.make_layout((D, S, HV), stride=(1, D * HV, D)) else: q_layout = cute.make_layout( (S, D, (H, B)), @@ -393,27 +399,27 @@ def __call__( stride=(D * H, 1, (D, D * H * S)), ) v_layout = cute.make_layout( - (D, S, (H, B)), - stride=(1, D * H, (D, D * H * S)), + (D, S, (HV, B)), + stride=(1, D * HV, (D, D * HV * S)), ) o_layout = cute.make_layout( - (D, S, (H, B)), - stride=(1, D * H, (D, D * H * S)), + (D, S, (HV, B)), + stride=(1, D * HV, (D, D * HV * S)), ) q = cute.make_tensor(q_in.iterator, q_layout) k = cute.make_tensor(k_in.iterator, k_layout) v = cute.make_tensor(v_in.iterator, v_layout) o = cute.make_tensor(o_in.iterator, o_layout) - # Initial state / final state: [B, H, D, D] in BHVK layout (K-contiguous) - # CuTe shape (V, K, (H, B)) with strides (D, 1, ...) for K-contiguous access. + # Initial state / final state: [B, HV, D, D] in BHVK layout (K-contiguous) + # CuTe shape (V, K, (HV, B)) with strides (D, 1, ...) for K-contiguous access. # When has_initial_state / output_final_state is False, None is passed # and the parameter is eliminated at compile time via const_expr guards. - # For varlen: state pool is [pool_size, H, D, D]. We use B (=N) as the + # For varlen: state pool is [pool_size, HV, D, D]. We use B (=N) as the # pool dimension — strides are correct regardless of actual pool_size. fstate_layout = cute.make_layout( - (D, D, (H, B)), - stride=(D, 1, (D * D, D * D * H)), + (D, D, (HV, B)), + stride=(D, 1, (D * D, D * D * HV)), ) if cutlass.const_expr(self.has_initial_state): initial_state = cute.make_tensor(initial_state_in.iterator, fstate_layout) @@ -739,8 +745,8 @@ class SharedStorage: sm_count = _torch.cuda.get_device_properties(0).multi_processor_count self.grid = (sm_count, 1, 1) elif cutlass.const_expr(self.is_varlen): - # Varlen grid: (1, H, N) where B = N = num_sequences - self.grid = (1, H, B) + # Varlen grid: (1, HV, N) where B = N = num_sequences + self.grid = (1, HV, B) else: self.grid = self._compute_grid( o_shape=cute.shape(o), @@ -1001,26 +1007,30 @@ def kernel( B, S = problem_size H = self.H + HV = self.HV D = self.D C = self.chunk_size scale = cutlass.Float32(self.scale) + qk_group_size = HV // H # ===================== Block indices ===================== if cutlass.const_expr(self.is_varlen): if cutlass.const_expr(self.persistent): # 1D grid work decode: persistent (grid=SM_count) - total_work_units = H * B + total_work_units = HV * B num_iters = Int32(0) # not used, while loop controls iteration # Pre-initialize variables reassigned inside persistent loop (CuTe DSL requirement) hidx = Int32(0) + i_h = Int32(0) bidx = Int32(0) bos = Int32(0) eos = Int32(0) seq_len = Int32(0) state_idx = Int32(0) else: - # Non-persistent varlen: 3D grid (1, H, N) + # Non-persistent varlen: 3D grid (1, HV, N) (_, hidx, bidx) = cute.arch.block_idx() + i_h = hidx // qk_group_size bos = cu_seqlens[bidx] eos = cu_seqlens[bidx + 1] seq_len = eos - bos @@ -1028,6 +1038,7 @@ def kernel( num_iters = Int32(1) else: (_, hidx, bidx) = cute.arch.block_idx() + i_h = hidx // qk_group_size seq_len = S state_idx = bidx num_iters = Int32(1) @@ -1051,7 +1062,7 @@ def kernel( block_decay = Float32(0.0) else: # Non-varlen and non-persistent varlen: hidx known at CTA start - decay_tensor = cute.make_tensor(decay, cute.make_layout(H)) + decay_tensor = cute.make_tensor(decay, cute.make_layout(HV)) decay_s = decay_tensor[hidx] # Block-level decay: λ^C for inter-chunk state accumulation block_decay = cute.exp(-decay_s * cutlass.Float32(C), fastmath=self.use_fast_math) @@ -1252,8 +1263,9 @@ def kernel( while should_continue: # --- Work decode (persistent only) --- if cutlass.const_expr(self.is_varlen and self.persistent): - hidx = work_idx % H - bidx = work_idx // H + hidx = work_idx % HV + i_h = hidx // qk_group_size + bidx = work_idx // HV bos = cu_seqlens[bidx] eos = cu_seqlens[bidx + 1] seq_len = eos - bos @@ -1266,7 +1278,7 @@ def kernel( tma_tensor_q_use = cute.domain_offset((bos, 0, 0), tma_tensor_q) # K: (S, D, H) → offset S (mode 0) by bos tma_tensor_k_use = cute.domain_offset((bos, 0, 0), tma_tensor_k) - # V: (D, S, H) → offset S (mode 1) by bos + # V: (D, S, HV) → offset S (mode 1) by bos tma_tensor_v_use = cute.domain_offset((0, bos, 0), tma_tensor_v) else: tma_tensor_q_use = tma_tensor_q @@ -1282,7 +1294,7 @@ def kernel( self.qk_mma_tiler, qk_tiled_mma, operand_mode="A", - hidx=hidx, + hidx=i_h, bidx=bidx, debug_name="Q", ) @@ -1294,7 +1306,7 @@ def kernel( self.qk_mma_tiler, qk_tiled_mma, operand_mode="B", - hidx=hidx, + hidx=i_h, bidx=bidx, debug_name="K", ) @@ -1384,7 +1396,7 @@ def kernel( while should_continue: # --- Work decode (MMA only needs seq_len) --- if cutlass.const_expr(self.is_varlen and self.persistent): - bidx_mma = work_idx // H + bidx_mma = work_idx // HV seq_len = cu_seqlens[bidx_mma + 1] - cu_seqlens[bidx_mma] for chunk_start in cutlass.range(0, seq_len, C, unroll=0): @@ -1699,8 +1711,9 @@ def kernel( while should_continue: # --- Work decode (persistent only) --- if cutlass.const_expr(self.is_varlen and self.persistent): - hidx = work_idx % H - bidx = work_idx // H + hidx = work_idx % HV + i_h = hidx // qk_group_size + bidx = work_idx // HV bos = cu_seqlens[bidx] eos = cu_seqlens[bidx + 1] seq_len = eos - bos @@ -1708,7 +1721,7 @@ def kernel( # Load per-head decay parameter to register (s_h > 0) # For persistent: hidx was decoded above; for non-persistent: hidx from block_idx - decay_tensor_cuda = cute.make_tensor(decay, cute.make_layout(H)) + decay_tensor_cuda = cute.make_tensor(decay, cute.make_layout(HV)) decay_s_cuda = decay_tensor_cuda[hidx] block_decay = cute.exp(-decay_s_cuda * cutlass.Float32(C), fastmath=self.use_fast_math) @@ -2032,8 +2045,8 @@ def kernel( while should_continue: # --- Work decode (persistent only) --- if cutlass.const_expr(self.is_varlen and self.persistent): - hidx = work_idx % H - bidx = work_idx // H + hidx = work_idx % HV + bidx = work_idx // HV bos = cu_seqlens[bidx] eos = cu_seqlens[bidx + 1] seq_len = eos - bos @@ -2081,14 +2094,14 @@ def kernel( tOrO = cute.make_fragment_like(tOsO, self.io_dtype) cute.autovec_copy(tOsO, tOrO) - o_chunk_raw = o_tensor.iterator + (bos + chunk_start) * D * H + hidx * D + o_chunk_raw = o_tensor.iterator + (bos + chunk_start) * D * HV + hidx * D o_chunk_ptr = cute.make_ptr( self.io_dtype, o_chunk_raw.toint(), cute.AddressSpace.gmem, assumed_align=16, ) - o_stride_c = D * H + o_stride_c = D * HV gO_chunk = cute.make_tensor( o_chunk_ptr, cute.make_layout( @@ -2803,11 +2816,22 @@ def make_thread_cooperative_group(size: int): # Compile cache + TVM-FFI API # --------------------------------------------------------------------------- -# Internal cache: maps (has_initial_state, output_final_state, H, D, scale, chunk_size) → compiled_fn +# Internal cache: maps (has_initial_state, output_final_state, H, HV, D, scale, chunk_size) → compiled_fn _kernel_cache: dict = {} -def _compile_single_variant(has_initial_state, output_final_state, H, D, scale, chunk_size): +def _normalize_gva_decay(decay: torch.Tensor, H: int, HV: int) -> torch.Tensor: + """Return a contiguous [HV] decay tensor, accepting [H] decay for grouped cases.""" + if decay.ndim != 1: + raise ValueError(f"decay must be a 1D tensor, got shape {tuple(decay.shape)}") + if decay.shape[0] == HV: + return decay.contiguous() + if decay.shape[0] == H and HV != H: + return decay.repeat_interleave(HV // H).contiguous() + raise ValueError(f"decay must have shape ({HV},) or ({H},), got {tuple(decay.shape)}") + + +def _compile_single_variant(has_initial_state, output_final_state, H, HV, D, scale, chunk_size): """Compile one kernel variant. Returns the compiled TVM-FFI callable. Uses make_fake_compact_tensor and make_fake_stream for compilation with @@ -2822,6 +2846,7 @@ def _compile_single_variant(has_initial_state, output_final_state, H, D, scale, has_initial_state=has_initial_state, output_final_state=output_final_state, H=H, + HV=HV, K=D, V=D, scale=scale, @@ -2831,7 +2856,7 @@ def _compile_single_variant(has_initial_state, output_final_state, H, D, scale, sym_b = cute.sym_int() sym_s = cute.sym_int() - # Q, K, V, O: (B, S, H, D) row-major bf16 + # Q/K: (B, S, H, D); V/O: (B, S, HV, D) row-major bf16 q_fake = make_fake_compact_tensor( cutlass.BFloat16, (sym_b, sym_s, H, D), @@ -2846,29 +2871,29 @@ def _compile_single_variant(has_initial_state, output_final_state, H, D, scale, ) v_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_b, sym_s, H, D), + (sym_b, sym_s, HV, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) o_fake = make_fake_compact_tensor( cutlass.BFloat16, - (sym_b, sym_s, H, D), + (sym_b, sym_s, HV, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) - # decay: (H,) float32 + # decay: (HV,) float32 decay_fake = make_fake_compact_tensor( cutlass.Float32, - (H,), + (HV,), assumed_align=128, ) - # initial_state / final_state: (B, H, D, D) float32 or None + # initial_state / final_state: (B, HV, D, D) float32 or None h0_fake = ( make_fake_compact_tensor( cutlass.Float32, - (sym_b, H, D, D), + (sym_b, HV, D, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) @@ -2878,7 +2903,7 @@ def _compile_single_variant(has_initial_state, output_final_state, H, D, scale, ht_fake = ( make_fake_compact_tensor( cutlass.Float32, - (sym_b, H, D, D), + (sym_b, HV, D, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) @@ -2929,7 +2954,7 @@ def _compile_single_variant(has_initial_state, output_final_state, H, D, scale, return compiled_fn -def _get_compiled_kernel(has_initial_state, output_final_state, H, D, scale, chunk_size): +def _get_compiled_kernel(has_initial_state, output_final_state, H, HV, D, scale, chunk_size): """Get a compiled kernel with on-demand (lazy) compilation. Each variant is compiled exactly once and cached. Compilation is deferred @@ -2938,14 +2963,15 @@ def _get_compiled_kernel(has_initial_state, output_final_state, H, D, scale, chu where a subsequent cute.compile can invalidate previously compiled but not-yet-executed functions. - Cache key: (has_initial_state, output_final_state, H, D, scale, chunk_size, USE_FAST_MATH) + Cache key: (has_initial_state, output_final_state, H, HV, D, scale, chunk_size, USE_FAST_MATH) """ - key = (has_initial_state, output_final_state, H, D, scale, chunk_size, USE_FAST_MATH) + key = (has_initial_state, output_final_state, H, HV, D, scale, chunk_size, USE_FAST_MATH) if key not in _kernel_cache: _kernel_cache[key] = _compile_single_variant( has_initial_state, output_final_state, H, + HV, D, scale, chunk_size, @@ -2971,23 +2997,34 @@ def lightning_attn_fwd( sym_int() is used for B and S so a single compilation handles all batch-size / sequence-length combinations. - Cache key: (has_initial_state, output_final_state, H, D, scale, chunk_size) + Cache key: (has_initial_state, output_final_state, H, HV, D, scale, chunk_size) Args: Q: (B, S, H, D) bf16 query K: (B, S, H, D) bf16 key - V: (B, S, H, D) bf16 value - decay: (H,) f32 per-head decay coefficients + V: (B, S, HV, D) bf16 value. HV > H enables GVA. + decay: (HV,) f32 per-value-head decay coefficients; (H,) is accepted and expanded scale: attention scale factor (default: 1.0) - initial_state: (B, H, D, D) f32 initial state in BHVK layout, or None + initial_state: (B, HV, D, D) f32 initial state in BHVK layout, or None output_final_state: whether to output final state chunk_size: chunk size (default: 64) Returns: - (O, ht): output tensor (B,S,H,D) bf16, final state (B,H,D,D) f32 in BHVK layout or None + (O, ht): output tensor (B,S,HV,D) bf16, final state (B,HV,D,D) f32 in BHVK layout or None """ B, S, H, D = Q.shape - O = torch.zeros_like(Q) + if K.shape != Q.shape: + raise ValueError(f"K must have the same shape as Q, got K={tuple(K.shape)}, Q={tuple(Q.shape)}") + if V.ndim != 4 or V.shape[0] != B or V.shape[1] != S or V.shape[3] != D: + raise ValueError(f"V must have shape (B, S, HV, D), got {tuple(V.shape)}") + HV = V.shape[2] + if HV < H or HV % H != 0: + raise ValueError(f"HV ({HV}) must be >= H ({H}) and divisible by H") + decay = _normalize_gva_decay(decay, H, HV) + if initial_state is not None and initial_state.shape != (B, HV, D, D): + raise ValueError(f"initial_state must have shape {(B, HV, D, D)}, got {tuple(initial_state.shape)}") + + O = torch.zeros_like(V) has_initial_state = initial_state is not None @@ -2995,13 +3032,14 @@ def lightning_attn_fwd( has_initial_state, output_final_state, H, + HV, D, scale, chunk_size, ) if output_final_state: - ht = torch.zeros(B, H, D, D, dtype=torch.float32, device=Q.device) + ht = torch.zeros(B, HV, D, D, dtype=torch.float32, device=Q.device) else: ht = None @@ -3036,7 +3074,7 @@ def lightning_attn_fwd( _varlen_kernel_cache: dict = {} -def _compile_single_variant_varlen(H, D, scale, chunk_size, persistent=True): +def _compile_single_variant_varlen(H, HV, D, scale, chunk_size, persistent=True): """Compile one varlen kernel variant. Returns the compiled TVM-FFI callable. Varlen kernel always has initial state and output_final_state (INPLACE_UPDATE). @@ -3048,6 +3086,7 @@ def _compile_single_variant_varlen(H, D, scale, chunk_size, persistent=True): has_initial_state=True, output_final_state=True, H=H, + HV=HV, K=D, V=D, scale=scale, @@ -3059,7 +3098,7 @@ def _compile_single_variant_varlen(H, D, scale, chunk_size, persistent=True): sym_n = cute.sym_int() # N: number of sequences sym_t = cute.sym_int() # T: total packed tokens - # Q, K, V, O: [1, T, H, D] row-major bf16 + # Q/K: [1, T, H, D]; V/O: [1, T, HV, D] row-major bf16 # For varlen, B=1 in the physical tensor but we view as (T, D, H) q_fake = make_fake_compact_tensor( cutlass.BFloat16, @@ -3075,29 +3114,29 @@ def _compile_single_variant_varlen(H, D, scale, chunk_size, persistent=True): ) v_fake = make_fake_compact_tensor( cutlass.BFloat16, - (1, sym_t, H, D), + (1, sym_t, HV, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) o_fake = make_fake_compact_tensor( cutlass.BFloat16, - (1, sym_t, H, D), + (1, sym_t, HV, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) - # decay: (H,) float32 + # decay: (HV,) float32 decay_fake = make_fake_compact_tensor( cutlass.Float32, - (H,), + (HV,), assumed_align=128, ) - # State pool: [pool_size, H, D, D] float32 — always present for varlen + # State pool: [pool_size, HV, D, D] float32 — always present for varlen # Use sym_n as pool dimension (actual pool may be larger, strides are correct) h0_fake = make_fake_compact_tensor( cutlass.Float32, - (sym_n, H, D, D), + (sym_n, HV, D, D), stride_order=(3, 2, 1, 0), assumed_align=128, ) @@ -3149,15 +3188,16 @@ def _compile_single_variant_varlen(H, D, scale, chunk_size, persistent=True): return compiled_fn -def _get_compiled_kernel_varlen(H, D, scale, chunk_size, persistent=True): +def _get_compiled_kernel_varlen(H, HV, D, scale, chunk_size, persistent=True): """Get a compiled varlen kernel with on-demand compilation. - Cache key: (H, D, scale, chunk_size, persistent, USE_FAST_MATH) + Cache key: (H, HV, D, scale, chunk_size, persistent, USE_FAST_MATH) """ - key = (H, D, scale, chunk_size, persistent, USE_FAST_MATH) + key = (H, HV, D, scale, chunk_size, persistent, USE_FAST_MATH) if key not in _varlen_kernel_cache: _varlen_kernel_cache[key] = _compile_single_variant_varlen( H, + HV, D, scale, chunk_size, @@ -3191,11 +3231,11 @@ def lightning_attn_fwd_varlen( Args: Q: (1, T, H, D) bf16 query — packed tokens from all sequences K: (1, T, H, D) bf16 key - V: (1, T, H, D) bf16 value - decay: (H,) f32 per-head decay coefficients + V: (1, T, HV, D) bf16 value. HV > H enables GVA. + decay: (HV,) f32 per-value-head decay coefficients; (H,) is accepted and expanded cu_seqlens: (N+1,) int32 cumulative sequence lengths scale: attention scale factor (default: 1.0) - state_pool: (pool_size, H, D, D) f32 state pool in BHVK layout, or None + state_pool: (pool_size, HV, D, D) f32 state pool in BHVK layout, or None If None, a zero state pool is allocated with pool_size=N. States are updated in-place (INPLACE_UPDATE). initial_state_indices: (N,) int32 indices into state_pool per sequence. @@ -3203,15 +3243,25 @@ def lightning_attn_fwd_varlen( chunk_size: chunk size (default: 64) Returns: - (O, state_pool): output tensor (1,T,H,D) bf16, updated state pool (pool_size,H,D,D) f32 + (O, state_pool): output tensor (1,T,HV,D) bf16, updated state pool (pool_size,HV,D,D) f32 """ _, T, H, D = Q.shape + if K.shape != Q.shape: + raise ValueError(f"K must have the same shape as Q, got K={tuple(K.shape)}, Q={tuple(Q.shape)}") + if V.ndim != 4 or V.shape[0] != 1 or V.shape[1] != T or V.shape[3] != D: + raise ValueError(f"V must have shape (1, T, HV, D), got {tuple(V.shape)}") + HV = V.shape[2] + if HV < H or HV % H != 0: + raise ValueError(f"HV ({HV}) must be >= H ({H}) and divisible by H") + decay = _normalize_gva_decay(decay, H, HV) N = cu_seqlens.shape[0] - 1 - O = torch.zeros_like(Q) + O = torch.zeros_like(V) # Allocate state pool if not provided if state_pool is None: - state_pool = torch.zeros(N, H, D, D, dtype=torch.float32, device=Q.device) + state_pool = torch.zeros(N, HV, D, D, dtype=torch.float32, device=Q.device) + elif state_pool.ndim != 4 or state_pool.shape[1:] != (HV, D, D): + raise ValueError(f"state_pool must have shape (pool_size, {HV}, {D}, {D}), got {tuple(state_pool.shape)}") # Default indices: identity mapping if initial_state_indices is None: @@ -3221,7 +3271,7 @@ def lightning_attn_fwd_varlen( cu_seqlens = cu_seqlens.to(torch.int32) initial_state_indices = initial_state_indices.to(torch.int32) - compiled_fn = _get_compiled_kernel_varlen(H, D, scale, chunk_size, persistent=persistent) + compiled_fn = _get_compiled_kernel_varlen(H, HV, D, scale, chunk_size, persistent=persistent) # Workspace for persistent kernel atomic counter (zeroed before each call) workspace = torch.zeros(1, dtype=torch.int32, device=Q.device) @@ -3253,6 +3303,7 @@ def main(): parser.add_argument("--batch_size", type=int, default=2, help="Batch size") parser.add_argument("--seq_len", type=int, default=4096, help="Sequence length") parser.add_argument("--num_heads", type=int, default=64, help="Number of heads") + parser.add_argument("--num_v_heads", type=int, default=None, help="Number of value heads (default: num_heads)") parser.add_argument("--head_dim", type=int, default=128, help="Head dimension") parser.add_argument("--chunk_size", type=int, default=64, help="Chunk size") parser.add_argument("--decay", type=float, default=0.95, help="Decay factor") @@ -3267,6 +3318,7 @@ def main(): print(f" Batch size: {args.batch_size}") print(f" Sequence length: {args.seq_len}") print(f" Number of heads: {args.num_heads}") + print(f" Number of value heads: {args.num_v_heads or args.num_heads}") print(f" Head dimension: {args.head_dim}") print(f" Chunk size: {args.chunk_size}") print(f" Decay factor: {args.decay}") @@ -3281,14 +3333,15 @@ def main(): # Create inputs B, S, H, D = args.batch_size, args.seq_len, args.num_heads, args.head_dim + HV = args.num_v_heads if args.num_v_heads is not None else H # Input tensors in format [B, S, H, D] Q = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) K = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) - V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) + V = torch.randn(B, S, HV, D, device="cuda", dtype=torch.bfloat16) - # Per-head decay coefficients [H] - decay = torch.full((H,), args.decay, device="cuda", dtype=torch.float32) + # Per-value-head decay coefficients [HV] + decay = torch.full((HV,), args.decay, device="cuda", dtype=torch.float32) scale = 1.0 / (D**0.5) @@ -3307,7 +3360,7 @@ def main(): compilation_time = time.time() - start_time print(f"Compilation + first run time: {compilation_time:.4f} seconds") - print(f"B, S, H, D: {(B, S, H, D)}") + print(f"B, S, H, HV, D: {(B, S, H, HV, D)}") # Warmup (uses cached kernel — no recompilation) for _ in range(args.warmup_iterations): diff --git a/tests/test_la_decode.py b/tests/test_la_decode.py index b8708336..39ad3d67 100644 --- a/tests/test_la_decode.py +++ b/tests/test_la_decode.py @@ -48,45 +48,59 @@ def torch_la_decode_ref(q, k, v, state, decay_scales, scale): Pure PyTorch reference for single-token linear attention decode. Args: - q, k, v: [B, H, D] bf16 - state: [B, H, D, D] fp32 (K x V layout) - decay_scales: [H] fp32 (positive values; kernel does exp(-decay)) + q, k: [B, H, D] bf16 + v: [B, HV, D] bf16 + state: [B, HV, D, D] fp32 (K x V layout) + decay_scales: [HV] or [H] fp32 (positive values; kernel does exp(-decay)) scale: float Returns: - o: [B, H, D] bf16 - state_new: [B, H, D, D] fp32 + o: [B, HV, D] bf16 + state_new: [B, HV, D, D] fp32 """ B, H, D = q.shape + HV = v.shape[1] + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + group_size = HV // H + if group_size > 1: + q = q.repeat_interleave(group_size, dim=1) + k = k.repeat_interleave(group_size, dim=1) + if decay_scales.shape[0] == H and HV != H: + decay_scales = decay_scales.repeat_interleave(group_size) + q_f = q.float() * scale k_f = k.float() v_f = v.float() - decay = torch.exp(-decay_scales).view(1, H, 1, 1) # [1, H, 1, 1] - state_new = state * decay + k_f.unsqueeze(-1) * v_f.unsqueeze(-2) # [B,H,D,D] - o = torch.einsum("bhk,bhkv->bhv", q_f, state_new) # [B,H,D] + decay = torch.exp(-decay_scales).view(1, HV, 1, 1) # [1, HV, 1, 1] + state_new = state * decay + k_f.unsqueeze(-1) * v_f.unsqueeze(-2) # [B,HV,D,D] + o = torch.einsum("bhk,bhkv->bhv", q_f, state_new) # [B,HV,D] return o.to(torch.bfloat16), state_new # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -def make_inputs(B, H, D, device="cuda", seed=42): +def make_inputs(B, H, D, HV=None, device="cuda", seed=42): torch.manual_seed(seed) + HV = H if HV is None else HV q = torch.randn(B, H, D, device=device, dtype=torch.bfloat16) k = torch.randn(B, H, D, device=device, dtype=torch.bfloat16) - v = torch.randn(B, H, D, device=device, dtype=torch.bfloat16) - state = torch.randn(B, H, D, D, device=device, dtype=torch.float32) * 0.01 + v = torch.randn(B, HV, D, device=device, dtype=torch.bfloat16) + state = torch.randn(B, HV, D, D, device=device, dtype=torch.float32) * 0.01 return q, k, v, state def run_la_decode(q, k, v, state_4d, decay_scales, scale): """Run la_decode with proper state layout conversion.""" - B, H, D, _ = state_4d.shape - # la_decode kernel expects BHVK layout: [B*H, V, K] - # Reference/test state is BHKV: [B, H, K, V] → transpose to BHVK - state_cute = state_4d.clone().transpose(-1, -2).contiguous().reshape(B * H, D, D) - out = torch.zeros(B, H, D, device=q.device, dtype=torch.bfloat16) + B, H, D = q.shape + HV = v.shape[1] + assert state_4d.shape == (B, HV, D, D) + + # la_decode kernel expects BHVK layout: [B*HV, V, K] + # Reference/test state is BHKV: [B, HV, K, V] -> transpose to BHVK + state_cute = state_4d.clone().transpose(-1, -2).contiguous().reshape(B * HV, D, D) + out = torch.zeros(B, HV, D, device=q.device, dtype=torch.bfloat16) s_offsets = torch.arange(B, device=q.device, dtype=torch.int32) linear_attention_decode( @@ -108,7 +122,7 @@ def run_la_decode(q, k, v, state_4d, decay_scales, scale): V_SPLIT_DIM=D, ) # Convert output state back from BHVK to BHKV for comparison - state_out = state_cute.reshape(B, H, D, D).transpose(-1, -2).contiguous() + state_out = state_cute.reshape(B, HV, D, D).transpose(-1, -2).contiguous() return out, state_out @@ -158,6 +172,27 @@ def test_different_heads(H): assert state_rmse / (state_max + 1e-8) < 0.001, f"H={H}: state mismatch" +@pytest.mark.parametrize("B", [2, 33]) +@pytest.mark.parametrize("decay_head_space", ["qk", "value"]) +def test_gva_output_vs_torch_ref(B, decay_head_space): + H, HV, D = 8, 16, 128 + scale = D**-0.5 + decay_hv = 0.5 * torch.arange(HV, device="cuda", dtype=torch.float32) / HV + decay_scales = decay_hv[:: HV // H] if decay_head_space == "qk" else decay_hv + + q, k, v, state = make_inputs(B, H, D, HV=HV) + o_ref, state_ref = torch_la_decode_ref(q, k, v, state, decay_scales, scale) + o_cute, state_cute = run_la_decode(q, k, v, state, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((o_cute.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + assert rmse / (max_ref + 1e-8) < 0.01, f"GVA B={B}: output mismatch" + + state_rmse = torch.sqrt(torch.mean((state_cute - state_ref) ** 2)).item() + state_max = torch.abs(state_ref).max().item() + assert state_rmse / (state_max + 1e-8) < 0.001, f"GVA B={B}: state mismatch" + + def test_zero_decay(): """With decay=0, state_new = state_old + k⊗v (no decay applied).""" B, H, D = 2, 32, 128 @@ -229,21 +264,55 @@ def test_vs_fla(B): # --------------------------------------------------------------------------- +@pytest.mark.skipif(not HAS_FLA, reason="fla not available") +@pytest.mark.parametrize("B", [2, 33]) +def test_gva_vs_fla(B): + H, HV, D = 8, 16, 128 + scale = D**-0.5 + g_gamma = -(8 / HV * 0.5) * torch.arange(HV, device="cuda", dtype=torch.float32) + decay_scales = -g_gamma + group_size = HV // H + + q, k, v, state = make_inputs(B, H, D, HV=HV) + + q_4d = q.repeat_interleave(group_size, dim=1).unsqueeze(1) + k_4d = k.repeat_interleave(group_size, dim=1).unsqueeze(1) + v_4d = v.unsqueeze(1) + with torch.no_grad(): + o_fla, _ = fused_recurrent_fwd( + q_4d, + k_4d, + v_4d, + g_gamma=g_gamma, + scale=scale, + initial_state=state.clone(), + output_final_state=True, + ) + o_fla = o_fla.squeeze(1).to(torch.bfloat16) + + o_cute, _ = run_la_decode(q, k, v, state, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((o_cute.float() - o_fla.float()) ** 2)).item() + max_ref = torch.abs(o_fla.float()).max().item() + assert rmse / (max_ref + 1e-8) < 0.005, f"GVA B={B}: vs fla mismatch, rel_rmse={rmse / (max_ref + 1e-8):.6f}" + + # End-to-End Prefill -> Decode Test # --------------------------------------------------------------------------- @pytest.mark.sm100_only -def test_prefill_decode_e2e(): +@pytest.mark.parametrize("H, HV", [(8, 8), (4, 8)]) +def test_prefill_decode_e2e(H, HV): """Verify prefill output state passes directly into decode without transpose.""" from cula.ops.lightning.prefill_sm100 import lightning_attn_fwd - B, S, H, D = 2, 64, 8, 128 + B, S, D = 2, 64, 128 scale = D**-0.5 - decay_scales = 0.5 * torch.arange(H, device="cuda", dtype=torch.float32) / H + decay_scales = 0.5 * torch.arange(HV, device="cuda", dtype=torch.float32) / HV # Dummy prefill tokens q_pre = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) k_pre = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) - v_pre = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) + v_pre = torch.randn(B, S, HV, D, device="cuda", dtype=torch.bfloat16) # 1. Run Prefill (Generates BHVK ht) _, ht = lightning_attn_fwd(q_pre, k_pre, v_pre, decay_scales, scale=scale, output_final_state=True) @@ -254,7 +323,7 @@ def test_prefill_decode_e2e(): # Dummy decode tokens q_dec = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) k_dec = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) - v_dec = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v_dec = torch.randn(B, HV, D, device="cuda", dtype=torch.bfloat16) # 2. Run Decode (run_la_decode handles BHKV→BHVK internally) out_dec, state_new = run_la_decode(q_dec, k_dec, v_dec, ht_kv, decay_scales, scale) diff --git a/tests/test_la_decode_pool.py b/tests/test_la_decode_pool.py index d3205328..e65a42b8 100644 --- a/tests/test_la_decode_pool.py +++ b/tests/test_la_decode_pool.py @@ -32,12 +32,21 @@ def torch_la_decode_ref(q, k, v, state, decay_scales, scale): - """Pure PyTorch reference — state is [B, H, K, V] (BHKV).""" + """Pure PyTorch reference; state is [B, HV, K, V] (BHKV).""" B, H, D = q.shape + HV = v.shape[1] + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + group_size = HV // H + if group_size > 1: + q = q.repeat_interleave(group_size, dim=1) + k = k.repeat_interleave(group_size, dim=1) + if decay_scales.shape[0] == H and HV != H: + decay_scales = decay_scales.repeat_interleave(group_size) + q_f = q.float() * scale k_f = k.float() v_f = v.float() - decay = torch.exp(-decay_scales).view(1, H, 1, 1) + decay = torch.exp(-decay_scales).view(1, HV, 1, 1) state_new = state * decay + k_f.unsqueeze(-1) * v_f.unsqueeze(-2) o = torch.einsum("bhk,bhkv->bhv", q_f, state_new) return o.to(torch.bfloat16), state_new @@ -47,15 +56,16 @@ def run_la_decode_with_pool(q, k, v, state_pool_4d, s_offsets, decay_scales, sca """ Run la_decode with a state pool and arbitrary offsets. - state_pool_4d: [pool_size, H, K, V] — the full pool (BHKV layout) + state_pool_4d: [pool_size, HV, K, V] — the full pool (BHKV layout) s_offsets: [B] — which pool slot each batch element uses """ B, H, D = q.shape + HV = v.shape[1] pool_size = state_pool_4d.shape[0] - # la_decode expects BHVK layout: [pool_size*H, V, K] - state_cute = state_pool_4d.clone().transpose(-1, -2).contiguous().reshape(pool_size * H, D, D) - out = torch.zeros(B, H, D, device=q.device, dtype=torch.bfloat16) + # la_decode expects BHVK layout: [pool_size*HV, V, K] + state_cute = state_pool_4d.clone().transpose(-1, -2).contiguous().reshape(pool_size * HV, D, D) + out = torch.zeros(B, HV, D, device=q.device, dtype=torch.bfloat16) linear_attention_decode( q, @@ -76,7 +86,7 @@ def run_la_decode_with_pool(q, k, v, state_pool_4d, s_offsets, decay_scales, sca V_SPLIT_DIM=D, ) - state_out = state_cute.reshape(pool_size, H, D, D).transpose(-1, -2).contiguous() + state_out = state_cute.reshape(pool_size, HV, D, D).transpose(-1, -2).contiguous() return out, state_out @@ -145,6 +155,35 @@ def test_non_identity_offsets(): assert rel_err < 0.01, f"Non-identity offsets {offsets}: rel_err={rel_err:.6f}" +def test_gva_non_identity_offsets(): + """GVA with offsets: q/k use H heads while v/state/out use HV heads.""" + B = 4 + POOL_SIZE = 6 + H, HV, D = 4, 8, 128 + scale = D**-0.5 + decay_scales = 0.3 * torch.arange(HV, device="cuda", dtype=torch.float32) / HV + + torch.manual_seed(42) + q = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, H, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, HV, D, device="cuda", dtype=torch.bfloat16) + state_pool = torch.randn(POOL_SIZE, HV, D, D, device="cuda", dtype=torch.float32) * 0.1 + + offsets = [2, 0, 5, 1] + s_offsets = torch.tensor(offsets, device="cuda", dtype=torch.int32) + + out, _ = run_la_decode_with_pool(q, k, v, state_pool, s_offsets, decay_scales, scale) + + state_selected = state_pool[s_offsets.long()] + o_ref, _ = torch_la_decode_ref(q, k, v, state_selected, decay_scales, scale) + + rmse = torch.sqrt(torch.mean((out.float() - o_ref.float()) ** 2)).item() + max_ref = torch.abs(o_ref.float()).max().item() + rel_err = rmse / (max_ref + 1e-8) + + assert rel_err < 0.01, f"GVA non-identity offsets {offsets}: rel_err={rel_err:.6f}" + + # --------------------------------------------------------------------------- # Test 3: Reversed offsets (another non-identity pattern) # --------------------------------------------------------------------------- diff --git a/tests/test_lightning_attn.py b/tests/test_lightning_attn.py index 46c59b99..547c599b 100644 --- a/tests/test_lightning_attn.py +++ b/tests/test_lightning_attn.py @@ -66,16 +66,16 @@ def run_cute_kernel( Uses TVM-FFI compile cache: first call per config compiles, subsequent reuse. Args: - Q, K, V: (B, S, H, D) bfloat16 tensors on CUDA - decay: (H,) float32 per-head decay parameter s (s > 0) + Q, K: (B, S, H, D), V: (B, S, HV, D) bfloat16 tensors on CUDA + decay: (HV,) or (H,) float32 per-head decay parameter s (s > 0) scale: attention scale factor chunk_size: chunk size C - initial_state: (B, H, D, D) float32 or None + initial_state: (B, HV, D, D) float32 or None output_final_state: whether to allocate and return final state Returns: - O: (B, S, H, D) bfloat16 output - ht: (B, H, D, D) float32 final state (or None) + O: (B, S, HV, D) bfloat16 output + ht: (B, HV, D, D) float32 final state (or None) """ O, ht = lightning_attn_fwd( Q, @@ -105,15 +105,15 @@ def run_cute_kernel_varlen( """Run the CuTeDSL varlen kernel. Args: - Q, K, V: (1, T, H, D) bfloat16 — packed sequences - decay: (H,) float32 + Q, K: (1, T, H, D), V: (1, T, HV, D) bfloat16 — packed sequences + decay: (HV,) or (H,) float32 cu_seqlens: (N+1,) int32 - state_pool: (pool_size, H, D, D) float32 or None + state_pool: (pool_size, HV, D, D) float32 or None initial_state_indices: (N,) int32 or None Returns: - O: (1, T, H, D) bfloat16 - state_pool: (pool_size, H, D, D) float32 + O: (1, T, HV, D) bfloat16 + state_pool: (pool_size, HV, D, D) float32 """ O, sp = lightning_attn_fwd_varlen( Q, @@ -135,30 +135,52 @@ def run_cute_kernel_varlen( # --------------------------------------------------------------------------- +def _expand_qk_to_value_heads(Q, K, V): + """Expand Q/K from H heads to HV heads for GVA references.""" + H = Q.shape[2] + HV = V.shape[2] + assert HV >= H and HV % H == 0, f"HV ({HV}) must be >= H ({H}) and divisible by H" + if HV == H: + return Q, K + group_size = HV // H + return Q.repeat_interleave(group_size, dim=2), K.repeat_interleave(group_size, dim=2) + + +def _normalize_decay_to_value_heads(decay, H, HV): + if decay.shape[0] == HV: + return decay + if decay.shape[0] == H and HV != H: + return decay.repeat_interleave(HV // H) + raise ValueError(f"decay must have shape ({HV},) or ({H},), got {tuple(decay.shape)}") + + def pytorch_reference(Q, K, V, decay, chunk_size=64, scale=1.0, initial_state=None, output_final_state=False): """PyTorch reference for chunkwise linear attention with exponential decay. Args: - Q, K, V: (B, T, H, D) — any dtype, computed in float32 - decay: (H,) float32 per-head s (s >= 0) + Q, K: (B, T, H, D), V: (B, T, HV, D) — any dtype, computed in float32 + decay: (HV,) or (H,) float32 per-head s (s >= 0) chunk_size: C scale: scalar multiplier applied to final output - initial_state: (B, H, D, D) float32 or None + initial_state: (B, HV, D, D) float32 or None output_final_state: bool Returns: - O: (B, T, H, D) float32 - final_state: (B, H, D, D) float32 or None + O: (B, T, HV, D) float32 + final_state: (B, HV, D, D) float32 or None """ B, T, H, D = Q.shape + HV = V.shape[2] C = chunk_size + Q, K = _expand_qk_to_value_heads(Q, K, V) + decay = _normalize_decay_to_value_heads(decay, H, HV) Q, K, V = Q.float(), K.float(), V.float() - O = torch.zeros(B, T, H, D, device=Q.device, dtype=torch.float32) + O = torch.zeros(B, T, HV, D, device=Q.device, dtype=torch.float32) state = ( initial_state.clone().float() if initial_state is not None - else torch.zeros(B, H, D, D, device=Q.device, dtype=torch.float32) + else torch.zeros(B, HV, D, D, device=Q.device, dtype=torch.float32) ) num_chunks = (T + C - 1) // C @@ -176,21 +198,21 @@ def pytorch_reference(Q, K, V, decay, chunk_size=64, scale=1.0, initial_state=No pos_k = torch.arange(cl, device=Q.device).view(1, cl) dist = pos_q - pos_k # (cl, cl) - s = decay.view(1, H, 1, 1) + s = decay.view(1, HV, 1, 1) mask = torch.exp(-s * dist.unsqueeze(0).unsqueeze(0).float()) mask = mask * (pos_q >= pos_k).unsqueeze(0).unsqueeze(0).float() O_intra = torch.einsum("bhts,bshd->bthd", QK * mask, Vc) # --- inter-chunk: Q @ state with per-position decay --- pos_in = torch.arange(cl, device=Q.device).float() - per_pos = torch.exp(-decay.view(1, 1, H, 1) * (pos_in.view(1, -1, 1, 1) + 1.0)) + per_pos = torch.exp(-decay.view(1, 1, HV, 1) * (pos_in.view(1, -1, 1, 1) + 1.0)) O_inter = torch.einsum("bthd,bhde->bthe", Qc, state) * per_pos O[:, cs:ce] = (O_intra + O_inter) * scale # --- state update --- - block_decay = torch.exp(-decay.view(1, H, 1, 1) * C) - pos_w = torch.exp(-decay.view(1, 1, H, 1) * (C - 1 - pos_in.view(1, -1, 1, 1))) + block_decay = torch.exp(-decay.view(1, HV, 1, 1) * C) + pos_w = torch.exp(-decay.view(1, 1, HV, 1) * (C - 1 - pos_in.view(1, -1, 1, 1))) state = state * block_decay + torch.einsum("bthd,bthe->bhde", Kc * pos_w, Vc) return O, (state if output_final_state else None) @@ -270,16 +292,17 @@ def test_different_decay_values(): return False -def test_against_reference(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): +def test_against_reference(B=1, S=128, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): """Compare against PyTorch reference (exact match).""" + HV = H if HV is None else HV if verbose: - print(f"\nRef: B={B}, S={S}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nRef: B={B}, S={S}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) Q = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + V = torch.randn(B, S, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) O_ref, _ = pytorch_reference(Q, K, V, decay, chunk_size=C) O_ref_bf16 = O_ref.to(torch.bfloat16) @@ -291,22 +314,23 @@ def test_against_reference(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e- return passed -def test_initial_and_final_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): +def test_initial_and_final_state(B=1, S=128, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): """Test h0/ht against PyTorch reference. NOTE: This test is placed BEFORE FLA tests so that the (has_initial_state=True, output_final_state=True) kernel variant is compiled before any Triton/FLA code runs. Running Triton corrupts state needed by cute.compile. """ + HV = H if HV is None else HV if verbose: - print(f"\nh0/ht: B={B}, S={S}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nh0/ht: B={B}, S={S}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) Q = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) - h0 = torch.randn(B, H, D, D, device="cuda", dtype=torch.float32) * 0.01 + V = torch.randn(B, S, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) + h0 = torch.randn(B, HV, D, D, device="cuda", dtype=torch.float32) * 0.01 h0_vk = h0.transpose(-1, -2).contiguous() # BHVK for CuTe kernel O_ref, ht_ref = pytorch_reference( @@ -340,7 +364,7 @@ def test_initial_and_final_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, at return passed -def test_against_fla(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): +def test_against_fla(B=1, S=128, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): """Compare against FLA chunk_simple_gla using g_gamma = -s. FLA's g_gamma is the per-head log-decay (negative). Our decay parameter s @@ -350,20 +374,22 @@ def test_against_fla(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rto print("\n ⊘ SKIPPED: fla library not available") return True + HV = H if HV is None else HV if verbose: - print(f"\nFLA: B={B}, S={S}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nFLA: B={B}, S={S}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) Q = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 + V = torch.randn(B, S, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + Q_fla, K_fla = _expand_qk_to_value_heads(Q, K, V) # Our decay s -> FLA g_gamma = -s - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) g_gamma = -decay # FLA reference (scale=1.0 to match our kernel) - O_fla, _ = chunk_simple_gla(Q, K, V, g_gamma=g_gamma, scale=1.0) + O_fla, _ = chunk_simple_gla(Q_fla, K_fla, V, g_gamma=g_gamma, scale=1.0) # Our kernel O_cute, _ = run_cute_kernel(Q, K, V, decay, scale=1.0, chunk_size=C) @@ -383,29 +409,31 @@ def test_against_fla(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rto return passed -def test_against_fla_with_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): +def test_against_fla_with_state(B=1, S=128, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True): """Compare h0/ht against FLA chunk_simple_gla.""" if not HAS_FLA: print("\n ⊘ SKIPPED: fla library not available") return True + HV = H if HV is None else HV if verbose: - print(f"\nFLA h0/ht: B={B}, S={S}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nFLA h0/ht: B={B}, S={S}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) Q = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - h0 = torch.randn(B, H, D, D, device="cuda", dtype=torch.float32) * 0.01 + V = torch.randn(B, S, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + Q_fla, K_fla = _expand_qk_to_value_heads(Q, K, V) + h0 = torch.randn(B, HV, D, D, device="cuda", dtype=torch.float32) * 0.01 h0_vk = h0.transpose(-1, -2).contiguous() # BHVK for CuTe kernel - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) g_gamma = -decay # FLA (expects BHKV state) O_fla, ht_fla = chunk_simple_gla( - Q, - K, + Q_fla, + K_fla, V, g_gamma=g_gamma, scale=1.0, @@ -440,16 +468,17 @@ def test_against_fla_with_state(B=1, S=128, H=4, D=128, C=64, decay_val=0.1, ato # =========================================================================== -def test_varlen_single_seq(H=4, S=128, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True) -> bool: +def test_varlen_single_seq(H=4, HV=None, S=128, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True) -> bool: """Varlen with a single sequence vs non-varlen reference.""" + HV = H if HV is None else HV if verbose: - print(f"\nVarlen single: S={S}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nVarlen single: S={S}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) Q = torch.randn(1, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(1, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(1, S, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + V = torch.randn(1, S, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) # Non-varlen reference O_ref, ht_ref = run_cute_kernel(Q, K, V, decay, chunk_size=C, output_final_state=True) @@ -465,12 +494,13 @@ def test_varlen_single_seq(H=4, S=128, D=128, C=64, decay_val=0.1, atol=5e-3, rt return passed -def test_varlen_multi_seq(seq_lens=None, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True) -> bool: +def test_varlen_multi_seq(seq_lens=None, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True) -> bool: """Varlen with multiple packed sequences vs per-sequence non-varlen reference.""" + HV = H if HV is None else HV if seq_lens is None: seq_lens = [128, 64, 192] # all multiples of C if verbose: - print(f"\nVarlen multi: seqs={seq_lens}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nVarlen multi: seqs={seq_lens}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) T = sum(seq_lens) @@ -482,8 +512,8 @@ def test_varlen_multi_seq(seq_lens=None, H=4, D=128, C=64, decay_val=0.1, atol=5 Q = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + V = torch.randn(1, T, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) O_var, sp = run_cute_kernel_varlen(Q, K, V, decay, cu_seqlens, chunk_size=C) @@ -504,12 +534,15 @@ def test_varlen_multi_seq(seq_lens=None, H=4, D=128, C=64, decay_val=0.1, atol=5 return all_pass -def test_varlen_with_initial_state(seq_lens=None, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True) -> bool: +def test_varlen_with_initial_state( + seq_lens=None, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True +) -> bool: """Varlen with initial state from state pool (non-contiguous indices).""" + HV = H if HV is None else HV if seq_lens is None: seq_lens = [128, 64] if verbose: - print(f"\nVarlen h0: seqs={seq_lens}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nVarlen h0: seqs={seq_lens}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) T = sum(seq_lens) @@ -521,12 +554,12 @@ def test_varlen_with_initial_state(seq_lens=None, H=4, D=128, C=64, decay_val=0. Q = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + V = torch.randn(1, T, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) # State pool with 3 slots, use indices [2, 0] — BHVK layout for CuTe pool_size = 3 - state_pool = torch.randn(pool_size, H, D, D, dtype=torch.float32, device="cuda").transpose(-1, -2).contiguous() * 0.01 + state_pool = torch.randn(pool_size, HV, D, D, dtype=torch.float32, device="cuda").transpose(-1, -2).contiguous() * 0.01 indices = torch.tensor([2, 0], dtype=torch.int32, device="cuda") O_var, sp = run_cute_kernel_varlen( @@ -569,13 +602,14 @@ def test_varlen_with_initial_state(seq_lens=None, H=4, D=128, C=64, decay_val=0. def test_varlen_against_pytorch_ref( - seq_lens=None, H=4, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True + seq_lens=None, H=4, HV=None, D=128, C=64, decay_val=0.1, atol=5e-3, rtol=5e-2, verbose=True ) -> bool: """Varlen against the PyTorch reference with initial state.""" + HV = H if HV is None else HV if seq_lens is None: seq_lens = [128, 192] if verbose: - print(f"\nVarlen vs ref: seqs={seq_lens}, H={H}, D={D}, C={C}, decay={decay_val}") + print(f"\nVarlen vs ref: seqs={seq_lens}, H={H}, HV={HV}, D={D}, C={C}, decay={decay_val}") torch.manual_seed(42) T = sum(seq_lens) @@ -588,10 +622,10 @@ def test_varlen_against_pytorch_ref( Q = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 K = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - V = torch.randn(1, T, H, D, device="cuda", dtype=torch.bfloat16) * 0.1 - decay = torch.full((H,), decay_val, device="cuda", dtype=torch.float32) + V = torch.randn(1, T, HV, D, device="cuda", dtype=torch.bfloat16) * 0.1 + decay = torch.full((HV,), decay_val, device="cuda", dtype=torch.float32) - state_pool = torch.randn(N, H, D, D, dtype=torch.float32, device="cuda") * 0.01 + state_pool = torch.randn(N, HV, D, D, dtype=torch.float32, device="cuda") * 0.01 state_pool_vk = state_pool.transpose(-1, -2).contiguous() # BHVK for CuTe O_var, sp = run_cute_kernel_varlen( @@ -675,6 +709,7 @@ def main(): ("Decay 0.2", dict(B=1, S=128, H=4, D=128, C=64, decay_val=0.2)), ("Decay 0.5", dict(B=1, S=128, H=4, D=128, C=64, decay_val=0.5)), ("Batch", dict(B=2, S=128, H=4, D=128, C=64, decay_val=0.1)), + ("GVA H2-HV4", dict(B=1, S=128, H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"Ref {tag}", test_against_reference(**kw, verbose=args.verbose))) @@ -694,6 +729,7 @@ def main(): ("Decay 0.2", dict(B=1, S=128, H=4, D=128, C=64, decay_val=0.2)), ("Decay 0.5", dict(B=1, S=128, H=4, D=128, C=64, decay_val=0.5)), ("Batch", dict(B=2, S=128, H=4, D=128, C=64, decay_val=0.1)), + ("GVA H2-HV4", dict(B=1, S=128, H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"FLA {tag}", test_against_fla(**kw, verbose=args.verbose))) @@ -706,6 +742,7 @@ def main(): ("Small", dict(B=1, S=64, H=2, D=128, C=64, decay_val=0.1)), ("Multi-chunk", dict(B=1, S=256, H=4, D=128, C=64, decay_val=0.1)), ("Batch", dict(B=2, S=128, H=4, D=128, C=64, decay_val=0.2)), + ("GVA H2-HV4", dict(B=1, S=128, H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"h0/ht {tag}", test_initial_and_final_state(**kw, verbose=args.verbose))) @@ -719,6 +756,7 @@ def main(): ("Small", dict(B=1, S=64, H=2, D=128, C=64, decay_val=0.1)), ("Multi-chunk", dict(B=1, S=256, H=4, D=128, C=64, decay_val=0.1)), ("Batch", dict(B=2, S=128, H=4, D=128, C=64, decay_val=0.2)), + ("GVA H2-HV4", dict(B=1, S=128, H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"FLA h0/ht {tag}", test_against_fla_with_state(**kw, verbose=args.verbose))) @@ -731,6 +769,7 @@ def main(): ("Single seq", dict(H=4, S=128, D=128, C=64, decay_val=0.1)), ("Single long", dict(H=4, S=256, D=128, C=64, decay_val=0.1)), ("Decay 0.5", dict(H=4, S=128, D=128, C=64, decay_val=0.5)), + ("GVA single", dict(H=2, HV=4, S=128, D=128, C=64, decay_val=0.1)), ]: results.append((f"Varlen {tag}", test_varlen_single_seq(**kw, verbose=args.verbose))) @@ -738,12 +777,14 @@ def main(): ("Multi 3-seq", dict(seq_lens=[128, 64, 192], H=4, D=128, C=64, decay_val=0.1)), ("Multi 2-seq", dict(seq_lens=[256, 128], H=4, D=128, C=64, decay_val=0.1)), ("Multi decay", dict(seq_lens=[128, 128], H=4, D=128, C=64, decay_val=0.5)), + ("GVA multi", dict(seq_lens=[128, 64], H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"Varlen {tag}", test_varlen_multi_seq(**kw, verbose=args.verbose))) for tag, kw in [ ("h0 indirect", dict(seq_lens=[128, 64], H=4, D=128, C=64, decay_val=0.1)), ("h0 decay 0.2", dict(seq_lens=[128, 64], H=4, D=128, C=64, decay_val=0.2)), + ("GVA h0", dict(seq_lens=[128, 64], H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"Varlen {tag}", test_varlen_with_initial_state(**kw, verbose=args.verbose))) @@ -753,6 +794,7 @@ def main(): for tag, kw in [ ("vs ref 2-seq", dict(seq_lens=[128, 192], H=4, D=128, C=64, decay_val=0.1)), ("vs ref decay", dict(seq_lens=[64, 128], H=4, D=128, C=64, decay_val=0.5)), + ("GVA vs ref", dict(seq_lens=[64, 128], H=2, HV=4, D=128, C=64, decay_val=0.1)), ]: results.append((f"Varlen {tag}", test_varlen_against_pytorch_ref(**kw, verbose=args.verbose))) From 326a307956609984868357d99ba88a26c162d6b8 Mon Sep 17 00:00:00 2001 From: cher <117337477+cherhh@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:14:33 +0800 Subject: [PATCH 33/34] [Fix] Wire cudac bindings (#105) * [Fix] Wire cudac bindings * Avoid pybind object reuse --------- Co-authored-by: cheheng.ch --- csrc/api/kda_sm100.cu | 6 ------ csrc/api/kda_sm90.cu | 5 ----- csrc/api/pybind_sm100.cu | 1 + csrc/api/pybind_sm90.cu | 1 + setup.py | 6 ++++-- 5 files changed, 6 insertions(+), 13 deletions(-) create mode 100644 csrc/api/pybind_sm100.cu create mode 100644 csrc/api/pybind_sm90.cu diff --git a/csrc/api/kda_sm100.cu b/csrc/api/kda_sm100.cu index 020d90ca..0a689ec5 100644 --- a/csrc/api/kda_sm100.cu +++ b/csrc/api/kda_sm100.cu @@ -189,9 +189,3 @@ ChunkKDAFwdRecompWU( kda::sm100::run_kda_fwd_recomp_w_u_sm100(params, at::cuda::getCurrentCUDAStream()); } - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "cuLA SM100/SM103 kernels"; - m.def("chunk_kda_fwd_intra_cuda", &ChunkKDAFwdIntra); - m.def("recompute_w_u_cuda", &ChunkKDAFwdRecompWU); -} \ No newline at end of file diff --git a/csrc/api/kda_sm90.cu b/csrc/api/kda_sm90.cu index d9a4de06..bd7ba5ee 100644 --- a/csrc/api/kda_sm90.cu +++ b/csrc/api/kda_sm90.cu @@ -224,8 +224,3 @@ kda_fwd_prefill( return {output, output_state}; } - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "cuLA SM90 kernels"; - m.def("kda_fwd_prefill", &kda_fwd_prefill); -} diff --git a/csrc/api/pybind_sm100.cu b/csrc/api/pybind_sm100.cu new file mode 100644 index 00000000..22164b17 --- /dev/null +++ b/csrc/api/pybind_sm100.cu @@ -0,0 +1 @@ +#include "pybind.cu" diff --git a/csrc/api/pybind_sm90.cu b/csrc/api/pybind_sm90.cu new file mode 100644 index 00000000..22164b17 --- /dev/null +++ b/csrc/api/pybind_sm90.cu @@ -0,0 +1 @@ +#include "pybind.cu" diff --git a/setup.py b/setup.py index 78c61e5c..a187764a 100644 --- a/setup.py +++ b/setup.py @@ -164,15 +164,16 @@ def get_nvcc_thread_args(): if not DISABLE_SM100 or not DISABLE_SM103: sm100_arch_flags = [] if not DISABLE_SM100: - sm100_arch_flags.extend(["-gencode", "arch=compute_100a,code=sm_100a"]) + sm100_arch_flags.extend(["-gencode", "arch=compute_100a,code=sm_100a", "-DCULA_SM100_ENABLED"]) if not DISABLE_SM103: - sm100_arch_flags.extend(["-gencode", "arch=compute_103a,code=sm_103a"]) + sm100_arch_flags.extend(["-gencode", "arch=compute_103a,code=sm_103a", "-DCULA_SM103_ENABLED"]) ext_modules.append( CUDAExtension( name="cula._cudac_sm100", sources=[ "csrc/api/kda_sm100.cu", + "csrc/api/pybind_sm100.cu", "csrc/kda/sm100/kda_fwd_sm100.cu", ], extra_compile_args={ @@ -195,6 +196,7 @@ def get_nvcc_thread_args(): name="cula._cudac_sm90", sources=[ "csrc/api/kda_sm90.cu", + "csrc/api/pybind_sm90.cu", "csrc/kda/sm90/kda_fwd_sm90.cu", "csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu", ], From 568789c59edead8e5852838b8ef44a424048d74f Mon Sep 17 00:00:00 2001 From: yz262713 Date: Mon, 20 Jul 2026 16:01:28 +0800 Subject: [PATCH 34/34] Add SM90 fused WY-DqKG backward --- README.md | 3 + benchmarks/bench_kda_bwd_wy_dqkg_sm90.py | 597 ++++ cula/kda/chunk_bwd.py | 146 +- cula/kda/chunk_fwd.py | 103 +- cula/kda/chunk_intra.py | 59 +- cula/ops/__init__.py | 5 + cula/ops/chunk_wy_dqkg_sm90.py | 3625 ++++++++++++++++++++++ tests/conftest.py | 15 +- tests/test_chunk_wy_dqkg_sm90.py | 689 ++++ tests/test_kda_chunk_sm90.py | 96 + 10 files changed, 5307 insertions(+), 31 deletions(-) create mode 100644 benchmarks/bench_kda_bwd_wy_dqkg_sm90.py create mode 100644 cula/ops/chunk_wy_dqkg_sm90.py create mode 100644 tests/test_chunk_wy_dqkg_sm90.py create mode 100644 tests/test_kda_chunk_sm90.py diff --git a/README.md b/README.md index 09418811..8b3dbe55 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,9 @@ python benchmarks/generate_benchmark_md.py # Hopper (SM90) python benchmarks/generate_benchmark_hopper_md.py + +# Hopper fused WY-DqKG backward vs FLA v0.5.0 +python benchmarks/bench_kda_bwd_wy_dqkg_sm90.py --mode both --heads 32 ``` ## Tests diff --git a/benchmarks/bench_kda_bwd_wy_dqkg_sm90.py b/benchmarks/bench_kda_bwd_wy_dqkg_sm90.py new file mode 100644 index 00000000..0d4d4e40 --- /dev/null +++ b/benchmarks/bench_kda_bwd_wy_dqkg_sm90.py @@ -0,0 +1,597 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +bench_kda_bwd_wy_dqkg_sm90.py — Benchmark: SM90 CuTe DSL vs FLA Triton + for chunk_kda_bwd_wy_dqkg kernel (Hopper) + +Modes: + - Fixed-length: B=1,2 with various T + - Varlen: variable-length sequences with different distributions + +Usage: + python benchmarks/bench_kda_bwd_wy_dqkg_sm90.py \ + [--mode fixed|varlen|both] [--heads 4 32] + + /usr/local/cuda-12.9/bin/ncu --profile-from-start off --set full \ + -o ncu_reports/ \ + python benchmarks/bench_kda_bwd_wy_dqkg_sm90.py --ncu \ + --mode varlen --heads 32 --total-len 16384 --num-seqs 20 --dist random +""" + +import argparse +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch +from einops import rearrange +from fla.modules.l2norm import l2norm_fwd +from fla.ops.kda.chunk_bwd import chunk_kda_bwd_wy_dqkg_fused as fla_bwd +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices + +from benchmarks.utils import ( + SEED, + build_varlen_configs, + exclusive_cumsum, + gen_random, + gen_skewed, + gen_uniform, + set_seed, +) +from cula.ops.chunk_wy_dqkg_sm90 import chunk_kda_bwd_wy_dqkg_fused as sm90_bwd + +torch.backends.cuda.matmul.allow_tf32 = True + +H_DEFAULT = 32 +K = 128 +V = 128 +BT = 64 +BK = 32 +BV = 64 +MIN_OCC = 2 +DTYPE = torch.bfloat16 +DEVICE = torch.device("cuda") +WARMUP = 25 +N_ITERS = 100 +NCU_MODE = False + +FIXED_CONFIGS = [ + (1, 256), + (1, 512), + (1, 1024), + (1, 2048), + (1, 4096), + (1, 8192), + (1, 16384), + (1, 32768), + (2, 512), + (2, 1024), + (2, 2048), + (2, 4096), + (2, 8192), + (2, 16384), + (2, 32768), +] +VARLEN_NUM_SEQS_LIST = (10, 20) +VARLEN_TOTAL_LENS = (4096, 8192, 16384) +VARLEN_DISTS = ("uniform", "random", "skewed") + + +def benchmark_fixed_configs(): + return list(FIXED_CONFIGS) + + +def benchmark_varlen_configs(): + return build_varlen_configs( + num_seqs_list=VARLEN_NUM_SEQS_LIST, + total_lens=VARLEN_TOTAL_LENS, + dists=VARLEN_DISTS, + ) + + +# ============================================================ +# Runners +# ============================================================ +def prepare_bwd_wy_dqkg_fused_inputs( + B: int, + T: int, + H: int, + K: int, + V: int, + chunk_size: int = BT, + device: torch.device | str = DEVICE, + seed: int = SEED, + cu_seqlens: torch.Tensor | None = None, + dtype: torch.dtype = DTYPE, +) -> dict: + """Prepare inputs for FLA and SM90 WY-DqKG fused backward runners.""" + scale = K**-0.5 + set_seed(seed) + + q = torch.randn(B, T, H, K, dtype=dtype, device=device) + k = torch.randn(B, T, H, K, dtype=dtype, device=device) + v = torch.randn(B, T, H, V, dtype=dtype, device=device) + g_raw = torch.randn(B, T, H, K, dtype=dtype, device=device) + beta = torch.randn(B, T, H, dtype=torch.float32, device=device).sigmoid() + + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + + A_log = torch.randn(H, dtype=torch.float32, device=device) + dt_bias = torch.randn(H * K, dtype=torch.float32, device=device) + + v_new = torch.randn(B, T, H, V, dtype=dtype, device=device) + do = torch.randn(B, T, H, V, dtype=dtype, device=device) + dv = torch.randn(B, T, H, V, dtype=dtype, device=device) + A = torch.randn(B, T, H, chunk_size, dtype=dtype, device=device) * 0.1 + + if cu_seqlens is not None: + cu_seqlens = cu_seqlens.int() + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = chunk_indices.shape[0] + + # Explicit cu_seqlens means the input is represented as one flattened + # varlen stream. Fixed-length B>1 is just uniform varlen here. + if B != 1: + q, k, v, g_raw, beta = map( + lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), + (q, k, v, g_raw, beta), + ) + v_new, do, dv, A = map( + lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), + (v_new, do, dv, A), + ) + + h = torch.randn(1, NT, H, K, V, dtype=dtype, device=device) * 0.01 + dh = torch.randn(1, NT, H, K, V, dtype=dtype, device=device) * 0.01 + else: + NT = (T + chunk_size - 1) // chunk_size + chunk_indices = None + h = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 + dh = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 + + g = kda_gate_chunk_cumsum( + g=g_raw, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=-5.0, + ) + + return dict( + q=q, + k=k, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=A, + h=h, + dh=dh, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + +def run_sm90(inputs: dict): + """Run the SM90 CuTe DSL kernel.""" + return sm90_bwd( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + v_new=inputs["v_new"], + g=inputs["g"], + beta=inputs["beta"], + A=inputs["A"], + h=inputs["h"], + do=inputs["do"], + dh=inputs["dh"], + dv=inputs["dv"], + scale=inputs["scale"], + cu_seqlens=inputs["cu_seqlens"], + chunk_size=BT, + chunk_indices=inputs["chunk_indices"], + bk=BK, + bv=BV, + min_occupancy=MIN_OCC, + ) + + +def run_fla(inputs: dict): + """Run the FLA Triton baseline.""" + return fla_bwd( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + v_new=inputs["v_new"], + g=inputs["g"], + beta=inputs["beta"], + A=inputs["A"], + h=inputs["h"], + do=inputs["do"], + dh=inputs["dh"], + dv=inputs["dv"], + scale=inputs["scale"], + cu_seqlens=inputs["cu_seqlens"], + chunk_size=BT, + chunk_indices=inputs["chunk_indices"], + transpose_state_layout=False, + ) + + +# ============================================================ +# Helpers +# ============================================================ +def time_kernel(fn, warmup=None, n_iters=None): + if warmup is None: + warmup = 1 if NCU_MODE else WARMUP + if n_iters is None: + n_iters = 1 if NCU_MODE else N_ITERS + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(n_iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / n_iters + + +def accuracy_stats(ref, out): + """Compute err_ratio, relative max diff, and mean absolute difference.""" + ref_f = ref.float() + out_f = out.float() + diff = (ref_f - out_f).abs() + err = diff.flatten().pow(2).mean().sqrt().item() + base = ref_f.flatten().pow(2).mean().sqrt().item() + err_ratio = err / (base + 1e-8) + max_diff = diff.max().item() + denom = ref_f.abs().max().item() + rel_max = max_diff / denom if denom > 0 else 0.0 + mean_diff = diff.mean().item() + return err_ratio, rel_max, mean_diff + + +# Both SM90 and FLA return (dq, dk, dv, db, dg, dA). +OUT_MAP = {"dq": 0, "dk": 1, "dv": 2, "db": 3, "dg": 4, "dA": 5} +ACC_KEYS = ["dq", "dk", "dv", "db", "dg", "dA"] + + +def compute_accuracy(sm90_out, fla_out): + """Compute per-output accuracy stats.""" + acc = {} + for name in ACC_KEYS: + s = sm90_out[OUT_MAP[name]] + f = fla_out[OUT_MAP[name]] + if s.shape != f.shape: + f = f.reshape(s.shape) + err_ratio, rel_max, mean_diff = accuracy_stats(f, s) + acc[name] = {"err_ratio": err_ratio, "rel_max": rel_max, "mean_diff": mean_diff} + return acc + + +def make_profile_seq_lens(args): + profile_mode = "varlen" if args.mode == "both" else args.mode + if profile_mode == "fixed": + return [args.total_len] * args.batch, args.batch, args.total_len, profile_mode + if args.dist == "uniform": + seq_lens = gen_uniform(args.num_seqs, args.total_len) + elif args.dist == "random": + seq_lens = gen_random(args.num_seqs, args.total_len, seed=SEED) + elif args.dist == "skewed": + seq_lens = gen_skewed(args.num_seqs, args.total_len) + else: + raise ValueError(f"unknown dist: {args.dist}") + return seq_lens, 1, args.total_len, profile_mode + + +def run_ncu_profile(args): + """Run one SM90-only config bracketed by cudaProfilerStart/Stop.""" + seq_lens, batch, T, profile_mode = make_profile_seq_lens(args) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=DEVICE) + inputs = prepare_bwd_wy_dqkg_fused_inputs( + B=batch, + T=T, + H=args.heads[0], + K=K, + V=V, + chunk_size=BT, + device=DEVICE, + seed=SEED, + cu_seqlens=cu_seqlens, + ) + + print( + f"[NCU profiler] mode={profile_mode} dist={args.dist} H={args.heads[0]} " + f"T={T} seqs={len(seq_lens)} min={min(seq_lens)} max={max(seq_lens)} " + f"BK={BK} BV={BV} OCC={MIN_OCC} warmup={args.profile_warmup} " + f"profile_iters={args.profile_iters}", + flush=True, + ) + + for _ in range(args.profile_warmup): + run_sm90(inputs) + torch.cuda.synchronize() + + torch.cuda.cudart().cudaProfilerStart() + for _ in range(args.profile_iters): + run_sm90(inputs) + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStop() + print("[NCU profiler] done") + + +# ============================================================ +# Fixed-length benchmark +# ============================================================ +def bench_fixed(configs, H: int): + print(f"\n{'=' * 120}") + print(f" Fixed-Length Benchmark: SM90 CuTe DSL vs FLA Triton (H={H}, K={K}, V={V}, BT={BT})") + print(f"{'=' * 120}") + results = [] + + for B, T in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + seq_lens = [T] * B + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=DEVICE) + + inputs = prepare_bwd_wy_dqkg_fused_inputs( + B=B, + T=T, + H=H, + K=K, + V=V, + chunk_size=BT, + device=DEVICE, + seed=SEED, + cu_seqlens=cu_seqlens, + ) + + sm90_out = run_sm90(inputs) + fla_out = run_fla(inputs) + torch.cuda.synchronize() + + acc = compute_accuracy(sm90_out, fla_out) + + ms_fla = time_kernel(lambda inp=inputs: run_fla(inp)) + ms_sm90 = time_kernel(lambda inp=inputs: run_sm90(inp)) + speedup = ms_fla / ms_sm90 if ms_sm90 > 0 else float("inf") + + results.append( + { + "B": B, + "T": T, + "accuracy": acc, + "ms_fla": ms_fla, + "ms_sm90": ms_sm90, + "speedup": speedup, + } + ) + + del inputs + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Varlen benchmark +# ============================================================ +def bench_varlen(configs, H: int): + print(f"\n{'=' * 120}") + print(f" Varlen Benchmark: SM90 CuTe DSL vs FLA Triton (H={H}, K={K}, V={V}, BT={BT})") + print(f"{'=' * 120}") + results = [] + + for seq_lens, total_len, dist in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + T = total_len + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=DEVICE) + + inputs = prepare_bwd_wy_dqkg_fused_inputs( + B=1, + T=T, + H=H, + K=K, + V=V, + chunk_size=BT, + device=DEVICE, + seed=SEED, + cu_seqlens=cu_seqlens, + ) + + sm90_out = run_sm90(inputs) + fla_out = run_fla(inputs) + torch.cuda.synchronize() + + acc = compute_accuracy(sm90_out, fla_out) + + ms_fla = time_kernel(lambda inp=inputs: run_fla(inp)) + ms_sm90 = time_kernel(lambda inp=inputs: run_sm90(inp)) + speedup = ms_fla / ms_sm90 if ms_sm90 > 0 else float("inf") + + n_seqs = len(seq_lens) + min_l, max_l = min(seq_lens), max(seq_lens) + avg_l = T // n_seqs + tag = f"{dist:>7s} {n_seqs:>2d}seqs T={T} [{min_l}..{max_l}] avg={avg_l}" + + results.append( + { + "tag": tag, + "dist": dist, + "T_total": T, + "n_seqs": n_seqs, + "accuracy": acc, + "ms_fla": ms_fla, + "ms_sm90": ms_sm90, + "speedup": speedup, + } + ) + + del inputs + torch.cuda.empty_cache() + + return results + + +# ============================================================ +# Report +# ============================================================ +def print_report(fixed_results, varlen_results, H: int): + sep = "=" * 160 + wu = 1 if NCU_MODE else WARMUP + ni = 1 if NCU_MODE else N_ITERS + print(f"\n{sep}") + print(" BENCHMARK REPORT: chunk_kda_bwd SM90 vs FLA Triton") + print(f" H={H} K={K} V={V} BT={BT} BK={BK} BV={BV} OCC={MIN_OCC} dtype=bf16 Warmup={wu} Iters={ni}") + print(sep) + + acc_header = " ".join(f"{k:>10s}" for k in ACC_KEYS) + + if fixed_results: + print("\n [Fixed-Length]") + print(f" {'─' * 145}") + print(f" {'B':>3s} {'T':>5s} │ {'FLA(ms)':>9s} {'SM90(ms)':>9s} {'Speedup':>8s} │ {'':>10s}{acc_header}") + print(f" {'─' * 145}") + + for r in fixed_results: + rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in ACC_KEYS) + err_ratio_vals = " ".join(f"{r['accuracy'].get(k, {}).get('err_ratio', 0.0):10.6f}" for k in ACC_KEYS) + print( + f" {r['B']:3d} {r['T']:5d} │ " + f"{r['ms_fla']:9.4f} {r['ms_sm90']:9.4f} {r['speedup']:7.2f}x │ " + f"{'rel_max:':>10s}{rel_max_vals}" + ) + print(f" {'':3s} {'':5s} │ {'':9s} {'':9s} {'':8s} │ {'err_ratio:':>10s}{err_ratio_vals}") + print(f" {'─' * 145}") + + if varlen_results: + print("\n [Varlen]") + print(f" {'─' * 160}") + print(f" {'Config':>45s} │ {'FLA(ms)':>9s} {'SM90(ms)':>9s} {'Speedup':>8s} │ {'':>10s}{acc_header}") + print(f" {'─' * 160}") + + for r in varlen_results: + rel_max_vals = " ".join(f"{r['accuracy'].get(k, {}).get('rel_max', 0.0):10.6f}" for k in ACC_KEYS) + err_ratio_vals = " ".join(f"{r['accuracy'].get(k, {}).get('err_ratio', 0.0):10.6f}" for k in ACC_KEYS) + print( + f" {r['tag']:>45s} │ " + f"{r['ms_fla']:9.4f} {r['ms_sm90']:9.4f} {r['speedup']:7.2f}x │ " + f"{'rel_max:':>10s}{rel_max_vals}" + ) + print(f" {'':>45s} │ {'':9s} {'':9s} {'':8s} │ {'err_ratio:':>10s}{err_ratio_vals}") + print(f" {'─' * 160}") + + print(f"\n{sep}\n") + + +# ============================================================ +# Main +# ============================================================ +def main(): + global NCU_MODE, BK, BV, MIN_OCC + + parser = argparse.ArgumentParser(description="Benchmark SM90 bwd kernel vs FLA Triton") + parser.add_argument( + "--mode", + type=str, + default="both", + choices=["fixed", "varlen", "both"], + help="Which benchmark mode to run (default: both)", + ) + parser.add_argument("--heads", nargs="+", type=int, default=[H_DEFAULT]) + parser.add_argument("--bk", type=int, default=32, choices=[32, 64]) + parser.add_argument("--bv", type=int, default=64, choices=[32, 64]) + parser.add_argument("--occ", type=int, default=2, choices=[1, 2]) + parser.add_argument( + "--ncu", + action="store_true", + help="Run one SM90-only workload inside cudaProfilerStart/Stop for NCU.", + ) + parser.add_argument("--batch", type=int, default=1, help="Fixed-length profiling batch size.") + parser.add_argument("--total-len", type=int, default=16384, help="Profiling sequence/token length.") + parser.add_argument("--num-seqs", type=int, default=20, help="Varlen profiling sequence count.") + parser.add_argument( + "--dist", + type=str, + default="random", + choices=["uniform", "random", "skewed"], + help="Varlen profiling length distribution.", + ) + parser.add_argument( + "--profile-warmup", + "--warmup", + dest="profile_warmup", + type=int, + default=1, + help="Warmup iterations before the NCU profiler region.", + ) + parser.add_argument( + "--profile-iters", + type=int, + default=1, + help="SM90 iterations inside the NCU profiler region.", + ) + args = parser.parse_args() + + BK = args.bk + BV = args.bv + MIN_OCC = args.occ + + gpu_name = torch.cuda.get_device_name(0) + print(f"GPU: {gpu_name}") + print(f"K={K}, V={V}, BT={BT}, BK={BK}, BV={BV}, OCC={MIN_OCC}, dtype={DTYPE}") + + if args.ncu: + NCU_MODE = True + run_ncu_profile(args) + return + + fixed_configs = benchmark_fixed_configs() + varlen_configs = benchmark_varlen_configs() + + for H in args.heads: + fixed_res, varlen_res = [], [] + + if args.mode in ("fixed", "both"): + fixed_res = bench_fixed(fixed_configs, H) + + if args.mode in ("varlen", "both"): + varlen_res = bench_varlen(varlen_configs, H) + + print_report(fixed_res, varlen_res, H) + + print("All benchmarks done.") + + +if __name__ == "__main__": + main() diff --git a/cula/kda/chunk_bwd.py b/cula/kda/chunk_bwd.py index 40e56a86..0d2a7889 100644 --- a/cula/kda/chunk_bwd.py +++ b/cula/kda/chunk_bwd.py @@ -14,19 +14,22 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -"""SM100 modular chunk KDA backward orchestration""" +"""SM90/SM100 modular chunk KDA backward orchestration.""" +import functools import importlib import torch import triton import triton.language as tl from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h as _triton_chunk_gated_delta_rule_fwd_h from fla.ops.cp import FLACPContext from fla.ops.cp.chunk_delta_h import ( chunk_gated_delta_rule_bwd_dhu_pre_process, expand_h0, ) +from fla.ops.kda.chunk_bwd import recompute_w_u_fwd as recompute_w_u_fwd_triton from fla.ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices from fla.ops.utils.constant import RCP_LN2 @@ -37,14 +40,95 @@ check_shared_mem, ) -import cula.cudac as cula_cuda from cula.kda.chunk_intra import chunk_kda_bwd_intra from cula.ops.kda.sm100.bwd_wy_dqkg import chunk_kda_bwd_wy_dqkg_fused as chunk_kda_bwd_wy_dqkg_fused_cutedsl -from cula.utils import prepare_uniform_cu_seqlens +from cula.utils import get_device_sm_version, prepare_uniform_cu_seqlens _delta_h_mod = importlib.import_module("cula.ops.kda.sm100.delta_h") chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h + +@functools.cache +def _get_cula_cuda(): + return importlib.import_module("cula.cudac") + + +def _sm90_triton_chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + persistent: bool = True, + _no_cp: bool = False, + cu_seqlens_cpu: torch.Tensor | None = None, +): + """FLA Triton delta-H adapter for the SM90 backward path.""" + del persistent, _no_cp + return _triton_chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=False, + ) + + +def _get_chunk_delta_h_fwd(device: torch.device): + major, minor = get_device_sm_version(device) + if major == 9 and minor == 0: + return _sm90_triton_chunk_gated_delta_rule_fwd_h + if major == 10 and minor in (0, 3): + return chunk_gated_delta_rule_fwd_h + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) + + +def _select_recompute_w_u_backend(device: torch.device): + major, minor = get_device_sm_version(device) + if major == 9 and minor == 0: + return "triton" + if major == 10 and minor in (0, 3): + return "cuda_extension" + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) + + +def _select_chunk_kda_bwd_wy_dqkg_fused(q: torch.Tensor, v: torch.Tensor): + major, minor = get_device_sm_version(q.device) + if major == 9 and minor == 0: + if q.shape[2] == v.shape[2]: + return importlib.import_module("cula.ops.chunk_wy_dqkg_sm90").chunk_kda_bwd_wy_dqkg_fused + # The SM90 CuTe DSL kernel currently supports MHA only. Keep FLA/Triton + # behavior for grouped-value attention until the kernel gains GVA. + return chunk_kda_bwd_wy_dqkg_fused + if major == 10 and minor in (0, 3): + return chunk_kda_bwd_wy_dqkg_fused_cutedsl + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) + + BK_LIST = [32, 64] if check_shared_mem() else [16, 32] BV_LIST = [64, 128] if check_shared_mem("ampere") else [16, 32] NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] @@ -484,18 +568,46 @@ def chunk_kda_bwd( chunk_indices=chunk_indices, lower_bound=lower_bound, ) - reset_cu_seqlens = False - if cu_seqlens is None: - reset_cu_seqlens = True - cu_seqlens = prepare_uniform_cu_seqlens(B, T, q.device, torch.int32) - if chunk_indices is None and cu_seqlens is not None: + if cu_seqlens is not None and chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - # w, u, kg, qg all live in h_v head space. - w = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) - u = torch.empty_like(v) - qg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) if q is not None else None - kg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) - cula_cuda.recompute_w_u_cuda(k, v, beta, Akk, g, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q, qg) + reset_cu_seqlens = False + recompute_backend = _select_recompute_w_u_backend(q.device) + if recompute_backend == "cuda_extension": + if cu_seqlens is None: + reset_cu_seqlens = True + cu_seqlens = prepare_uniform_cu_seqlens(B, T, q.device, torch.int32) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # w, u, kg, qg all live in h_v head space. + w = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) + u = torch.empty_like(v) + qg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) if q is not None else None + kg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) + _get_cula_cuda().recompute_w_u_cuda( + k, + v, + beta, + Akk, + g, + cu_seqlens, + chunk_indices, + w, + u, + kg, + chunk_size, + q, + qg, + ) + else: + w, u, qg, kg = recompute_w_u_fwd_triton( + k=k, + v=v, + beta=beta, + A=Akk, + q=q, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) if cp_context is not None: # Restore the full initial_state tensor from the compressed version. # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. @@ -504,7 +616,8 @@ def chunk_kda_bwd( cu_seqlens = None chunk_indices = None # TODO: update to support only varlen (1,T,H,D) format - h, v_new, _ = chunk_gated_delta_rule_fwd_h( + chunk_delta_h_fwd = _get_chunk_delta_h_fwd(q.device) + h, v_new, _ = chunk_delta_h_fwd( k=kg, w=w, u=u, @@ -570,7 +683,8 @@ def chunk_kda_bwd( transpose_state_layout=transpose_state_layout, ) - dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused_cutedsl( + chunk_kda_bwd_wy_dqkg_fused_impl = _select_chunk_kda_bwd_wy_dqkg_fused(q, v) + dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused_impl( q=q, k=k, v=v, diff --git a/cula/kda/chunk_fwd.py b/cula/kda/chunk_fwd.py index f73a03c1..3694cedd 100644 --- a/cula/kda/chunk_fwd.py +++ b/cula/kda/chunk_fwd.py @@ -12,22 +12,112 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SM100 modular chunk KDA forward orchestration""" +"""SM90/SM100 modular chunk KDA forward orchestration.""" -import torch +import functools -# from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h +import torch +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h as _triton_chunk_gated_delta_rule_fwd_h from fla.ops.cp import FLACPContext from fla.ops.cp.chunk_delta_h import ( chunk_gated_delta_rule_fwd_h_pre_process, compress_h0, ) +from fla.ops.gla.chunk import chunk_gla_fwd_o_gk as _triton_chunk_gla_fwd_o_gk from fla.ops.kda.gate import kda_gate_chunk_cumsum from fla.ops.utils import chunk_local_cumsum from fla.ops.utils.constant import RCP_LN2 from cula.kda.chunk_intra import chunk_kda_fwd_intra -from cula.utils import assert_blackwell +from cula.utils import get_device_sm_version + + +def _sm90_triton_chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + persistent: bool = True, + _no_cp: bool = False, + cu_seqlens_cpu: torch.Tensor | None = None, + use_intracard_cp=None, +): + """FLA Triton delta-H adapter for the SM90 forward path.""" + del persistent, _no_cp, use_intracard_cp + return _triton_chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=False, + ) + + +def _sm90_triton_chunk_gla_fwd_o( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor | None, + A: torch.Tensor, + scale: float, + chunk_size: int = 64, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + is_varlen: bool = False, + persistent: bool = True, +): + """FLA Triton output adapter for the SM90 forward path.""" + del is_varlen, persistent + h_fla = h.flatten(0, 1) if h.dim() == 5 else h + out = _triton_chunk_gla_fwd_o_gk( + q=q, + v=v, + g=g, + A=A, + h=h_fla, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=False, + ) + if o is not None: + o.copy_(out) + return o + return out + + +@functools.cache +def _get_fwd_kernels_for_sm(major: int, minor: int): + if major == 9 and minor == 0: + return _sm90_triton_chunk_gated_delta_rule_fwd_h, _sm90_triton_chunk_gla_fwd_o + if major == 10 and minor in (0, 3): + from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h + from cula.ops.kda.sm100.fwd_o import chunk_gla_fwd_o + + return chunk_gated_delta_rule_fwd_h, chunk_gla_fwd_o + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) def chunk_kda_fwd( @@ -55,10 +145,7 @@ def chunk_kda_fwd( use_tf32_inverse: bool = True, unified_gref: bool = False, # Set True for ~5% extra perf (slightly lower precision) ): - assert_blackwell(q.device) - - from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h - from cula.ops.kda.sm100.fwd_o import chunk_gla_fwd_o + chunk_gated_delta_rule_fwd_h, chunk_gla_fwd_o = _get_fwd_kernels_for_sm(*get_device_sm_version(q.device)) # Apply gate activation g_org = None diff --git a/cula/kda/chunk_intra.py b/cula/kda/chunk_intra.py index f864849c..7f462f82 100644 --- a/cula/kda/chunk_intra.py +++ b/cula/kda/chunk_intra.py @@ -14,17 +14,26 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -"""SM100 modular chunk KDA intra-chunk wrapper and helpers""" +"""SM90/SM100 modular chunk KDA intra-chunk wrapper and helpers.""" + +import functools +import importlib import torch import triton import triton.language as tl +from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra as chunk_kda_bwd_intra_triton +from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as chunk_kda_fwd_intra_triton from fla.ops.utils import prepare_chunk_indices from fla.ops.utils.op import exp2, gather from fla.utils import IS_GATHER_SUPPORTED, autotune_cache_kwargs -import cula.cudac as cula_cuda -from cula.utils import prepare_uniform_cu_seqlens +from cula.utils import get_device_sm_version, prepare_uniform_cu_seqlens + + +@functools.cache +def _get_cula_cuda(): + return importlib.import_module("cula.cudac") @triton.heuristics( @@ -321,6 +330,26 @@ def chunk_kda_fwd_intra( unified_gref: bool = False, # Set True for ~5% extra perf (slightly lower precision) ): assert safe_gate, "Only safe_gate=True is supported in chunk_kda_fwd_intra for now" + major, minor = get_device_sm_version(q.device) + if major == 9 and minor == 0: + return chunk_kda_fwd_intra_triton( + q=q, + k=k, + v=v, + gk=gk, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + disable_recompute=disable_recompute, + ) + if not (major == 10 and minor in (0, 3)): + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) B, T, H, K = k.shape # GVA: g/beta/v live in h_v head space; q/k live in h_qk head space. HV = v.size(2) @@ -343,6 +372,7 @@ def chunk_kda_fwd_intra( Akk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) tile_counter = torch.zeros(1, dtype=torch.int32, device=q.device) + cula_cuda = _get_cula_cuda() cula_cuda.chunk_kda_fwd_intra_cuda( q, k, gk, beta, cu_seqlens, chunk_indices, Aqk, Akk, tile_counter, scale, chunk_size, use_tf32_inverse, unified_gref ) @@ -376,6 +406,29 @@ def chunk_kda_bwd_intra( chunk_size: int = 64, safe_gate: bool = False, ): + major, minor = get_device_sm_version(q.device) + if major == 9 and minor == 0: + return chunk_kda_bwd_intra_triton( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dk=dk, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + if not (major == 10 and minor in (0, 3)): + raise RuntimeError( + f"Unsupported CUDA compute capability sm_{major}{minor}. " + "Only sm90a (Hopper) and Blackwell (SM100/SM103) are supported." + ) B, T, H, K, HV = *k.shape, g.shape[2] BT = chunk_size BC = min(16, BT) diff --git a/cula/ops/__init__.py b/cula/ops/__init__.py index 60c60de0..0d177a18 100644 --- a/cula/ops/__init__.py +++ b/cula/ops/__init__.py @@ -13,6 +13,7 @@ # limitations under the License. __all__ = [ + "chunk_kda_bwd_wy_dqkg_fused", "kda_decode", "kda_decode_mtp", "kda_decode_mtp_recurrent", @@ -22,6 +23,10 @@ ] _LAZY = { + "chunk_kda_bwd_wy_dqkg_fused": ( + "cula.ops.chunk_wy_dqkg_sm90", + "chunk_kda_bwd_wy_dqkg_fused", + ), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), diff --git a/cula/ops/chunk_wy_dqkg_sm90.py b/cula/ops/chunk_wy_dqkg_sm90.py new file mode 100644 index 00000000..d01684c0 --- /dev/null +++ b/cula/ops/chunk_wy_dqkg_sm90.py @@ -0,0 +1,3625 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +chunk_wy_dqkg_sm90 — Hopper SM90 WGMMA implementation of KDA chunkwise +backward fused WY DqKG path. + +Architecture: Hopper warp specialization (1 DMA WG + 1 MMA WG). + DMA WG (warps 0-3): warp 0 = TMA G2S load, warp 1 = TMA S2G store + MMA WG (warps 4-7): WGMMA + r2s epilogue (writes SMEM only) + +Computes: + dq, dk, dv, db, dg, dA in FLA-compatible output order. +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.hopper_helpers as sm90_utils +import torch +from cutlass.cute.nvgpu import cpasync, warpgroup +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream +from cutlass.cute.typing import Int32, Int64 + + +def make_thread_cooperative_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +@cute.jit +def smem_load_f32x4_sw128(raw_ptr: cute.Pointer, row: Int32, col_base: Int32): + """ + Load 4 consecutive float32 from SMEM with 128B swizzle layout. + Layout: tile_to_shape(epi_smem_atom, (BT, BK)) where BK=64. + Atom row width = 32 elements. Two outer blocks for BK=64. + Swizzle: 128B → elem_xor = ((row & 7) << 2). + col_base must be 4-aligned. + """ + c_inner = col_base & Int32(31) + c_outer = col_base >> Int32(5) + swizzled_inner = c_inner ^ ((row & Int32(7)) << Int32(2)) + elem_offset = row * Int32(32) + swizzled_inner + c_outer * Int32(2048) + aligned_ptr = cute.make_ptr( + cutlass.Float32, + (raw_ptr + elem_offset).toint(), + cute.AddressSpace.smem, + assumed_align=16, + ) + t = cute.make_tensor(aligned_ptr, cute.make_layout((4,), stride=(1,))) + return t.load() + + +@cute.jit +def gmem_load_f32x4(gmem_addr: Int64): + """Load 4 contiguous fp32 values from 16-byte-aligned GMEM.""" + ptr = cute.make_ptr( + cutlass.Float32, + gmem_addr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + t = cute.make_tensor(ptr, cute.make_layout((4,), stride=(1,))) + return t.load() + + +@cute.jit +def gmem_store_f32x4(gmem_addr: Int64, val): + """Store 4 contiguous fp32 values to 16-byte-aligned GMEM.""" + ptr = cute.make_ptr( + cutlass.Float32, + gmem_addr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + t = cute.make_tensor(ptr, cute.make_layout((4,), stride=(1,))) + t.store(val) + + +@cute.jit +def copy_partial_epi_tile(tiled_copy, thr_copy, tOs, tOc, tOr, gmem_tile, rows: Int32): + """SMEM -> REG -> GMEM for a 64x32 epilogue tile with per-row mask.""" + tOg = thr_copy.partition_D(gmem_tile) + for m1 in cutlass.range_constexpr(cute.size(tOs.shape[1])): + row = tOc[(0, 0), m1, 0][0] + if row < rows: + cute.autovec_copy(tOs[(None, m1, None)], tOr[(None, m1, None)]) + cute.copy(tiled_copy, tOr[(None, m1, None)], tOg[(None, m1, None)]) + + +@cute.jit +def copy_partial_epi_tile_gmem_f32( + tiled_copy, + thr_copy, + tOs, + tOc, + tOr, + gmem_iter: cute.Pointer, + chunk_row_base: Int32, + row_stride: Int32, + head_idx: Int32, + head_stride: Int32, + col_base: Int32, + rows: Int32, +): + """Build a 64x32 fp32 GMEM tile view and copy only valid rows.""" + gmem_ptr = cute.make_ptr( + cutlass.Float32, + (gmem_iter + chunk_row_base * row_stride + head_idx * head_stride + col_base).toint(), + cute.AddressSpace.gmem, + assumed_align=16, + ) + stride_t = cute.assume(row_stride, divby=4) + gmem_tile = cute.make_tensor( + gmem_ptr, + cute.make_layout((64, 32), stride=(stride_t, 1)), + ) + copy_partial_epi_tile(tiled_copy, thr_copy, tOs, tOc, tOr, gmem_tile, rows) + + +@cute.jit +def copy_partial_epi_tile_gmem_bf16( + tiled_copy, + thr_copy, + tOs, + tOc, + tOr, + gmem_iter: cute.Pointer, + chunk_row_base: Int32, + row_stride: Int32, + head_idx: Int32, + head_stride: Int32, + col_base: Int32, + rows: Int32, +): + """Build a 64x32 bf16 GMEM tile view and copy only valid rows.""" + gmem_ptr = cute.make_ptr( + cutlass.BFloat16, + (gmem_iter + chunk_row_base * row_stride + head_idx * head_stride + col_base).toint(), + cute.AddressSpace.gmem, + assumed_align=16, + ) + stride_t = cute.assume(row_stride, divby=8) + gmem_tile = cute.make_tensor( + gmem_ptr, + cute.make_layout((64, 32), stride=(stride_t, 1)), + ) + copy_partial_epi_tile(tiled_copy, thr_copy, tOs, tOc, tOr, gmem_tile, rows) + + +USE_FAST_MATH = True +DEBUG_PRINT = False + +COMPILE_OPTIONS = "--enable-tvm-ffi --generate-line-info --ptxas-options '--verbose'" + +BFloat16 = cutlass.BFloat16 +Float32 = cutlass.Float32 + +_torch_to_cutlass_dtype = { + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, +} + + +@cute.jit +def smem_load_bf16x8_sw128(raw_ptr: cute.Pointer, row: Int32, col_base: Int32): + """Load 8 consecutive bf16 from SMEM with K_SW128 (Swizzle<3,4,3>) layout. + raw_ptr: bf16 SMEM base pointer for one stage + row: row index in [0, BK=64) + col_base: 8-aligned column index in [0, BV=64) + """ + swizzled = col_base ^ ((row & Int32(7)) << Int32(3)) + elem_off = row * Int32(64) + swizzled + aligned_ptr = cute.make_ptr( + BFloat16, + (raw_ptr + elem_off).toint(), + cute.AddressSpace.smem, + assumed_align=16, + ) + smem_t = cute.make_tensor(aligned_ptr, cute.make_layout((8,), stride=(1,))) + rmem_t = cute.make_fragment_like(smem_t) + cute.autovec_copy(smem_t, rmem_t) + return rmem_t + + +# ── Named barrier IDs ── +# barrier 0 is reserved by CUDA runtime (sync_threads), do not use. +BARRIER_DW_READY = 1 # dw stmatrix visible to all MMA warps (128 thr) +BARRIER_DG_COMPUTE = 2 # intra-MMA-WG sync for dgk/dg epilogue (128 thr) +BARRIER_DB_SYNC = 3 # sDb read/write synchronization (128 thr) + + +class ChunkKdaBwdWyDqkgFusedSM90: + """Hopper SM90 WGMMA kernel for KDA chunkwise WY DqKG backward.""" + + def __init__( + self, + chunk_size: int = 64, + head_dim_k: int = 128, + head_dim_v: int = 128, + acc_dtype: type[cutlass.Numeric] = cutlass.Float32, + io_dtype: type[cutlass.Numeric] = cutlass.BFloat16, + scale: float = 1.0, + min_occupancy: int = 2, + use_fast_math: bool = True, + bk: int = 32, + bv: int = 64, + ): + assert chunk_size == 64, "chunk_size must be 64" + assert head_dim_k == 128 and head_dim_v == 128 + assert bk in (32, 64), "bk must be 32 or 64" + assert bv in (32, 64), "bv must be 32 or 64" + assert head_dim_k % bk == 0 + assert head_dim_v % bv == 0 + + self.use_fast_math = use_fast_math + self.chunk_size = chunk_size + self.head_dim_k = head_dim_k + self.head_dim_v = head_dim_v + self.acc_dtype = acc_dtype + self.io_dtype = io_dtype + self.scale = scale + self.min_occupancy = min_occupancy + + self.BT = chunk_size # 64 + # head_dim_k is always 128, accessed via self.head_dim_k + self.BK = bk + self.BV = bv + self.num_v_tiles = (head_dim_v + self.BV - 1) // self.BV # 2 + self.num_k_iters = self.head_dim_k // self.BK # 2 + self.vloop_gemm_tiler = (self.BT, self.BK, self.BV) # M=BT, N=BK, K=BV + self.dv2_gemm_tiler = (self.BT, self.BV, self.BT) # M=BT, N=BV, K=BT + self.vloop_stage = max(2, self.num_v_tiles) + + self.threads_per_warp = 32 + self.num_warps_per_warp_group = 4 + self.num_threads_per_warp_group = 128 + self.num_dma_warp_groups = 1 + self.num_mma_warp_groups = 1 + self.threads_per_cta = self.num_threads_per_warp_group * (self.num_dma_warp_groups + self.num_mma_warp_groups) + + self.load_register_requirement = 40 + if self.min_occupancy >= 2: + self.mma_register_requirement = 200 + else: + self.mma_register_requirement = 256 + + self.persistent = True + hardware_info = cutlass.utils.HardwareInfo() + self.num_sm = hardware_info.get_device_multiprocessor_count() + self.buffer_align_bytes = 1024 + + # Epilogue tile: (BT, 32) — per k_iter writes BK/32 = 2 epi-tiles + self.epi_tile = (self.BT, 32) + self.epi_stage = 1 + self.num_epi_tiles = self.BK // self.epi_tile[1] # 2 + self.num_dA_epi_tiles = self.BT // self.epi_tile[1] # 2 + self.num_dv2_epi_tiles = self.BV // self.epi_tile[1] # 2 + + def _compute_grid(self, B: int, T: int, H: int, total_nt: Int32 | None = None): + assert total_nt is not None + total_tiles = total_nt * H + grid_x = cutlass.min(Int32(self.num_sm * self.min_occupancy), total_tiles) + return (grid_x, Int32(1), Int32(1)) + + @cute.jit + def __call__( + self, + do_in: cute.Tensor, # [B, T, H, V] bf16 + h_in: cute.Tensor, # [B, NT, H, K, V] bf16 + vnew_in: cute.Tensor, # [B, T, H, V] bf16 + dh_in: cute.Tensor, # [B, NT, H, K, V] bf16 + g_in: cute.Tensor, # [B, T, H, K] fp32 — gating + q_in: cute.Tensor, # [B, T, H, K] bf16 — query (for dg) + k_in: cute.Tensor, # [B, T, H, K] bf16 — key (for kg) + dq_in: cute.Tensor, # [B, T, H, K] fp32 + dk_in: cute.Tensor, # [B, T, H, K] fp32 + dg_in: cute.Tensor, # [B, T, H, K] fp32 + dv_in: cute.Tensor, # [B, T, H, V] bf16 + v_in: cute.Tensor, # [B, T, H, V] bf16 + A_in: cute.Tensor, # [B, T, H, BT] bf16 — intra-chunk attn + dA_out: cute.Tensor, # [B, T, H, BT] fp32 + dv2_out: cute.Tensor, # [B, T, H, V] bf16 + db_out: cute.Tensor, # [B, T, H] fp32 — gradient of beta + beta_in: cute.Tensor, # [B, T, H] fp32/bf16 — per-token decay scalar + cu_seqlens_in: cute.Tensor, # [N+1] int32 + chunk_indices_in: cute.Tensor, # [NT, 2] int32 + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + total_nt: Int32, + stream, + ): + do_ptr = do_in.iterator + h_ptr = h_in.iterator + vnew_ptr = vnew_in.iterator + dh_ptr = dh_in.iterator + g_ptr = g_in.iterator + q_ptr = q_in.iterator + k_ptr = k_in.iterator + dv_ptr = dv_in.iterator + v_ptr = v_in.iterator + A_ptr = A_in.iterator + dq_ptr = dq_in.iterator + dk_ptr = dk_in.iterator + dg_ptr = dg_in.iterator + dA_ptr = dA_out.iterator + dv2_ptr = dv2_out.iterator + db_ptr = db_out.iterator + beta_ptr = beta_in.iterator + + B, T, H, K, V = problem_size + BT, BV = self.BT, self.BV + data_B = Int32(1) + NT = total_nt + + # ===================== GMEM layouts ===================== + tv_layout = cute.make_layout( + (T, V, (H, data_B)), + stride=(H * V, 1, (V, T * H * V)), + ) + do = cute.make_tensor(do_ptr, tv_layout) + vnew = cute.make_tensor(vnew_ptr, tv_layout) + + h_layout = cute.make_layout( + (K, V, (NT, H)), + stride=(V, 1, (H * K * V, K * V)), + ) + h = cute.make_tensor(h_ptr, h_layout) + dh = cute.make_tensor(dh_ptr, h_layout) + + # q layout: bf16 [B, T, H, K] — query for dg computation + q_tk_layout = cute.make_layout( + (T, K, (H, data_B)), + stride=(H * K, 1, (K, T * H * K)), + ) + q = cute.make_tensor(q_ptr, q_tk_layout) + + # k layout: bf16 [B, T, H, K] — key for kg computation + k = cute.make_tensor(k_ptr, q_tk_layout) + + # g layout: fp32 [B, T, H, K] — same shape as dq + g_tk_layout = cute.make_layout( + (T, K, (H, data_B)), + stride=(H * K, 1, (K, T * H * K)), + ) + g = cute.make_tensor(g_ptr, g_tk_layout) + + dqk_layout = cute.make_layout( + (T, K, (H, data_B)), + stride=(H * K, 1, (K, T * H * K)), + ) + dq = cute.make_tensor(dq_ptr, dqk_layout) + dk = cute.make_tensor(dk_ptr, dqk_layout) + dg = cute.make_tensor(dg_ptr, dqk_layout) + + dv = cute.make_tensor(dv_ptr, tv_layout) + v = cute.make_tensor(v_ptr, tv_layout) + + # A^T: transposed GMEM view (BT, T) — first dim contiguous + a_t_layout = cute.make_layout( + (BT, T, (H, data_B)), + stride=(1, H * BT, (BT, T * H * BT)), + ) + A_attn = cute.make_tensor(A_ptr, a_t_layout) + + dA_layout = cute.make_layout( + (T, BT, (H, data_B)), + stride=(H * BT, 1, (BT, T * H * BT)), + ) + dA = cute.make_tensor(dA_ptr, dA_layout) + + # dv2: bf16 [B, T, H, V] — same layout as do/vnew + dv2_layout = cute.make_layout( + (T, V, (H, data_B)), + stride=(H * V, 1, (V, T * H * V)), + ) + dv2 = cute.make_tensor(dv2_ptr, dv2_layout) + + # beta: fp32 [B, T, H] — per-token decay scalar (no K dim) + beta_layout = cute.make_layout( + (T, (H, data_B)), + stride=(H, (1, T * H)), + ) + beta_gmem = cute.make_tensor(beta_ptr, beta_layout) + + # db: fp32 [B, T, H] — gradient of beta (same layout as beta) + db_gmem = cute.make_tensor(db_ptr, beta_layout) + + # ===================== TiledMMA ===================== + vloop_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + self.acc_dtype, + (1, 1, 1), + tiler_mn=(self.BT, self.BK), + ) + # 64×16 MMA layout — only used to derive ldmatrix tiled_copy for + # chunked q loading (reduces register pressure in dq*q elementwise) + q16_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + self.acc_dtype, + (1, 1, 1), + tiler_mn=(self.BT, 16), + ) + + # dwkg GEMM: dw(BT,BK) @ kg(BT,BK)^T → (BT,BT) + # A from registers (converted from vloop C layout), B from SMEM + dwkg_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + self.acc_dtype, + (1, 1, 1), + (self.BT, self.BT), + warpgroup.OperandSource.RMEM, + ) + dkgb_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.MN, + warpgroup.OperandMajorMode.K, + self.acc_dtype, + (1, 1, 1), + (self.BT, self.BK), + ) + # dA GEMM: dv(BT,BV) @ v^T(BV,BT) → (BT,BT), always m64n64 + dA_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + self.acc_dtype, + (1, 1, 1), + tiler_mn=(self.BT, self.BT), + ) + # dA post GEMM: sA(MN-major) @ scratch(K-major) → (BT,BT), always m64n64 + dA_post1_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.MN, + warpgroup.OperandMajorMode.K, + self.acc_dtype, + (1, 1, 1), + (self.BT, self.BT), + ) + # dv2 GEMM: A(BT,BT) @ dv(BT,BV) → (BT,BV) + # A from sA (COL_MAJOR → MN-major), B from sDv_col (BV,BT) COL_MAJOR MN-major + # sDv_col is a COL_MAJOR read view of buf_dv (same pattern as sA_row) + dv2_tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.io_dtype, + self.io_dtype, + warpgroup.OperandMajorMode.MN, + warpgroup.OperandMajorMode.MN, + self.acc_dtype, + (1, 1, 1), + (self.BT, self.BV), + ) + + # ===================== SMEM layouts ===================== + tv_smem_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.BV, + ), + self.io_dtype, + ) + tv_smem_layout_staged = cute.tile_to_shape( + tv_smem_atom, + cute.append((self.BT, self.BV), self.vloop_stage), + order=(0, 1, 2), + ) + + # COL_MAJOR read view of buf_dv for dv2 B-operand (same pattern as sA_row) + dv_col_smem_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.COL_MAJOR, + self.io_dtype, + self.BV, + ), + self.io_dtype, + ) + dv_col_smem_layout_staged = cute.tile_to_shape( + dv_col_smem_atom, + cute.append((self.BV, self.BT), self.vloop_stage), + order=(0, 1, 2), + ) + + kv_smem_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.BV, + ), + self.io_dtype, + ) + kv_smem_layout_staged = cute.tile_to_shape( + kv_smem_atom, + cute.append((self.BK, self.BV), self.vloop_stage), + order=(0, 1, 2), + ) + + epi_smem_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.acc_dtype, + self.epi_tile[1], + ), + self.acc_dtype, + ) + epi_smem_layout_staged = cute.tile_to_shape( + epi_smem_atom, + cute.append(self.epi_tile, self.epi_stage), + order=(0, 1, 2), + ) + + epi_smem_atom_bf16 = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.epi_tile[1], + ), + self.io_dtype, + ) + epi_smem_layout_staged_bf16 = cute.tile_to_shape( + epi_smem_atom_bf16, + cute.append(self.epi_tile, self.epi_stage), + order=(0, 1, 2), + ) + + # ===================== TMA atoms ===================== + tv_smem_no_stage = cute.slice_(tv_smem_layout_staged, (None, None, 0)) + kv_smem_no_stage = cute.slice_(kv_smem_layout_staged, (None, None, 0)) + epi_smem_no_stage = cute.tile_to_shape(epi_smem_atom, self.epi_tile, order=(0, 1)) + epi_smem_no_stage_bf16 = cute.tile_to_shape(epi_smem_atom_bf16, self.epi_tile, order=(0, 1)) + + tma_atom_do, tma_tensor_do = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + do, + tv_smem_no_stage, + (BT, BV), + ) + tma_atom_h, tma_tensor_h = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + h, + kv_smem_no_stage, + (self.BK, BV), + ) + tma_atom_vnew, tma_tensor_vnew = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + vnew, + tv_smem_no_stage, + (BT, BV), + ) + tma_atom_dh, tma_tensor_dh = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + dh, + kv_smem_no_stage, + (self.BK, BV), + ) + tma_atom_dq, tma_tensor_dq = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dq, + epi_smem_no_stage, + self.epi_tile, + ) + tma_atom_dk, tma_tensor_dk = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dk, + epi_smem_no_stage, + self.epi_tile, + ) + tma_atom_dg, tma_tensor_dg = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dg, + epi_smem_no_stage, + self.epi_tile, + ) + tma_atom_dg_reduce, tma_tensor_dg_reduce = cpasync.make_tiled_tma_atom( + cpasync.CopyReduceBulkTensorTileS2GOp(reduction_kind=cute.ReductionOp.ADD), + dg, + epi_smem_no_stage, + self.epi_tile, + ) + tma_atom_dv, tma_tensor_dv = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + dv, + tv_smem_no_stage, + (BT, BV), + ) + tma_atom_v, tma_tensor_v = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + v, + tv_smem_no_stage, + (BT, BV), + ) + tma_atom_dv2, tma_tensor_dv2 = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dv2, + epi_smem_no_stage_bf16, + self.epi_tile, + ) + tma_atom_dA, tma_tensor_dA = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dA, + epi_smem_no_stage, + self.epi_tile, + ) + + # g TMA: tile = (BT, 32), SMEM covers (BT, BK) — loaded per k_iter + self.g_tma_tile = (BT, 32) + self.num_g_tma_tiles_per_k = self.BK // self.g_tma_tile[1] # 2 = 64/32 + self.g_smem_layout = cute.tile_to_shape( + epi_smem_atom, + (BT, self.BK), + order=(0, 1), + ) + g_smem_layout = self.g_smem_layout + + # dg accumulator: (BT, BK) x f32, 2-stage — one stage per commit + # (2 k_iters × 2 parts = 4 commits per wu_iter). + self.num_dg_stages = 1 + self.dg_smem_layout = g_smem_layout # single-stage layout for TMA/read + self.dg_smem_layout_staged = cute.tile_to_shape( + epi_smem_atom, + cute.append((BT, self.BK), self.num_dg_stages), + order=(0, 1, 2), + ) + self.dg_smem_layout_write = cute.tile_to_shape( + epi_smem_atom, + (BT, 32, self.num_epi_tiles), + order=(0, 1, 2), + ) + self.dg_smem_layout_write_staged = cute.tile_to_shape( + epi_smem_atom, + cute.append((BT, 32, self.num_epi_tiles), self.num_dg_stages), + order=(0, 1, 2, 3), + ) + tma_atom_g, tma_tensor_g = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + g, + epi_smem_no_stage, + self.g_tma_tile, + ) + + # q TMA: bf16 (BT, BK), same tile (BT, 32) but bf16 swizzle atom + # q SMEM layout: (BT, BK) bf16, no stage — loaded once per wu_iter + q_smem_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.epi_tile[1], + ), + self.io_dtype, + ) + self.tk_smem_layout = cute.tile_to_shape( + q_smem_atom, + (BT, self.BK), + order=(0, 1), + ) + # q TMA tile: bf16 (BT, 32) — same shape as g TMA tile + q_smem_tma_slice = cute.tile_to_shape( + q_smem_atom, + self.g_tma_tile, + order=(0, 1), + ) + tma_atom_q, tma_tensor_q = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + q, + q_smem_tma_slice, + self.g_tma_tile, + ) + + # k TMA: bf16 (BT, BK), same layout/tile as q — loaded once per wu_iter + tma_atom_k, tma_tensor_k = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + k, + q_smem_tma_slice, + self.g_tma_tile, + ) + + # A SMEM: (BT, BT) COL_MAJOR — matches A^T GMEM (first dim contiguous) + # MN-major MMA A-operand also expects first dim contiguous + A_smem_atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.COL_MAJOR, + self.io_dtype, + BT, + ), + self.io_dtype, + ) + self.A_smem_layout = cute.tile_to_shape( + A_smem_atom, + (BT, BT), + order=(0, 1), + ) + # ROW_MAJOR read view of same buf_A — sA_row[i,j] = sA[j,i] + # Used as K-major B-operand for GEMM 2 in dA post-processing + A_smem_atom_row = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + BT, + ), + self.io_dtype, + ) + self.A_smem_layout_row = cute.tile_to_shape( + A_smem_atom_row, + (BT, BT), + order=(0, 1), + ) + tma_atom_A, tma_tensor_A = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + A_attn, + self.A_smem_layout, + (BT, BT), + ) + + # dw scratch: two views of the same buffer for stmatrix.trans write + WGMMA read + # Write view: (BT, BK) — stmatrix.trans writes MMA C(BT,BK) M-major + # Read view: (BK, BT) — dkgb B-operand K-major (BT contiguous) + dw_smem_atom_write = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.COL_MAJOR, + self.io_dtype, + BT, + ), + self.io_dtype, + ) + self.dw_smem_layout_write = cute.tile_to_shape( + dw_smem_atom_write, + (BT, self.BK), + order=(0, 1), + ) + dw_smem_atom_read = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + BT, + ), + self.io_dtype, + ) + self.dw_smem_layout_read = cute.tile_to_shape( + dw_smem_atom_read, + (self.BK, BT), + order=(0, 1), + ) + # Wide views for dA post-processing M matrix (BT×BT) + self.dw_smem_layout_write_wide = cute.tile_to_shape( + dw_smem_atom_write, + (BT, BT), + order=(0, 1), + ) + self.dw_smem_layout_read_wide = cute.tile_to_shape( + dw_smem_atom_read, + (BT, BT), + order=(0, 1), + ) + # buf_dw sized to max(narrow, wide) — wide needed when BK < BT + # Keep narrow read view as dw_smem_layout for the sDw tensor creation + self.dw_smem_layout = self.dw_smem_layout_read + self.dw_buf_cosize = max( + cute.cosize(self.dw_smem_layout_read), + cute.cosize(self.dw_smem_layout_read_wide), + ) + + # kg scratch: (BT, BK) x bf16, no stage — holds kg = k * exp2(g) + self.kg_smem_layout = cute.tile_to_shape( + q_smem_atom, + (BT, self.BK), + order=(0, 1), + ) + + # ===================== TMA byte counts ===================== + self.tma_bytes_tv = cute.size_in_bytes(self.io_dtype, tv_smem_no_stage) + self.tma_bytes_kv = cute.size_in_bytes(self.io_dtype, kv_smem_no_stage) + # g: num_g_tma_tiles_per_k TMA copies arrive on single barrier per k_iter + g_gmem_dtype = cutlass.Float32 + self.tma_bytes_g_single = cute.size_in_bytes(g_gmem_dtype, epi_smem_no_stage) + self.tma_bytes_g = self.tma_bytes_g_single * self.num_g_tma_tiles_per_k + # q/k: bf16 TMA tiles per k_iter (BK / 32 = 2 tiles per acquire) + self.tma_bytes_q_single = cute.size_in_bytes(self.io_dtype, q_smem_tma_slice) + self.num_q_tma_tiles_per_kiter = self.BK // self.g_tma_tile[1] # 2 + self.tma_bytes_q = self.tma_bytes_q_single * self.num_q_tma_tiles_per_kiter + self.tma_bytes_k = self.tma_bytes_q # same layout as q + # A: bf16 (BT, BT) = 8KB — one full TMA load per wu_iter + self.tma_bytes_A = cute.size_in_bytes(self.io_dtype, self.A_smem_layout) + + # ===================== SharedStorage ===================== + # ===================== SMEM budget (BK=32, OCC=2) ===================== + # buf_epi: (64, 32) x f32, 1 stage, 8 KB (also bf16 view for dv2) + # buf_tv: (64, 64) x bf16, 2 stages, 16 KB (shared: do/v/vnew) + # buf_h: (32, 64) x bf16, 2 stages, 8 KB + # buf_dh: (32, 64) x bf16, 2 stages, 8 KB + # buf_dv: (64, 64) x bf16, 2 stages, 16 KB + # buf_g: (64, 32) x f32, 1 stage, 8 KB + # buf_q: (64, 32) x bf16, 1 stage, 4 KB + # buf_k: (64, 32) x bf16, 1 stage, 4 KB + # buf_A: (64, 64) x bf16, 1 stage, 8 KB + # buf_dw: (64, 64) x bf16, 1 stage, 8 KB + # buf_kg: (64, 32) x bf16, 1 stage, 4 KB + # buf_dg: (64, 32) x f32, 1 stage, 8 KB + # buf_kdk: (64, 32) x f32, 1 stage, 8 KB + # buf_dgk_hdh: (128,) x f32, 1 stage, 0.5 KB + # buf_db: (64,) x f32, 1 stage, 0.25 KB + # s_beta: (64,) x f32, 1 stage, 0.25 KB + # + barriers + 1KB alignment padding per buffer + # ───────────────────────────────────────────────── + # BK=32 total: 113.66 KB (NCU measured), OCC=2 limit = 114 KB + # BK=64 total: ~151 KB (OCC=1 only) + @cute.struct + class SharedStorage: + bar_load_tv: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_load_h: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_load_dh: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_load_dv: cute.struct.MemRange[Int64, self.vloop_stage * 2] + bar_load_g: cute.struct.MemRange[Int64, 1 * 2] + bar_load_q: cute.struct.MemRange[Int64, 1 * 2] + bar_load_k: cute.struct.MemRange[Int64, 1 * 2] + bar_load_A: cute.struct.MemRange[Int64, 1 * 2] + bar_epi_ready: cute.struct.MemRange[Int64, self.epi_stage * 2] + bar_epi_done: cute.struct.MemRange[Int64, self.epi_stage * 2] + bar_dg_ready: cute.struct.MemRange[Int64, self.num_dg_stages * 2] + bar_dgk_hdh_ready: cute.struct.MemRange[Int64, self.num_k_iters * 2] + buf_epi: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(epi_smem_layout_staged)], + self.buffer_align_bytes, + ] + buf_tv: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(tv_smem_layout_staged)], + self.buffer_align_bytes, + ] + buf_h: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(kv_smem_layout_staged)], + self.buffer_align_bytes, + ] + buf_dh: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(kv_smem_layout_staged)], + self.buffer_align_bytes, + ] + buf_dv: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(tv_smem_layout_staged)], + self.buffer_align_bytes, + ] + buf_g: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(g_smem_layout)], + self.buffer_align_bytes, + ] + buf_q: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(self.tk_smem_layout)], + self.buffer_align_bytes, + ] + buf_k: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(self.tk_smem_layout)], + self.buffer_align_bytes, + ] + buf_A: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(self.A_smem_layout)], + self.buffer_align_bytes, + ] + buf_dw: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, self.dw_buf_cosize], + self.buffer_align_bytes, + ] + buf_kg: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(self.kg_smem_layout)], + self.buffer_align_bytes, + ] + buf_dg: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(self.dg_smem_layout_staged)], + self.buffer_align_bytes, + ] + # kdk scratch: (BT, BK) f32 for dgk column reduction + buf_kdk: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(g_smem_layout)], + self.buffer_align_bytes, + ] + # dgk_hdh cache: BK fp32 entries (512 bytes), computed by warp 2. + buf_dgk_hdh: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.head_dim_k], + 128, + ] + # db accumulator: BT fp32 entries (256 bytes) + buf_db: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BT], + 128, + ] + bar_load_beta: cute.struct.MemRange[Int64, 1 * 2] + s_beta: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.BT], + 128, + ] + + self.shared_storage = SharedStorage + + cu_seqlens = cute.make_tensor(cu_seqlens_in.iterator, cute.make_layout((B + 1,))) + chunk_indices = cute.make_tensor( + chunk_indices_in.iterator, + cute.make_layout((NT, 2), stride=(2, 1)), + ) + + grid = self._compute_grid(B, T, H, total_nt=NT) + + self._kernel( + vloop_tiled_mma, + q16_tiled_mma, + dwkg_tiled_mma, + dkgb_tiled_mma, + dv2_tiled_mma, + dA_tiled_mma, + dA_post1_tiled_mma, + tma_atom_do, + tma_tensor_do, + tma_atom_h, + tma_tensor_h, + tma_atom_vnew, + tma_tensor_vnew, + tma_atom_dh, + tma_tensor_dh, + tma_atom_dq, + tma_tensor_dq, + tma_atom_dk, + tma_tensor_dk, + tma_atom_dg, + tma_tensor_dg, + tma_atom_dg_reduce, + tma_tensor_dg_reduce, + tma_atom_dv, + tma_tensor_dv, + tma_atom_v, + tma_tensor_v, + tma_atom_dv2, + tma_tensor_dv2, + tma_atom_dA, + tma_tensor_dA, + tma_atom_A, + tma_tensor_A, + tma_atom_g, + tma_tensor_g, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + g, + beta_gmem, + db_gmem, + dq, + dk, + dg, + dA, + dv2, + tv_smem_layout_staged, + dv_col_smem_layout_staged, + kv_smem_layout_staged, + epi_smem_layout_staged, + epi_smem_layout_staged_bf16, + self.g_smem_layout, + self.tk_smem_layout, + self.A_smem_layout, + self.A_smem_layout_row, + self.dw_smem_layout, + self.dw_smem_layout_write, + self.dw_smem_layout_read_wide, + self.dw_smem_layout_write_wide, + self.kg_smem_layout, + self.dg_smem_layout, + self.dg_smem_layout_staged, + self.dg_smem_layout_write, + self.dg_smem_layout_write_staged, + cu_seqlens, + chunk_indices, + problem_size, + NT, + ).launch( + grid=grid, + block=(self.threads_per_cta, 1, 1), + stream=stream, + min_blocks_per_mp=self.min_occupancy, + ) + + # --------------------------------------------------------------- + # --------------------------------------------------------------- + # C→A layout conversion + # --------------------------------------------------------------- + @staticmethod + def convert_c_layout_to_a_layout(c, a): + """Convert accumulator C layout to A operand layout for RMEM-sourced WGMMA. + + Handles nested strides like (1,2,(4,16)) from non-coalesced WGMMA atoms. + """ + c_stride_0 = c.stride[0] + if isinstance(c_stride_0[2], tuple): + c_stride_0_flat = (c_stride_0[0], c_stride_0[1], c_stride_0[2][0]) + inner_base_stride = c_stride_0[2][0] + else: + c_stride_0_flat = c_stride_0 + inner_base_stride = c_stride_0[2] + + return cute.make_layout( + (a, c.shape[1], (c.shape[2], cute.size(c, mode=[0]) // cute.size(a))), + stride=( + c_stride_0_flat, + c.stride[1], + (c.stride[2], cute.size(a, mode=[2]) * inner_base_stride), + ), + ) + + @cute.jit + def make_acc_into_op(self, acc, tiled_mma_target, negate=False): + """Convert fp32 accumulator (C layout) → bf16 RMEM tensor (A layout). + + Follows FMHA's make_acc_into_op pattern: + 1. Compute A layout from C layout via convert_c_layout_to_a_layout + 2. Allocate RMEM tensor with A layout + 3. Write acc values (as bf16, optionally negated) through C-layout view + """ + a_layout = self.convert_c_layout_to_a_layout(acc.layout, tiled_mma_target.tv_layout_A.shape[1]) + operand = cute.make_rmem_tensor(a_layout, cutlass.BFloat16) + operand_as_acc = cute.make_tensor(operand.iterator, acc.layout) + for i in cutlass.range_constexpr(cute.size(acc)): + val = acc[i] + if negate: + val = -val + operand_as_acc[i] = cutlass.BFloat16(val) + return operand + + # --------------------------------------------------------------- + # Register-level row reduction helpers (from FMHA softmax pattern) + # --------------------------------------------------------------- + @staticmethod + def _layout_separate(thr, src, ref): + lt = cute.make_layout(()) + ge = cute.make_layout(()) + for k, v in enumerate(ref): + if cutlass.const_expr(v < thr): + lt = cute.append(lt, src[k]) + else: + ge = cute.append(ge, src[k]) + r = None + if cutlass.const_expr(cute.rank(lt) == 1): + r = cute.append(lt, ge) + else: + r = cute.append(cute.append(cute.make_layout(()), lt), ge) + return r + + @staticmethod + @cute.jit + def _layout_acc_mn(tiled_mma, acc): + separated = ChunkKdaBwdWyDqkgFusedSM90._layout_separate( + tiled_mma.shape_mnk[0], acc[0], tiled_mma.tv_layout_C.stride[1] + ) + V_M = separated[0] + V_N = separated[1] + V_M1 = None + V_N1 = None + if cutlass.const_expr(cute.rank(V_M) == 1): + V_M1 = cute.append(V_M, acc[1]) + else: + V_M1 = cute.append(cute.append(cute.make_layout(()), V_M), acc[1]) + if cutlass.const_expr(cute.rank(V_N) == 1): + V_N1 = cute.append(V_N, acc[2]) + else: + V_N1 = cute.append(cute.append(cute.make_layout(()), V_N), acc[2]) + r = None + if cutlass.const_expr(cute.rank(V_M1) == 1): + r = cute.append(V_M1, V_N1) + else: + r = cute.append(cute.append(cute.make_layout(()), V_M1), V_N1) + return r + + @staticmethod + @cute.jit + def _reduction_target_n(tiled_mma): + separated = ChunkKdaBwdWyDqkgFusedSM90._layout_separate( + tiled_mma.shape_mnk[0], + cute.make_layout(tiled_mma.tv_layout_C.shape[0]), + tiled_mma.tv_layout_C.stride[0], + ) + return separated[1] + + # --------------------------------------------------------------- + # TMA partition helpers + # --------------------------------------------------------------- + @cute.jit + def _tma_partition_A( + self, + tma_atom, + tma_tensor, + smem, + tile_shape, + tiled_mma, + batch_idx, + hidx, + ): + """Partition TMA tensor as MMA A-operand (M, K dims).""" + coord = (None, 0, None) + gX = cute.local_tile( + tma_tensor, + cute.slice_(tile_shape, coord), + (None, None, (hidx, batch_idx)), + ) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_A(gX) + tXsX, tXgX = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 2), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + @cute.jit + def _tma_partition_B( + self, + tma_atom, + tma_tensor, + smem, + tile_shape, + tiled_mma, + batch_idx, + hidx, + ): + """Partition TMA tensor as MMA B-operand (N, K dims).""" + coord = (0, None, None) + gX = cute.local_tile( + tma_tensor, + cute.slice_(tile_shape, coord), + (None, None, (hidx, batch_idx)), + ) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_B(gX) + tXsX, tXgX = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 2), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + # --------------------------------------------------------------- + # MMA copy factories + # --------------------------------------------------------------- + def _make_ldmatrix_copy_atom(self, transpose=False): + return cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + transpose=transpose, + num_matrices=4, + ), + self.io_dtype, + ) + + def _make_stmatrix_copy_atom(self, elem_ty, transpose=False): + return cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp( + transpose=transpose, + num_matrices=4, + ), + elem_ty, + ) + + def _make_r2s_tiled_copy(self, elem_ty_d, tiled_mma): + copy_atom_r2s = sm90_utils.sm90_get_smem_store_op( + utils.LayoutEnum.ROW_MAJOR, + elem_ty_d=elem_ty_d, + elem_ty_acc=self.acc_dtype, + ) + copy_atom_C = self._make_stmatrix_copy_atom(elem_ty_d) + tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma) + return cute.make_tiled_copy_S(copy_atom_r2s, tiled_copy_C_atom) + + def _make_stmatrix_r2s_tiled_copy(self, elem_ty_d, tiled_mma, transpose=False): + copy_atom = self._make_stmatrix_copy_atom(elem_ty_d, transpose=transpose) + tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom, tiled_mma) + return cute.make_tiled_copy_S(copy_atom, tiled_copy_C_atom) + + def _make_ldmatrix_c_tiled_copy(self, tiled_mma, transpose=False): + return cute.make_tiled_copy_C( + self._make_ldmatrix_copy_atom(transpose=transpose), + tiled_mma, + ) + + def _make_stmatrix_c_tiled_copy(self, elem_ty_d, tiled_mma, transpose=False): + return cute.make_tiled_copy_C( + self._make_stmatrix_copy_atom(elem_ty_d, transpose=transpose), + tiled_mma, + ) + + def _make_ldmatrix_a_tiled_copy(self, tiled_mma, transpose=False): + return cute.make_tiled_copy_A( + self._make_ldmatrix_copy_atom(transpose=transpose), + tiled_mma, + ) + + def _make_stmatrix_a_tiled_copy(self, elem_ty_d, tiled_mma, transpose=False): + return cute.make_tiled_copy_A( + self._make_stmatrix_copy_atom(elem_ty_d, transpose=transpose), + tiled_mma, + ) + + # --------------------------------------------------------------- + # Epilogue helper: r2s + signal store warp + # --------------------------------------------------------------- + @cute.jit + def _write_epi_tile( + self, + epi_idx, + tiled_copy_r2s_fp32, + tRS_rAcc, + tRS_sDq, + size_tRS_rD, + tRS_rD, + pipeline_epi_ready, + epi_ready_state, + ): + """Write one epi-tile from rmem to SMEM, then signal store warp.""" + pipeline_epi_ready.producer_acquire(epi_ready_state) + + # rmem chunk -> register staging + for epi_v in cutlass.range_constexpr(size_tRS_rD): + tRS_rD[epi_v] = tRS_rAcc[epi_idx * size_tRS_rD + epi_v] + + # r2s: register staging -> SMEM buffer + epi_buffer = epi_idx % cute.size(tRS_sDq, mode=[3]) + cute.copy( + tiled_copy_r2s_fp32, + tRS_rD, + tRS_sDq[(None, None, None, epi_buffer)], + ) + + # SMEM fence so store warp sees the writes + cute.arch.fence_view_async_shared() + + # Signal store warp that epi-tile is ready + pipeline_epi_ready.producer_commit(epi_ready_state) + + # --------------------------------------------------------------- + # Kernel body + # --------------------------------------------------------------- + @cute.kernel + def _kernel( + self, + vloop_tiled_mma: cute.TiledMma, + q16_tiled_mma: cute.TiledMma, + dwkg_tiled_mma: cute.TiledMma, + dkgb_tiled_mma: cute.TiledMma, + dv2_tiled_mma: cute.TiledMma, + dA_tiled_mma: cute.TiledMma, + dA_post1_tiled_mma: cute.TiledMma, + tma_atom_do: cute.CopyAtom, + tma_tensor_do: cute.Tensor, + tma_atom_h: cute.CopyAtom, + tma_tensor_h: cute.Tensor, + tma_atom_vnew: cute.CopyAtom, + tma_tensor_vnew: cute.Tensor, + tma_atom_dh: cute.CopyAtom, + tma_tensor_dh: cute.Tensor, + tma_atom_dq: cute.CopyAtom, + tma_tensor_dq: cute.Tensor, + tma_atom_dk: cute.CopyAtom, + tma_tensor_dk: cute.Tensor, + tma_atom_dg: cute.CopyAtom, + tma_tensor_dg: cute.Tensor, + tma_atom_dg_reduce: cute.CopyAtom, + tma_tensor_dg_reduce: cute.Tensor, + tma_atom_dv: cute.CopyAtom, + tma_tensor_dv: cute.Tensor, + tma_atom_v: cute.CopyAtom, + tma_tensor_v: cute.Tensor, + tma_atom_dv2: cute.CopyAtom, + tma_tensor_dv2: cute.Tensor, + tma_atom_dA: cute.CopyAtom, + tma_tensor_dA: cute.Tensor, + tma_atom_A: cute.CopyAtom, + tma_tensor_A: cute.Tensor, + tma_atom_g: cute.CopyAtom, + tma_tensor_g: cute.Tensor, + tma_atom_q: cute.CopyAtom, + tma_tensor_q: cute.Tensor, + tma_atom_k: cute.CopyAtom, + tma_tensor_k: cute.Tensor, + g_gmem: cute.Tensor, + beta_gmem: cute.Tensor, + db_gmem: cute.Tensor, + dq_gmem: cute.Tensor, + dk_gmem: cute.Tensor, + dg_gmem: cute.Tensor, + dA_gmem: cute.Tensor, + dv2_gmem: cute.Tensor, + tv_smem_layout_staged: cute.ComposedLayout, + dv_col_smem_layout_staged: cute.ComposedLayout, + kv_smem_layout_staged: cute.ComposedLayout, + epi_smem_layout_staged: cute.ComposedLayout, + epi_smem_layout_staged_bf16: cute.ComposedLayout, + g_smem_layout: cute.ComposedLayout, + tk_smem_layout: cute.ComposedLayout, + A_smem_layout: cute.ComposedLayout, + A_smem_layout_row: cute.ComposedLayout, # ROW_MAJOR read view of buf_A + dw_smem_layout: cute.ComposedLayout, # read view (BK, BT) + dw_smem_layout_write: cute.ComposedLayout, # write view (BT, BK) + dw_smem_layout_read_wide: cute.ComposedLayout, # read view (BT, BT) for dA post + dw_smem_layout_write_wide: cute.ComposedLayout, # write view (BT, BT) for dA post + kg_smem_layout: cute.ComposedLayout, + dg_smem_layout: cute.ComposedLayout, + dg_smem_layout_staged: cute.ComposedLayout, + dg_smem_layout_write: cute.ComposedLayout, + dg_smem_layout_write_staged: cute.ComposedLayout, + cu_seqlens: cute.Tensor, + chunk_indices: cute.Tensor, + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + NT: Int32, + ): + B, T, H, K, V = problem_size + BT, BV = self.BT, self.BV + + block_idx_x = cute.arch.block_idx()[0] + grid_dim_x = cute.arch.grid_dim()[0] + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + total_work_units = NT * H + num_iters = (total_work_units - block_idx_x + grid_dim_x - 1) // grid_dim_x + + # Warp assignment: WG0 warps 0-3 = DMA, WG1 warps 4-7 = MMA + load_warp_id = 0 + store_warp_id = 1 + + # Prefetch TMA descriptors + if warp_idx == load_warp_id: + cpasync.prefetch_descriptor(tma_atom_do) + cpasync.prefetch_descriptor(tma_atom_h) + cpasync.prefetch_descriptor(tma_atom_vnew) + cpasync.prefetch_descriptor(tma_atom_dh) + cpasync.prefetch_descriptor(tma_atom_dv) + cpasync.prefetch_descriptor(tma_atom_v) + cpasync.prefetch_descriptor(tma_atom_g) + cpasync.prefetch_descriptor(tma_atom_q) + cpasync.prefetch_descriptor(tma_atom_k) + cpasync.prefetch_descriptor(tma_atom_A) + if warp_idx == store_warp_id: + cpasync.prefetch_descriptor(tma_atom_dq) + cpasync.prefetch_descriptor(tma_atom_dk) + cpasync.prefetch_descriptor(tma_atom_dg) + cpasync.prefetch_descriptor(tma_atom_dg_reduce) + cpasync.prefetch_descriptor(tma_atom_dA) + cpasync.prefetch_descriptor(tma_atom_dv2) + + # ===================== SMEM allocation ===================== + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # ===================== Pipelines ===================== + # consumer_group size = num_mma_warps (not num_threads!). + # PipelineTmaAsync.consumer_release only arrives from is_signalling_thread + # (1 lane per warp for cluster_size=1), so the empty-barrier arrive_count + # must equal the number of warps, not threads. With arrive_count=128 the + # barrier never flips when DMA needs to reuse a stage (exposed by 2-deep). + num_mma_warps = self.num_threads_per_warp_group // self.threads_per_warp # 4 + num_h_dh_consumers = num_mma_warps + 1 # 5: MMA WG (4 warps) + warp 2 + pipeline_load_tv = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_tv.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_mma_warps), + tx_count=self.tma_bytes_tv, + ) + pipeline_load_h = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_h.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_h_dh_consumers), + tx_count=self.tma_bytes_kv, + ) + pipeline_load_dh = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_dh.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_h_dh_consumers), + tx_count=self.tma_bytes_kv, + ) + pipeline_load_dv = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_dv.data_ptr(), + num_stages=self.vloop_stage, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_mma_warps), + tx_count=self.tma_bytes_tv, + ) + + # Epilogue handshake: MMA WG (128 thr) -> store warp (1 thr) + # Must use PipelineAsync (not PipelineTmaAsync) because these are + # pure thread-to-thread notifications with no TMA involvement. + # PipelineTmaAsync.producer_commit() is a noop (designed for TMA + # hardware auto-arrive), so it would never signal the consumer. + # g pipeline: 1-deep (full BT×BK loaded once per wu_iter, reused by both passes) + pipeline_load_g = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_g.data_ptr(), + num_stages=1, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_mma_warps), + tx_count=self.tma_bytes_g, + ) + # q pipeline: 1-deep (full BT×BK loaded once per wu_iter, same as g) + pipeline_load_q = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_q.data_ptr(), + num_stages=1, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_mma_warps), + tx_count=self.tma_bytes_q, + ) + # k pipeline: 1-deep (full BT×BK loaded once per wu_iter, same as q) + pipeline_load_k = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_k.data_ptr(), + num_stages=1, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_mma_warps), + tx_count=self.tma_bytes_k, + ) + # A pipeline: 1-deep (BT×BT loaded once per wu_iter) + pipeline_load_A = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.bar_load_A.data_ptr(), + num_stages=1, + producer_group=make_thread_cooperative_group(1), + consumer_group=make_thread_cooperative_group(num_mma_warps), + tx_count=self.tma_bytes_A, + ) + + # beta pipeline: 1-deep, warp 2 (32 threads) -> MMA WG (128 threads) + # PipelineAsync uses thread counts (not warp counts like PipelineTmaAsync) + pipeline_load_beta = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_load_beta.data_ptr(), + num_stages=1, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group(self.num_threads_per_warp_group), + ) + + pipeline_epi_ready = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_epi_ready.data_ptr(), + num_stages=self.epi_stage, + producer_group=make_thread_cooperative_group(self.num_threads_per_warp_group), + consumer_group=make_thread_cooperative_group(self.threads_per_warp), + ) + pipeline_epi_done = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_epi_done.data_ptr(), + num_stages=self.epi_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group(self.num_threads_per_warp_group), + ) + + # dg store pipeline: MMA WG (128 thr) → warp 3 (32 thr), single-stage + pipeline_dg_ready = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_dg_ready.data_ptr(), + num_stages=self.num_dg_stages, + producer_group=make_thread_cooperative_group(self.num_threads_per_warp_group), + consumer_group=make_thread_cooperative_group(self.threads_per_warp), + ) + # dgk_hdh pipeline: one slot per k_iter, warp 2 -> MMA WG. + pipeline_dgk_hdh_ready = pipeline.PipelineAsync.create( + barrier_storage=storage.bar_dgk_hdh_ready.data_ptr(), + num_stages=self.num_k_iters, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group(self.num_threads_per_warp_group), + ) + + # ===================== SMEM tensors ===================== + sDo = storage.buf_tv.get_tensor(tv_smem_layout_staged.outer, swizzle=tv_smem_layout_staged.inner) + sH = storage.buf_h.get_tensor(kv_smem_layout_staged.outer, swizzle=kv_smem_layout_staged.inner) + sVnew = storage.buf_tv.get_tensor(tv_smem_layout_staged.outer, swizzle=tv_smem_layout_staged.inner) + sDh = storage.buf_dh.get_tensor(kv_smem_layout_staged.outer, swizzle=kv_smem_layout_staged.inner) + sDv = storage.buf_dv.get_tensor(tv_smem_layout_staged.outer, swizzle=tv_smem_layout_staged.inner) + sDv_col = storage.buf_dv.get_tensor(dv_col_smem_layout_staged.outer, swizzle=dv_col_smem_layout_staged.inner) + sV = storage.buf_tv.get_tensor(tv_smem_layout_staged.outer, swizzle=tv_smem_layout_staged.inner) + sEpi = storage.buf_epi.get_tensor(epi_smem_layout_staged.outer, swizzle=epi_smem_layout_staged.inner) + sEpi_bf16 = storage.buf_epi.get_tensor( + epi_smem_layout_staged_bf16.outer, + swizzle=epi_smem_layout_staged_bf16.inner, + dtype=self.io_dtype, + ) + sG = storage.buf_g.get_tensor(g_smem_layout.outer, swizzle=g_smem_layout.inner) + sDg_staged = storage.buf_dg.get_tensor(dg_smem_layout_staged.outer, swizzle=dg_smem_layout_staged.inner) + sDg_write_staged = storage.buf_dg.get_tensor( + dg_smem_layout_write_staged.outer, + swizzle=dg_smem_layout_write_staged.inner, + ) + sQ = storage.buf_q.get_tensor(tk_smem_layout.outer, swizzle=tk_smem_layout.inner) + sK = storage.buf_k.get_tensor(tk_smem_layout.outer, swizzle=tk_smem_layout.inner) + sA = storage.buf_A.get_tensor(A_smem_layout.outer, swizzle=A_smem_layout.inner) + sA_row = storage.buf_A.get_tensor(A_smem_layout_row.outer, swizzle=A_smem_layout_row.inner) + sDw = storage.buf_dw.get_tensor(dw_smem_layout.outer, swizzle=dw_smem_layout.inner) + # Write view (BT, BK) for stmatrix.trans — same physical buffer + sDw_write = storage.buf_dw.get_tensor(dw_smem_layout_write.outer, swizzle=dw_smem_layout_write.inner) + # Wide views (BT, BT) for dA post-processing M matrix + sDw_read_wide = storage.buf_dw.get_tensor(dw_smem_layout_read_wide.outer, swizzle=dw_smem_layout_read_wide.inner) + sDw_write_wide = storage.buf_dw.get_tensor(dw_smem_layout_write_wide.outer, swizzle=dw_smem_layout_write_wide.inner) + sKg = storage.buf_kg.get_tensor(kg_smem_layout.outer, swizzle=kg_smem_layout.inner) + sKdk_write = storage.buf_kdk.get_tensor(dg_smem_layout_write.outer, swizzle=dg_smem_layout_write.inner) + sKdk_raw_ptr = cute.make_ptr( + cutlass.Float32, + storage.buf_kdk.data_ptr().toint(), + cute.AddressSpace.smem, + ) + sG_raw_ptr = cute.make_ptr( + cutlass.Float32, + storage.buf_g.data_ptr().toint(), + cute.AddressSpace.smem, + ) + sDg_raw_ptr = cute.make_ptr( + cutlass.Float32, + storage.buf_dg.data_ptr().toint(), + cute.AddressSpace.smem, + ) + sDgkHdh = cute.make_tensor( + cute.make_ptr(cutlass.Float32, storage.buf_dgk_hdh.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((self.head_dim_k,), stride=(1,)), + ) + sDb = cute.make_tensor( + cute.make_ptr(cutlass.Float32, storage.buf_db.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((BT,), stride=(1,)), + ) + sBeta = cute.make_tensor( + cute.make_ptr(cutlass.Float32, storage.s_beta.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((BT,), stride=(1,)), + ) + sH_base = storage.buf_h.data_ptr().toint() + sDh_base = storage.buf_dh.data_ptr().toint() + kv_bytes_per_stage = self.tma_bytes_kv + + # ===================== Warp specialization ===================== + warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group) + is_dma_warp_group = warp_group_idx < self.num_dma_warp_groups + + # ══════════════════════════════════════════════════════════════ + # DMA WARP GROUP (warps 0-3) + # ══════════════════════════════════════════════════════════════ + if is_dma_warp_group: + cute.arch.setmaxregister_decrease(self.load_register_requirement) + + # ── Warp 0: TMA G2S load (do, h, vnew, dh, g) ── + if warp_idx == load_warp_id: + load_tv_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_h_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_dh_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_dv_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.vloop_stage) + load_g_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + load_q_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + load_k_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + load_A_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_t = work_idx // H + head_idx = work_idx % H + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + seq_tok_offset = cu_seqlens[(batch_idx,)] + BK = self.BK + + # ── Load A^T (BT, T) via COL_MAJOR SMEM ── + gA = cute.local_tile( + cute.domain_offset( + (Int32(0), seq_tok_offset, (Int32(0), Int32(0))), + tma_tensor_A, + ), + (BT, BT), + (0, tile_idx, (head_idx, Int32(0))), + ) + gA_for_tma = cute.zipped_divide(gA, (BT, BT)) + sA_for_tma = cute.zipped_divide(sA, (BT, BT)) + bSA_sA, bSA_gA = cpasync.tma_partition( + tma_atom_A, + 0, + cute.make_layout(1), + sA_for_tma, + gA_for_tma, + ) + pipeline_load_A.producer_acquire(load_A_ps) + cute.copy( + tma_atom_A, + bSA_gA[(None, (0, 0))], + bSA_sA[(None, 0)], + tma_bar_ptr=pipeline_load_A.producer_get_barrier(load_A_ps), + ) + pipeline_load_A.producer_commit(load_A_ps) + load_A_ps.advance() + + # ── dA V-loop: load dv + v (no k_iter dependency) ── + for v_iter in cutlass.range(self.num_v_tiles): + tma_dv_v = cute.domain_offset( + (seq_tok_offset, v_iter * BV, (Int32(0), Int32(0))), + tma_tensor_dv, + ) + tDVsDv, tDVgDv = self._tma_partition_A( + tma_atom_dv, + tma_dv_v, + sDv, + self.vloop_gemm_tiler, + vloop_tiled_mma, + Int32(0), + head_idx, + ) + pipeline_load_dv.producer_acquire(load_dv_ps) + dv_bar_ptr = pipeline_load_dv.producer_get_barrier(load_dv_ps) + cute.copy( + tma_atom_dv, + tDVgDv[(None, tile_idx, 0)], + tDVsDv[(None, load_dv_ps.index)], + tma_bar_ptr=dv_bar_ptr, + ) + pipeline_load_dv.producer_commit(load_dv_ps) + load_dv_ps.advance() + + tma_v_v = cute.domain_offset( + (seq_tok_offset, v_iter * BV, (Int32(0), Int32(0))), + tma_tensor_v, + ) + tVsV, tVgV = self._tma_partition_A( + tma_atom_v, + tma_v_v, + sV, + self.vloop_gemm_tiler, + vloop_tiled_mma, + Int32(0), + head_idx, + ) + pipeline_load_tv.producer_acquire(load_tv_ps) + cute.copy( + tma_atom_v, + tVgV[(None, tile_idx, 0)], + tVsV[(None, load_tv_ps.index)], + tma_bar_ptr=pipeline_load_tv.producer_get_barrier(load_tv_ps), + ) + pipeline_load_tv.producer_commit(load_tv_ps) + load_tv_ps.advance() + + # ── Unified k_iter loop: dq + dw + dk ── + for k_iter in cutlass.range(self.num_k_iters): + # ── Load g (BT, BK) fp32 per k_iter ── + gG = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_g, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gG_for_tma = cute.zipped_divide(gG, self.g_tma_tile) + sG_for_tma = cute.zipped_divide(sG, self.g_tma_tile) + bSG_sG, bSG_gG = cpasync.tma_partition( + tma_atom_g, + 0, + cute.make_layout(1), + sG_for_tma, + gG_for_tma, + ) + pipeline_load_g.producer_acquire(load_g_ps) + for k_sub in cutlass.range_constexpr(self.num_g_tma_tiles_per_k): + cute.copy( + tma_atom_g, + bSG_gG[(None, (0, k_sub))], + bSG_sG[(None, k_sub)], + tma_bar_ptr=pipeline_load_g.producer_get_barrier(load_g_ps), + ) + pipeline_load_g.producer_commit(load_g_ps) + load_g_ps.advance() + + # ── Load q (BT, BK) bf16 per k_iter ── + gQ = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_q, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gQ_for_tma = cute.zipped_divide(gQ, self.g_tma_tile) + sQ_for_tma = cute.zipped_divide(sQ, self.g_tma_tile) + bSQ_sQ, bSQ_gQ = cpasync.tma_partition( + tma_atom_q, + 0, + cute.make_layout(1), + sQ_for_tma, + gQ_for_tma, + ) + pipeline_load_q.producer_acquire(load_q_ps) + for k_sub in cutlass.range_constexpr(self.num_q_tma_tiles_per_kiter): + cute.copy( + tma_atom_q, + bSQ_gQ[(None, (0, k_sub))], + bSQ_sQ[(None, k_sub)], + tma_bar_ptr=pipeline_load_q.producer_get_barrier(load_q_ps), + ) + pipeline_load_q.producer_commit(load_q_ps) + load_q_ps.advance() + + # ── Load k (BT, BK) bf16 per k_iter ── + gK = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_k, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gK_for_tma = cute.zipped_divide(gK, self.g_tma_tile) + sK_for_tma = cute.zipped_divide(sK, self.g_tma_tile) + bSK_sK, bSK_gK = cpasync.tma_partition( + tma_atom_k, + 0, + cute.make_layout(1), + sK_for_tma, + gK_for_tma, + ) + pipeline_load_k.producer_acquire(load_k_ps) + for k_sub in cutlass.range_constexpr(self.num_q_tma_tiles_per_kiter): + cute.copy( + tma_atom_k, + bSK_gK[(None, (0, k_sub))], + bSK_sK[(None, k_sub)], + tma_bar_ptr=pipeline_load_k.producer_get_barrier(load_k_ps), + ) + pipeline_load_k.producer_commit(load_k_ps) + load_k_ps.advance() + + # dq+dw merged V-loop: load do + dv + h[k_iter] + # h is loaded ONCE and reused by both dq and dw compute loops + for v_iter in cutlass.range(self.num_v_tiles): + tma_do_v = cute.domain_offset( + (seq_tok_offset, v_iter * BV, (Int32(0), Int32(0))), + tma_tensor_do, + ) + tDOsDo, tDOgDo = self._tma_partition_A( + tma_atom_do, + tma_do_v, + sDo, + self.vloop_gemm_tiler, + vloop_tiled_mma, + Int32(0), + head_idx, + ) + pipeline_load_tv.producer_acquire(load_tv_ps) + cute.copy( + tma_atom_do, + tDOgDo[(None, tile_idx, 0)], + tDOsDo[(None, load_tv_ps.index)], + tma_bar_ptr=pipeline_load_tv.producer_get_barrier(load_tv_ps), + ) + pipeline_load_tv.producer_commit(load_tv_ps) + load_tv_ps.advance() + + tma_dv_v = cute.domain_offset( + (seq_tok_offset, v_iter * BV, (Int32(0), Int32(0))), + tma_tensor_dv, + ) + tDVsDv, tDVgDv = self._tma_partition_A( + tma_atom_dv, + tma_dv_v, + sDv, + self.vloop_gemm_tiler, + vloop_tiled_mma, + Int32(0), + head_idx, + ) + pipeline_load_dv.producer_acquire(load_dv_ps) + dv_bar_ptr = pipeline_load_dv.producer_get_barrier(load_dv_ps) + cute.copy( + tma_atom_dv, + tDVgDv[(None, tile_idx, 0)], + tDVsDv[(None, load_dv_ps.index)], + tma_bar_ptr=dv_bar_ptr, + ) + pipeline_load_dv.producer_commit(load_dv_ps) + load_dv_ps.advance() + + tma_h_v = cute.domain_offset( + (k_iter * BK, v_iter * BV, (0, 0)), + tma_tensor_h, + ) + tHsH, tHgH = self._tma_partition_B( + tma_atom_h, + tma_h_v, + sH, + self.vloop_gemm_tiler, + vloop_tiled_mma, + head_idx, + i_t, + ) + pipeline_load_h.producer_acquire(load_h_ps) + cute.copy( + tma_atom_h, + tHgH[(None, 0, 0)], + tHsH[(None, load_h_ps.index)], + tma_bar_ptr=pipeline_load_h.producer_get_barrier(load_h_ps), + ) + pipeline_load_h.producer_commit(load_h_ps) + load_h_ps.advance() + + # dk V-loop: load vnew + dh[k_iter] + for v_iter in cutlass.range(self.num_v_tiles): + tma_vnew_v = cute.domain_offset( + (seq_tok_offset, v_iter * BV, (Int32(0), Int32(0))), + tma_tensor_vnew, + ) + tVNsVN, tVNgVN = self._tma_partition_A( + tma_atom_vnew, + tma_vnew_v, + sVnew, + self.vloop_gemm_tiler, + vloop_tiled_mma, + Int32(0), + head_idx, + ) + pipeline_load_tv.producer_acquire(load_tv_ps) + cute.copy( + tma_atom_vnew, + tVNgVN[(None, tile_idx, 0)], + tVNsVN[(None, load_tv_ps.index)], + tma_bar_ptr=pipeline_load_tv.producer_get_barrier(load_tv_ps), + ) + pipeline_load_tv.producer_commit(load_tv_ps) + load_tv_ps.advance() + + tma_dh_v = cute.domain_offset( + (k_iter * BK, v_iter * BV, (0, 0)), + tma_tensor_dh, + ) + tDHsDH, tDHgDH = self._tma_partition_B( + tma_atom_dh, + tma_dh_v, + sDh, + self.vloop_gemm_tiler, + vloop_tiled_mma, + head_idx, + i_t, + ) + pipeline_load_dh.producer_acquire(load_dh_ps) + cute.copy( + tma_atom_dh, + tDHgDH[(None, 0, 0)], + tDHsDH[(None, load_dh_ps.index)], + tma_bar_ptr=pipeline_load_dh.producer_get_barrier(load_dh_ps), + ) + pipeline_load_dh.producer_commit(load_dh_ps) + load_dh_ps.advance() + + # ── Warp 1: TMA S2G store (dv2, dq, dk, dg, dA) ── + elif warp_idx == store_warp_id: + epi_ready_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.epi_stage) + epi_done_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.epi_stage) + + sEpi_for_tma = cute.group_modes(sEpi, 0, 2) + sEpi_bf16_for_tma = cute.group_modes(sEpi_bf16, 0, 2) + + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 1, + ), + ) + + tidx_in_warp = cute.arch.thread_idx()[0] % Int32(32) + universal_copy_bits = 128 + + epi_copy_elems_f32 = universal_copy_bits // self.acc_dtype.width + atom_universal_copy_f32 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.acc_dtype, + num_bits_per_copy=universal_copy_bits, + ) + epi_thr_dim1_f32 = self.epi_tile[1] // epi_copy_elems_f32 + epi_thr_dim0_f32 = self.threads_per_warp // epi_thr_dim1_f32 + epi_thr_layout_f32 = cute.make_ordered_layout( + (epi_thr_dim0_f32, epi_thr_dim1_f32), + order=(1, 0), + ) + epi_val_layout_f32 = cute.make_layout((1, epi_copy_elems_f32)) + gmem_tiled_copy_epi_f32 = cute.make_tiled_copy_tv( + atom_universal_copy_f32, + epi_thr_layout_f32, + epi_val_layout_f32, + ) + epi_thr_copy_f32 = gmem_tiled_copy_epi_f32.get_slice(tidx_in_warp) + sEpi_stage = sEpi[(None, None, 0)] + tOsEpi_f32 = epi_thr_copy_f32.partition_S(sEpi_stage) + tOcEpi_f32 = epi_thr_copy_f32.partition_S(cute.make_identity_tensor(self.epi_tile)) + tOrEpi_f32 = cute.make_fragment_like(tOsEpi_f32, self.acc_dtype) + + epi_copy_elems_bf16 = universal_copy_bits // self.io_dtype.width + atom_universal_copy_bf16 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.io_dtype, + num_bits_per_copy=universal_copy_bits, + ) + epi_thr_dim1_bf16 = self.epi_tile[1] // epi_copy_elems_bf16 + epi_thr_dim0_bf16 = self.threads_per_warp // epi_thr_dim1_bf16 + epi_thr_layout_bf16 = cute.make_ordered_layout( + (epi_thr_dim0_bf16, epi_thr_dim1_bf16), + order=(1, 0), + ) + epi_val_layout_bf16 = cute.make_layout((1, epi_copy_elems_bf16)) + gmem_tiled_copy_epi_bf16 = cute.make_tiled_copy_tv( + atom_universal_copy_bf16, + epi_thr_layout_bf16, + epi_val_layout_bf16, + ) + epi_thr_copy_bf16 = gmem_tiled_copy_epi_bf16.get_slice(tidx_in_warp) + sEpi_bf16_stage = sEpi_bf16[(None, None, 0)] + tOsEpi_bf16 = epi_thr_copy_bf16.partition_S(sEpi_bf16_stage) + tOcEpi_bf16 = epi_thr_copy_bf16.partition_S(cute.make_identity_tensor(self.epi_tile)) + tOrEpi_bf16 = cute.make_fragment_like(tOsEpi_bf16, self.io_dtype) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_t = work_idx // H + head_idx = work_idx % H + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + seq_tok_offset = cu_seqlens[(batch_idx,)] + seq_end = cu_seqlens[(batch_idx + Int32(1),)] + seq_len = seq_end - seq_tok_offset + sub_seq_len = cutlass.min(Int32(BT), seq_len - tile_idx * Int32(BT)) + chunk_row_base = seq_tok_offset + tile_idx * Int32(BT) + BK = self.BK + + # ── Store dv2 epi-tiles (per v_iter, before k_iter loop) ── + for v_iter in cutlass.range(self.num_v_tiles): + gDv2 = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_dv2, + ), + (BT, self.BV), + (tile_idx, v_iter, (head_idx, Int32(0))), + ) + gDv2_for_tma = cute.zipped_divide(gDv2, self.epi_tile) + bSG_sEpi_dv2, bSG_gDv2 = cpasync.tma_partition( + tma_atom_dv2, + 0, + cute.make_layout(1), + sEpi_bf16_for_tma, + gDv2_for_tma, + ) + epi_tile_shape_dv2 = gDv2_for_tma.shape[1] + epi_tile_layout_dv2 = cute.make_layout(epi_tile_shape_dv2, stride=(epi_tile_shape_dv2[1], 1)) + + for epi_idx in cutlass.range_constexpr(self.num_dv2_epi_tiles): + pipeline_epi_done.producer_acquire(epi_done_ps) + pipeline_epi_ready.consumer_wait(epi_ready_cs) + if sub_seq_len == Int32(BT): + epi_buffer = epi_idx % cute.size(bSG_sEpi_dv2, mode=[1]) + gmem_coord = epi_tile_layout_dv2.get_hier_coord(epi_idx) + cute.copy( + tma_atom_dv2, + bSG_sEpi_dv2[(None, epi_buffer)], + bSG_gDv2[(None, gmem_coord)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + else: + gmem_col_base = v_iter * Int32(self.BV) + epi_idx * Int32(32) + copy_partial_epi_tile_gmem_bf16( + gmem_tiled_copy_epi_bf16, + epi_thr_copy_bf16, + tOsEpi_bf16, + tOcEpi_bf16, + tOrEpi_bf16, + dv2_gmem.iterator, + chunk_row_base, + H * V, + head_idx, + V, + gmem_col_base, + sub_seq_len, + ) + pipeline_epi_ready.consumer_release(epi_ready_cs) + epi_ready_cs.advance() + pipeline_epi_done.producer_commit(epi_done_ps) + epi_done_ps.advance() + + # ── Unified k_iter: Store dq + dk + dg epi-tiles ── + for k_iter in cutlass.range(self.num_k_iters): + # dq epi-tiles — hybrid: full chunk uses TMA bulk + # store; ragged partial chunk falls back to per-thread + # per-thread store with `row < sub_seq_len` mask. + gDq = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_dq, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gDq_for_tma = cute.zipped_divide(gDq, self.epi_tile) + bSG_sEpi_dq, bSG_gDq = cpasync.tma_partition( + tma_atom_dq, + 0, + cute.make_layout(1), + sEpi_for_tma, + gDq_for_tma, + ) + epi_tile_shape_dq = gDq_for_tma.shape[1] + epi_tile_layout_dq = cute.make_layout(epi_tile_shape_dq, stride=(epi_tile_shape_dq[1], 1)) + + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + pipeline_epi_done.producer_acquire(epi_done_ps) + pipeline_epi_ready.consumer_wait(epi_ready_cs) + if sub_seq_len == Int32(BT): + epi_buffer = epi_idx % cute.size(bSG_sEpi_dq, mode=[1]) + gmem_coord = epi_tile_layout_dq.get_hier_coord(epi_idx) + cute.copy( + tma_atom_dq, + bSG_sEpi_dq[(None, epi_buffer)], + bSG_gDq[(None, gmem_coord)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + else: + gmem_col_base = k_iter * Int32(BK) + epi_idx * Int32(32) + copy_partial_epi_tile_gmem_f32( + gmem_tiled_copy_epi_f32, + epi_thr_copy_f32, + tOsEpi_f32, + tOcEpi_f32, + tOrEpi_f32, + dq_gmem.iterator, + chunk_row_base, + H * K, + head_idx, + K, + gmem_col_base, + sub_seq_len, + ) + pipeline_epi_ready.consumer_release(epi_ready_cs) + epi_ready_cs.advance() + pipeline_epi_done.producer_commit(epi_done_ps) + epi_done_ps.advance() + + # dk epi-tiles + gDk = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_dk, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gDk_for_tma = cute.zipped_divide(gDk, self.epi_tile) + bSG_sEpi_dk, bSG_gDk = cpasync.tma_partition( + tma_atom_dk, + 0, + cute.make_layout(1), + sEpi_for_tma, + gDk_for_tma, + ) + epi_tile_shape_dk = gDk_for_tma.shape[1] + epi_tile_layout_dk = cute.make_layout(epi_tile_shape_dk, stride=(epi_tile_shape_dk[1], 1)) + + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + pipeline_epi_done.producer_acquire(epi_done_ps) + pipeline_epi_ready.consumer_wait(epi_ready_cs) + if sub_seq_len == Int32(BT): + epi_buffer = epi_idx % cute.size(bSG_sEpi_dk, mode=[1]) + gmem_coord = epi_tile_layout_dk.get_hier_coord(epi_idx) + cute.copy( + tma_atom_dk, + bSG_sEpi_dk[(None, epi_buffer)], + bSG_gDk[(None, gmem_coord)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + else: + gmem_col_base = k_iter * Int32(BK) + epi_idx * Int32(32) + copy_partial_epi_tile_gmem_f32( + gmem_tiled_copy_epi_f32, + epi_thr_copy_f32, + tOsEpi_f32, + tOcEpi_f32, + tOrEpi_f32, + dk_gmem.iterator, + chunk_row_base, + H * K, + head_idx, + K, + gmem_col_base, + sub_seq_len, + ) + pipeline_epi_ready.consumer_release(epi_ready_cs) + epi_ready_cs.advance() + pipeline_epi_done.producer_commit(epi_done_ps) + epi_done_ps.advance() + + # ── Store dA epi-tiles (after all k_iters) ── + gDA = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_dA, + ), + (BT, BT), + (tile_idx, 0, (head_idx, Int32(0))), + ) + gDA_for_tma = cute.zipped_divide(gDA, self.epi_tile) + bSG_sEpi_dA, bSG_gDA = cpasync.tma_partition( + tma_atom_dA, + 0, + cute.make_layout(1), + sEpi_for_tma, + gDA_for_tma, + ) + epi_tile_shape_dA = gDA_for_tma.shape[1] + epi_tile_layout_dA = cute.make_layout(epi_tile_shape_dA, stride=(epi_tile_shape_dA[1], 1)) + for epi_idx in cutlass.range_constexpr(self.num_dA_epi_tiles): + pipeline_epi_done.producer_acquire(epi_done_ps) + pipeline_epi_ready.consumer_wait(epi_ready_cs) + if sub_seq_len == Int32(BT): + epi_buffer = epi_idx % cute.size(bSG_sEpi_dA, mode=[1]) + gmem_coord = epi_tile_layout_dA.get_hier_coord(epi_idx) + cute.copy( + tma_atom_dA, + bSG_sEpi_dA[(None, epi_buffer)], + bSG_gDA[(None, gmem_coord)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + else: + gmem_col_base = epi_idx * Int32(32) + copy_partial_epi_tile_gmem_f32( + gmem_tiled_copy_epi_f32, + epi_thr_copy_f32, + tOsEpi_f32, + tOcEpi_f32, + tOrEpi_f32, + dA_gmem.iterator, + chunk_row_base, + H * Int32(BT), + head_idx, + Int32(BT), + gmem_col_base, + sub_seq_len, + ) + pipeline_epi_ready.consumer_release(epi_ready_cs) + epi_ready_cs.advance() + pipeline_epi_done.producer_commit(epi_done_ps) + epi_done_ps.advance() + + c_pipeline.producer_tail() + + # ── Warp 2: load beta + compute dgk_hdh from SMEM h/dh ── + elif warp_idx == 2: + load_beta_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + dgk_hdh_ready_ps = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_k_iters, + ) + h_wait_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + h_release_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + dh_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + lane_idx = tidx % 32 + BK = self.BK + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_t = work_idx // H + head_idx = work_idx % H + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + seq_tok_offset = cu_seqlens[(batch_idx,)] + seq_end = cu_seqlens[(batch_idx + Int32(1),)] + sub_seq_len = cutlass.min(Int32(BT), seq_end - seq_tok_offset - tile_idx * Int32(BT)) + chunk_tok_offset = seq_tok_offset + tile_idx * BT + + pipeline_load_beta.producer_acquire(load_beta_ps) + for i in cutlass.range_constexpr(2): # BT=64 / 32 threads = 2 + idx = lane_idx + i * 32 + if idx < sub_seq_len: + sBeta[(idx,)] = cutlass.Float32(beta_gmem[(chunk_tok_offset + idx, (head_idx, Int32(0)))]) + else: + sBeta[(idx,)] = cutlass.Float32(0.0) + cute.arch.fence_view_async_shared() + pipeline_load_beta.producer_commit(load_beta_ps) + load_beta_ps.advance() + + NUM_ROWS_PER_THREAD = self.BK // 32 + + for k_iter in cutlass.range(self.num_k_iters): + pipeline_dgk_hdh_ready.producer_acquire(dgk_hdh_ready_ps) + + dgk_partials = cute.make_rmem_tensor((NUM_ROWS_PER_THREAD,), Float32) + dgk_partials.fill(Float32(0.0)) + + # Phase 1: wait for h (don't release yet) + for v_iter in cutlass.range(self.num_v_tiles): + pipeline_load_h.consumer_wait(h_wait_cs) + h_wait_cs.advance() + + # Phase 2: wait for dh, read both h+dh, release both + for v_iter in cutlass.range(self.num_v_tiles): + pipeline_load_dh.consumer_wait(dh_cs) + + sH_raw = cute.make_ptr( + BFloat16, + sH_base + h_release_cs.index * kv_bytes_per_stage, + cute.AddressSpace.smem, + ) + sDh_raw = cute.make_ptr( + BFloat16, + sDh_base + dh_cs.index * kv_bytes_per_stage, + cute.AddressSpace.smem, + ) + + for col_chunk in cutlass.range_constexpr(self.BV // 8): + col_base = Int32(col_chunk * 8) + h_dh = cute.make_rmem_tensor((8,), Float32) + for r in cutlass.range_constexpr(NUM_ROWS_PER_THREAD): + row = lane_idx + Int32(r * 32) + h_vals = smem_load_bf16x8_sw128(sH_raw, row, col_base) + dh_vals = smem_load_bf16x8_sw128(sDh_raw, row, col_base) + h_dh.store(h_vals.load().to(Float32) * dh_vals.load().to(Float32)) + for j in cutlass.range_constexpr(8): + dgk_partials[r] = dgk_partials[r] + h_dh[j] + + pipeline_load_h.consumer_release(h_release_cs) + h_release_cs.advance() + pipeline_load_dh.consumer_release(dh_cs) + dh_cs.advance() + + for r in cutlass.range_constexpr(NUM_ROWS_PER_THREAD): + sDgkHdh[(k_iter * Int32(BK) + lane_idx + Int32(r * 32),)] = dgk_partials[r] + cute.arch.fence_view_async_shared() + pipeline_dgk_hdh_ready.producer_commit(dgk_hdh_ready_ps) + dgk_hdh_ready_ps.advance() + + # ── Warp 3: TMA S2G store/reduce for dg (single-stage) ── + elif warp_idx == 3: + dg_ready_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_dg_stages) + + tidx_in_warp = cute.arch.thread_idx()[0] % Int32(32) + universal_copy_bits = 128 + epi_copy_elems_f32 = universal_copy_bits // self.acc_dtype.width + atom_universal_copy_f32 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.acc_dtype, + num_bits_per_copy=universal_copy_bits, + ) + epi_thr_dim1_f32 = self.epi_tile[1] // epi_copy_elems_f32 + epi_thr_dim0_f32 = self.threads_per_warp // epi_thr_dim1_f32 + epi_thr_layout_f32 = cute.make_ordered_layout( + (epi_thr_dim0_f32, epi_thr_dim1_f32), + order=(1, 0), + ) + epi_val_layout_f32 = cute.make_layout((1, epi_copy_elems_f32)) + gmem_tiled_copy_epi_f32 = cute.make_tiled_copy_tv( + atom_universal_copy_f32, + epi_thr_layout_f32, + epi_val_layout_f32, + ) + epi_thr_copy_f32 = gmem_tiled_copy_epi_f32.get_slice(tidx_in_warp) + cEpi_f32 = cute.make_identity_tensor(self.epi_tile) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_t = work_idx // H + head_idx = work_idx % H + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + seq_tok_offset = cu_seqlens[(batch_idx,)] + seq_end = cu_seqlens[(batch_idx + Int32(1),)] + sub_seq_len = cutlass.min(Int32(BT), seq_end - seq_tok_offset - tile_idx * Int32(BT)) + chunk_row_base = seq_tok_offset + tile_idx * Int32(BT) + BK = self.BK + + for k_iter in cutlass.range(self.num_k_iters): + gDg = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_dg, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gDg_for_tma_tile = cute.zipped_divide(gDg, self.epi_tile) + epi_tile_shape_dg = gDg_for_tma_tile.shape[1] + epi_tile_layout_dg = cute.make_layout(epi_tile_shape_dg, stride=(epi_tile_shape_dg[1], 1)) + + gDg_r = cute.local_tile( + cute.domain_offset( + (seq_tok_offset, Int32(0), (Int32(0), Int32(0))), + tma_tensor_dg_reduce, + ), + (BT, BK), + (tile_idx, k_iter, (head_idx, Int32(0))), + ) + gDg_r_for_tma = cute.zipped_divide(gDg_r, self.epi_tile) + + # ── Part 1: regular TMA store (even stage) ── + sDg_cur = sDg_staged[(None, None, dg_ready_cs.index)] + tOsDg_f32 = epi_thr_copy_f32.partition_S(sDg_cur) + tOcDg_f32 = epi_thr_copy_f32.partition_S(cEpi_f32) + tOrDg_f32 = cute.make_fragment_like(tOsDg_f32, self.acc_dtype) + sDg_cur_for_tma = cute.zipped_divide(sDg_cur, self.epi_tile) + bSG_sDg, bSG_gDg = cpasync.tma_partition( + tma_atom_dg, + 0, + cute.make_layout(1), + sDg_cur_for_tma, + gDg_for_tma_tile, + ) + + pipeline_dg_ready.consumer_wait(dg_ready_cs) + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + gmem_coord = epi_tile_layout_dg.get_hier_coord(epi_idx) + if sub_seq_len == Int32(BT): + cute.copy( + tma_atom_dg, + bSG_sDg[(None, gmem_coord)], + bSG_gDg[(None, gmem_coord)], + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + gmem_col_base = k_iter * Int32(BK) + epi_idx * Int32(32) + copy_partial_epi_tile_gmem_f32( + gmem_tiled_copy_epi_f32, + epi_thr_copy_f32, + tOsDg_f32, + tOcDg_f32, + tOrDg_f32, + dg_gmem.iterator, + chunk_row_base, + H * K, + head_idx, + K, + gmem_col_base, + sub_seq_len, + ) + pipeline_dg_ready.consumer_release(dg_ready_cs) + dg_ready_cs.advance() + + # ── Part 2: TMA reduce_add (odd stage) ── + sDg_cur2 = sDg_staged[(None, None, dg_ready_cs.index)] + sDg_cur2_for_tma = cute.zipped_divide(sDg_cur2, self.epi_tile) + bSG_sDg_r, bSG_gDg_r = cpasync.tma_partition( + tma_atom_dg_reduce, + 0, + cute.make_layout(1), + sDg_cur2_for_tma, + gDg_r_for_tma, + ) + + pipeline_dg_ready.consumer_wait(dg_ready_cs) + if chunk_row_base + Int32(BT) <= T: + # For ragged but physically in-bounds chunks, OOB + # rows were zeroed before the SMEM write, so full + # tile reduce_add is a no-op for invalid rows. + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + gmem_coord = epi_tile_layout_dg.get_hier_coord(epi_idx) + cute.copy( + tma_atom_dg_reduce, + bSG_sDg_r[(None, gmem_coord)], + bSG_gDg_r[(None, gmem_coord)], + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + # Tail chunks cannot use a full 64-row TMA reduce_add: + # the physical packed buffer may end before the tile does. + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + gmem_col_base = k_iter * Int32(BK) + epi_idx * Int32(32) + smem_col_base = epi_idx * Int32(32) + for row_block in cutlass.range_constexpr(16): + row = Int32(row_block * 4) + tidx_in_warp // Int32(8) + if row < sub_seq_len: + col_low = (tidx_in_warp % Int32(8)) * Int32(4) + gmem_addr = ( + dg_gmem.iterator + + (chunk_row_base + row) * H * K + + head_idx * K + + gmem_col_base + + col_low + ).toint() + old_f32 = gmem_load_f32x4(gmem_addr) + add_f32 = smem_load_f32x4_sw128( + sDg_raw_ptr, + row, + smem_col_base + col_low, + ) + out_f32 = old_f32 + add_f32 + gmem_store_f32x4(gmem_addr, out_f32) + pipeline_dg_ready.consumer_release(dg_ready_cs) + dg_ready_cs.advance() + + # DMA WG done — load warp and store warp both finish here. + # No CTA-wide sync: store warp communicates with MMA WG + # asynchronously via bar_epi_ready / bar_epi_done mbarriers. + + # ══════════════════════════════════════════════════════════════ + # MMA WARP GROUP (warps 4-7) + # ══════════════════════════════════════════════════════════════ + else: + cute.arch.setmaxregister_increase(self.mma_register_requirement) + + mma_warp_group_thread_layout = cute.make_layout( + self.num_mma_warp_groups, + stride=self.num_threads_per_warp_group, + ) + thr_mma = vloop_tiled_mma.get_slice(mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups)) + + # Fragments for staged SMEM operands (dq path: do, h) + tCsDo = thr_mma.partition_A(sDo) + tCsH = thr_mma.partition_B(sH) + tCrDo = vloop_tiled_mma.make_fragment_A(tCsDo) + tCrH = vloop_tiled_mma.make_fragment_B(tCsH) + + # Fragments for dk path (vnew, dh) — separate SMEM buffers + tCsVnew = thr_mma.partition_A(sVnew) + tCsDh = thr_mma.partition_B(sDh) + tCrVnew = vloop_tiled_mma.make_fragment_A(tCsVnew) + tCrDh = vloop_tiled_mma.make_fragment_B(tCsDh) + + # Fragments for dw path (dv @ h → dw, using vloop_tiled_mma) + tCsDv = thr_mma.partition_A(sDv) + tCrDv = vloop_tiled_mma.make_fragment_A(tCsDv) + + # Fragments for dA path (dv, v) — always m64n64 via dA_tiled_mma + thr_dA = dA_tiled_mma.get_slice(mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups)) + tCsDv_dA = thr_dA.partition_A(sDv) + tCsV_dA = thr_dA.partition_B(sV) + tCrDv_dA = dA_tiled_mma.make_fragment_A(tCsDv_dA) + tCrV_dA = dA_tiled_mma.make_fragment_B(tCsV_dA) + num_k_blocks_dA = cute.size(tCrDv_dA, mode=[2]) + + sEpi_no_stage = cute.slice_(sEpi, (None, None, 0)) + tCsEpi = thr_mma.partition_C(sEpi_no_stage) + acc_shape = tCsEpi.shape + dq_acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + dw_acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + dk_acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + + # dwkg B-operand fragment: partition sKg for dwkg_tiled_mma + thr_dwkg = dwkg_tiled_mma.get_slice(mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups)) + tCsKg_B = thr_dwkg.partition_B(sKg) + tCrKg = dwkg_tiled_mma.make_fragment_B(tCsKg_B) + num_k_blocks_dwkg = cute.size(tCrKg, mode=[2]) + + # dA_acc: always 64×64 via dwkg partition_C + dA_acc_ref = thr_dwkg.partition_C(sA) + dA_acc = cute.make_rmem_tensor(dA_acc_ref.shape, self.acc_dtype) + + # dkgb GEMM: A(BT,BT) @ (-dw)(BT,BK) -> (BT,BK) + thr_dkgb = dkgb_tiled_mma.get_slice(mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups)) + tCsA_dkgb = thr_dkgb.partition_A(sA) + tCsDw_dkgb = thr_dkgb.partition_B(sDw) + tCrA_dkgb = dkgb_tiled_mma.make_fragment_A(tCsA_dkgb) + tCrDw_dkgb = dkgb_tiled_mma.make_fragment_B(tCsDw_dkgb) + num_k_blocks_dkgb = cute.size(tCrA_dkgb, mode=[2]) + dkgb_acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + + # dA post-processing GEMM 2: temp @ sA via dwkg_tiled_mma + # B-operand = sA_row (ROW_MAJOR view of buf_A, K-major) + tCsA_row_post = thr_dwkg.partition_B(sA_row) + tCrA_row_post = dwkg_tiled_mma.make_fragment_B(tCsA_row_post) + num_k_blocks_post2 = cute.size(tCrA_row_post, mode=[2]) + + # dA post-processing GEMM 1: sA @ M via dA_post1_tiled_mma (always m64n64) + thr_dA_post1 = dA_post1_tiled_mma.get_slice( + mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups) + ) + tCsA_post1 = thr_dA_post1.partition_A(sA) + tCsDw_post1 = thr_dA_post1.partition_B(sDw_read_wide) + tCrA_post1 = dA_post1_tiled_mma.make_fragment_A(tCsA_post1) + tCrDw_post1 = dA_post1_tiled_mma.make_fragment_B(tCsDw_post1) + num_k_blocks_post1 = cute.size(tCrA_post1, mode=[2]) + + # dv2 GEMM: A(BT,BT) @ dv(BT,BV) -> (BT,BV) + # B-operand uses sDv_col (BV,BT) MN-major: COL_MAJOR view of buf_dv + thr_dv2 = dv2_tiled_mma.get_slice(mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups)) + tCsA_dv2 = thr_dv2.partition_A(sA) + tCsDv_dv2 = thr_dv2.partition_B(sDv_col) + tCrA_dv2 = dv2_tiled_mma.make_fragment_A(tCsA_dv2) + tCrDv_dv2 = dv2_tiled_mma.make_fragment_B(tCsDv_dv2) + num_k_blocks_dv2 = cute.size(tCrA_dv2, mode=[2]) + dv2_acc_shape = thr_dv2.partition_C(sEpi_no_stage).shape + dv2_acc = cute.make_rmem_tensor(dv2_acc_shape, self.acc_dtype) + + num_k_blocks = cute.size(tCrDo, mode=[2]) + + load_tv_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + load_h_wait_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + load_h_release_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + load_dh_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + load_g_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + load_q_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + load_k_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + load_A_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + load_beta_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + load_dv_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.vloop_stage) + + # MMA copy setup + tiled_copy_r2s_fp32 = self._make_r2s_tiled_copy(self.acc_dtype, vloop_tiled_mma) + mma_tidx = tidx - self.num_threads_per_warp_group * self.num_dma_warp_groups + thr_copy_r2s = tiled_copy_r2s_fp32.get_slice(mma_tidx) + tRS_sEpi = thr_copy_r2s.partition_D(sEpi) + tRS_sDg = thr_copy_r2s.partition_D(sDg_write_staged) + tRS_sKdk_write = thr_copy_r2s.partition_D(sKdk_write) + + rD_shape = cute.shape(thr_copy_r2s.partition_S(sEpi)) + tRS_rD_layout = cute.make_layout(rD_shape[:3]) + tRS_rD = cute.make_rmem_tensor_like(tRS_rD_layout, self.acc_dtype) + size_tRS_rD = cute.size(tRS_rD) + + # R2S tiled copy for dw with stmatrix.trans (bf16, transpose=True) + # MMA C-layout (BT, BK) → SMEM (BK, BT) with BT contiguous + tiled_copy_r2s_dw = self._make_stmatrix_r2s_tiled_copy( + self.io_dtype, + vloop_tiled_mma, + transpose=True, + ) + thr_copy_r2s_dw = tiled_copy_r2s_dw.get_slice(mma_tidx) + tRS_sDw = thr_copy_r2s_dw.partition_D(sDw_write) + + # R2S for dA (based on dwkg m64n64 — always 64×64 C layout) + tiled_copy_r2s_dA_fp32 = self._make_r2s_tiled_copy(self.acc_dtype, dwkg_tiled_mma) + thr_copy_r2s_dA = tiled_copy_r2s_dA_fp32.get_slice(mma_tidx) + tRS_sEpi_dA = thr_copy_r2s_dA.partition_D(sEpi) + rD_shape_dA = cute.shape(thr_copy_r2s_dA.partition_S(sEpi)) + tRS_rD_dA = cute.make_rmem_tensor_like(cute.make_layout(rD_shape_dA[:3]), self.acc_dtype) + size_tRS_rD_dA = cute.size(tRS_rD_dA) + + # R2S dw-transpose for dA M write (based on dwkg m64n64) + tiled_copy_r2s_dA_dw = self._make_stmatrix_r2s_tiled_copy( + self.io_dtype, + dwkg_tiled_mma, + transpose=True, + ) + thr_copy_r2s_dA_dw = tiled_copy_r2s_dA_dw.get_slice(mma_tidx) + tRS_sDw_wide = thr_copy_r2s_dA_dw.partition_D(sDw_write_wide) + + # R2S tiled copy for dv2: dv2_tiled_mma C layout → epi SMEM (bf16) + tiled_copy_r2s_bf16 = self._make_r2s_tiled_copy(self.io_dtype, dv2_tiled_mma) + thr_copy_r2s_dv2 = tiled_copy_r2s_bf16.get_slice(mma_tidx) + tRS_sEpi_dv2 = thr_copy_r2s_dv2.partition_D(sEpi_bf16) + + rD_shape_dv2 = cute.shape(thr_copy_r2s_dv2.partition_S(sEpi_bf16)) + tRS_rD_layout_dv2 = cute.make_layout(rD_shape_dv2[:3]) + tRS_rD_dv2 = cute.make_rmem_tensor_like(tRS_rD_layout_dv2, self.io_dtype) + size_tRS_rD_dv2 = cute.size(tRS_rD_dv2) + + # S2R tiled copy for sV: ldmatrix bulk load, aligned with dv2 MMA C partition + tiled_copy_s2r_v = self._make_ldmatrix_c_tiled_copy(dv2_tiled_mma) + thr_copy_s2r_v = tiled_copy_s2r_v.get_slice(mma_tidx) + tSR_sV = thr_copy_s2r_v.partition_S(sV) + tSR_rV_shape = cute.slice_(tSR_sV.shape, (None, None, None, 0)) + tSR_rV = cute.make_rmem_tensor(tSR_rV_shape, self.io_dtype) + + # Visitor-style kg elementwise path: + # ldmatrix sK -> register, reuse exp_g from the dq gate, stmatrix -> sKg. + # Uses 64×16 chunked copies to reduce register pressure. + # 64×64 copies — kept for kg_load (dA post-processing, line ~2690) + tiled_copy_s2r_kg = self._make_ldmatrix_c_tiled_copy(vloop_tiled_mma) + thr_copy_s2r_kg = tiled_copy_s2r_kg.get_slice(mma_tidx) + + # 64×16 chunked copies for kg computation (k load + kg store) + tiled_copy_r2s_kg16 = self._make_stmatrix_c_tiled_copy( + self.io_dtype, + q16_tiled_mma, + ) + thr_copy_r2s_kg16 = tiled_copy_r2s_kg16.get_slice(mma_tidx) + + # 64×16 chunked q loading — reduces register pressure vs full 64×64 + tiled_copy_s2r_q16 = self._make_ldmatrix_c_tiled_copy(q16_tiled_mma) + thr_copy_s2r_q16 = tiled_copy_s2r_q16.get_slice(mma_tidx) + + # Partitions: 64×16 chunked for k, kg, q + tKG16_sK = thr_copy_s2r_q16.partition_S(sK) + tKG16_sKg = thr_copy_r2s_kg16.partition_D(sKg) + tQ16_sQ = thr_copy_s2r_q16.partition_S(sQ) + tile16_shape_s2r = cute.slice_(tQ16_sQ.shape, (None, None, 0)) + tile16_shape_r2s = cute.slice_(tKG16_sKg.shape, (None, None, 0)) + rK16 = cute.make_rmem_tensor(tile16_shape_s2r, self.io_dtype) + rKg16 = cute.make_rmem_tensor(tile16_shape_r2s, self.io_dtype) + tQ16_rQ = cute.make_rmem_tensor(tile16_shape_s2r, self.io_dtype) + num_q16_tiles = cute.size(tQ16_sQ.shape, mode=[2]) + size_per_q16 = cute.size(tile16_shape_s2r) + + tKG_sKg_load = thr_copy_s2r_kg.partition_S(sKg) + tKG_rKg_load = cute.make_rmem_tensor(tKG_sKg_load.shape, self.io_dtype) + + # Ragged-tail sA fixup via ldmatrix -> registers -> stmatrix. + mma_op_A_zero = cute.nvgpu.warp.MmaF16BF16Op( + ab_dtype=self.io_dtype, + acc_dtype=self.acc_dtype, + shape_mnk=(16, 8, 16), + ) + tiled_mma_A_zero = cute.make_tiled_mma( + mma_op_A_zero, + atom_layout_mnk=(4, 1, 1), + permutation_mnk=(BT, BT, BT), + ) + tiled_copy_s2r_A_zero = self._make_ldmatrix_a_tiled_copy( + tiled_mma_A_zero, + transpose=True, + ) + tiled_copy_r2s_A_zero = self._make_stmatrix_a_tiled_copy( + self.io_dtype, + tiled_mma_A_zero, + transpose=True, + ) + thr_copy_s2r_A_zero = tiled_copy_s2r_A_zero.get_slice(mma_tidx) + thr_copy_r2s_A_zero = tiled_copy_r2s_A_zero.get_slice(mma_tidx) + thr_mma_A_zero = tiled_mma_A_zero.get_slice(mma_tidx) + tAZ_sA = thr_copy_s2r_A_zero.partition_S(sA) + tAZ_sA_store = thr_copy_r2s_A_zero.partition_D(sA) + tAZ_rA_proto = thr_mma_A_zero.make_fragment_A(thr_mma_A_zero.partition_A(sA)) + tAZ_rA = cute.make_fragment_like(tAZ_rA_proto, self.io_dtype) + cA_zero = cute.make_identity_tensor((BT, BT)) + tAZ_cA = thr_mma_A_zero.partition_A(cA_zero) + + copy_atom_s2r_g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.acc_dtype, + num_bits_per_copy=64, + ) + tiled_copy_s2r_g = cute.make_tiled_copy_D( + copy_atom_s2r_g, + tiled_copy_r2s_fp32, + ) + thr_copy_s2r_g = tiled_copy_s2r_g.get_slice(mma_tidx) + tSG_sG = thr_copy_s2r_g.partition_D(sG) + tSG_rG = cute.make_rmem_tensor(tSG_sG.shape, self.acc_dtype) + + epi_ready_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.epi_stage) + epi_done_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.epi_stage) + dg_ready_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_dg_stages) + dgk_hdh_ready_cs = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_k_iters, + ) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_t = work_idx // H + head_idx = work_idx % H + batch_idx = chunk_indices[(i_t, 0)] + tile_idx = chunk_indices[(i_t, 1)] + seq_tok_offset = cu_seqlens[(batch_idx,)] + seq_end_wu = cu_seqlens[(batch_idx + Int32(1),)] + sub_seq_len = cutlass.min(Int32(BT), seq_end_wu - seq_tok_offset - tile_idx * Int32(BT)) + chunk_tok_offset = seq_tok_offset + tile_idx * BT + BK = self.BK + + # Wait for A and beta early; dgk_hdh is consumed per k_iter. + pipeline_load_A.consumer_wait(load_A_cs) + pipeline_load_beta.consumer_wait(load_beta_cs) + + # Initialize sDb accumulator for db computation + compute_tidx = tidx - Int32(self.num_threads_per_warp_group) + if compute_tidx < Int32(BT): + sDb[(compute_tidx,)] = cutlass.Float32(0.0) + # Partial chunk: zero sA OOB columns [sub_seq_len, BT) so the + # dkgb = sA @ (-dw) GEMM does not consume rows pulled in by the + # full-tile TMA load past this sequence's tail. + if sub_seq_len < Int32(BT): + tAZ_rA_copy = thr_copy_s2r_A_zero.retile(tAZ_rA) + cute.copy(tiled_copy_s2r_A_zero, tAZ_sA, tAZ_rA_copy) + for i in cutlass.range_constexpr(cute.size(tAZ_rA)): + if tAZ_cA[i][1] >= sub_seq_len: + tAZ_rA[i] = self.io_dtype(0.0) + tAZ_rA_store = thr_copy_r2s_A_zero.retile(tAZ_rA) + cute.copy(tiled_copy_r2s_A_zero, tAZ_rA_store, tAZ_sA_store) + cute.arch.fence_view_async_shared() + pipeline.NamedBarrier(barrier_id=BARRIER_DB_SYNC, num_threads=128).sync() + + # ═══════════════════════════════════════════════ + # dA = dv @ v^T + dv2 = A @ dv (no k_iter, no g-scaling) + # Each v_iter: dA accumulates, dv2 is complete per v_iter + # ═══════════════════════════════════════════════ + dA_acc.fill(0.0) + + cD_dv2 = cute.make_identity_tensor((BT, self.BV)) + + # Register-level db_v reduction setup (FMHA pattern) + # Use per-thread MMA slice (runtime tidx) for identity partition + thr_dv2_per_thread = dv2_tiled_mma.get_slice(mma_tidx) + dv2_mn_layout = self._layout_acc_mn(dv2_tiled_mma, dv2_acc.layout) + n_rows_dv2 = cute.size(dv2_mn_layout, mode=[0]) + tCcC_dv2 = thr_dv2_per_thread.partition_C(cD_dv2) + coord_mn_dv2 = cute.make_tensor(tCcC_dv2.iterator, self._layout_acc_mn(dv2_tiled_mma, tCcC_dv2.layout)) + db_v_prod = cute.make_rmem_tensor(dv2_acc.layout, self.acc_dtype) + db_v_prod_mn = cute.make_tensor(db_v_prod.iterator, dv2_mn_layout) + partial_db_v_regs = cute.make_rmem_tensor(cute.make_layout((n_rows_dv2,)), self.acc_dtype) + partial_db_v_regs.fill(cutlass.Float32(0.0)) + + # Pre-load beta values for all rows this thread owns + beta_regs_dv2 = cute.make_rmem_tensor(cute.make_layout((n_rows_dv2,)), self.acc_dtype) + for i in cutlass.range_constexpr(n_rows_dv2): + row = coord_mn_dv2[i, 0][0] + beta_regs_dv2[i] = cutlass.Float32(sBeta[(row,)]) + + for v_iter in cutlass.range(self.num_v_tiles): + pipeline_load_dv.consumer_wait(load_dv_cs) + pipeline_load_tv.consumer_wait(load_tv_cs) + + # dv2 = A @ dv (complete per v_iter, K=BT) — first + dv2_acc.fill(0.0) + dv2_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for k_block_idx in cutlass.range_constexpr(num_k_blocks_dv2): + k_coord_dv2 = (None, None, k_block_idx, load_dv_cs.index) + cute.gemm( + dv2_tiled_mma, + dv2_acc, + tCrA_dv2[(None, None, k_block_idx)], + tCrDv_dv2[k_coord_dv2], + dv2_acc, + ) + warpgroup.commit_group() + + # dA += dv @ v^T (accumulates across v_iters) — second, overlaps with dv2 epi + dA_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for k_block_idx in cutlass.range_constexpr(num_k_blocks_dA): + k_coord_dv = (None, None, k_block_idx, load_dv_cs.index) + k_coord_tv = (None, None, k_block_idx, load_tv_cs.index) + cute.gemm( + dA_tiled_mma, + dA_acc, + tCrDv_dA[k_coord_dv], + tCrV_dA[k_coord_tv], + dA_acc, + ) + warpgroup.commit_group() + + # wait(1): dv2 done, dA still in flight + warpgroup.wait_group(1) + + # Retile dv2_acc to S2R layout → element j aligns with tSR_rV[j] + tSR_rAcc_dv2 = tiled_copy_s2r_v.retile(dv2_acc) + # Element-wise product in S2R layout, write to db_v_prod (MMA C layout) + # retile is a zero-cost view, same underlying registers + tSR_rProd = tiled_copy_s2r_v.retile(db_v_prod) + # ── db_v: sum_v dv2_acc[t,v] * v[t,v] before beta scaling ── + # Bulk load V via ldmatrix → register, then element-wise product + cute.copy( + tiled_copy_s2r_v, + tSR_sV[(None, None, None, load_tv_cs.index)], + tSR_rV, + ) + for j in cutlass.range_constexpr(cute.size(tSR_rProd)): + tSR_rProd[j] = tSR_rAcc_dv2[j] * cutlass.Float32(tSR_rV[j]) + + pipeline.NamedBarrier(barrier_id=BARRIER_DB_SYNC, num_threads=128).sync() + + for i in cutlass.range_constexpr(n_rows_dv2): + partial_db_v_regs[i] = partial_db_v_regs[i] + db_v_prod_mn[i, None].load().reduce( + cute.ReductionOp.ADD, cutlass.Float32.zero, 0 + ) + + # ── dv2 *= beta, then write to epi pipeline (overlaps with dA GEMM) ── + dv2_mn_view = cute.make_tensor(dv2_acc.iterator, dv2_mn_layout) + n_cols_dv2 = cute.size(dv2_mn_layout, mode=[1]) + for i in cutlass.range_constexpr(n_rows_dv2): + for j in cutlass.range_constexpr(n_cols_dv2): + dv2_mn_view[i, j] = dv2_mn_view[i, j] * beta_regs_dv2[i] + tRS_rAcc_dv2 = tiled_copy_r2s_bf16.retile(dv2_acc) + + for epi_idx in cutlass.range_constexpr(self.num_dv2_epi_tiles): + if wu_iter > 0 or v_iter > 0 or epi_idx >= self.epi_stage: + pipeline_epi_done.consumer_wait(epi_done_cs) + pipeline_epi_done.consumer_release(epi_done_cs) + epi_done_cs.advance() + + pipeline_epi_ready.producer_acquire(epi_ready_ps) + for epi_v in cutlass.range_constexpr(size_tRS_rD_dv2): + tRS_rD_dv2[epi_v] = cutlass.BFloat16(tRS_rAcc_dv2[epi_idx * size_tRS_rD_dv2 + epi_v]) + epi_buffer = epi_idx % cute.size(tRS_sEpi_dv2, mode=[3]) + cute.copy( + tiled_copy_r2s_bf16, + tRS_rD_dv2, + tRS_sEpi_dv2[(None, None, None, epi_buffer)], + ) + cute.arch.fence_view_async_shared() + pipeline_epi_ready.producer_commit(epi_ready_ps) + epi_ready_ps.advance() + + # wait(0): dA done + warpgroup.wait_group(0) + + # Release dv/tv after dA GEMM completes (it reads dv+v SMEM) + pipeline_load_dv.consumer_release(load_dv_cs) + pipeline_load_tv.consumer_release(load_tv_cs) + load_dv_cs.advance() + load_tv_cs.advance() + + # ── db_v writeback: warp reduction + write to sDb ── + reduction_target_dv2 = self._reduction_target_n(dv2_tiled_mma) + red_rank_dv2 = cute.rank(reduction_target_dv2) + for r_idx in cutlass.range_constexpr(red_rank_dv2): + for i in cutlass.range_constexpr(n_rows_dv2): + partial_db_v_regs[i] = cute.arch.warp_reduction_sum( + partial_db_v_regs[i], + threads_in_group=reduction_target_dv2.shape[r_idx], + ) + for i in cutlass.range_constexpr(n_rows_dv2): + if coord_mn_dv2[i, 0][1] == 0: + row = coord_mn_dv2[i, 0][0] + sDb[(row,)] = partial_db_v_regs[i] + pipeline.NamedBarrier(barrier_id=BARRIER_DB_SYNC, num_threads=128).sync() + + # ═══════════════════════════════════════════════ + # Unified k_iter: dq + dw + dk + dA(-=dw@kg) + # ═══════════════════════════════════════════════ + + # db_k register reduction setup (same pattern as db_v) + dk_mn_layout = self._layout_acc_mn(vloop_tiled_mma, dk_acc.layout) + n_rows_dk = cute.size(dk_mn_layout, mode=[0]) + thr_dk_per_thread = vloop_tiled_mma.get_slice(mma_tidx) + cD_dk_acc = cute.make_identity_tensor((BT, BK)) + tCcC_dk = thr_dk_per_thread.partition_C(cD_dk_acc) + coord_mn_dk = cute.make_tensor( + tCcC_dk.iterator, + self._layout_acc_mn(vloop_tiled_mma, tCcC_dk.layout), + ) + dbk_prod = cute.make_rmem_tensor(dk_acc.layout, self.acc_dtype) + dbk_prod_mn = cute.make_tensor(dbk_prod.iterator, dk_mn_layout) + partial_db_k_regs = cute.make_rmem_tensor(cute.make_layout((n_rows_dk,)), self.acc_dtype) + partial_db_k_regs.fill(cutlass.Float32(0.0)) + + for k_iter in cutlass.range(self.num_k_iters): + # ── dq = scale * exp2(g) * sum_v do @ h^T ── + dq_acc.fill(0.0) + vloop_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + + for v_iter in cutlass.range(self.num_v_tiles): + pipeline_load_tv.consumer_wait(load_tv_cs) + pipeline_load_h.consumer_wait(load_h_wait_cs) + + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_coord_do = (None, None, k_block_idx, load_tv_cs.index) + k_coord_h = (None, None, k_block_idx, load_h_wait_cs.index) + cute.gemm( + vloop_tiled_mma, + dq_acc, + tCrDo[k_coord_do], + tCrH[k_coord_h], + dq_acc, + ) + warpgroup.commit_group() + + if v_iter > 0: + warpgroup.wait_group(1) + + pipeline_load_tv.consumer_release(load_tv_cs) + # h NOT released here — reused by dw loop below + load_tv_cs.advance() + load_h_wait_cs.advance() + + # Wait for g/k TMA before WGMMA completes — ldmatrix overlaps with WGMMA. + # q is only needed for dg_part1 and is loaded in a separate pass + # after dq is gated. + pipeline_load_g.consumer_wait(load_g_cs) + pipeline_load_k.consumer_wait(load_k_cs) + + cute.copy(tiled_copy_s2r_g, tSG_sG, tSG_rG) + + warpgroup.wait_group(0) + + tRS_rAcc = tiled_copy_r2s_fp32.retile(dq_acc) + tRS_rG = tiled_copy_r2s_fp32.retile(tSG_rG) + cD = cute.make_identity_tensor((BT, BK)) + tRS_cD = thr_copy_r2s.partition_D(cD) + + if cutlass.const_expr(DEBUG_PRINT): + print("=== dq_acc / R2S layout analysis ===") + print(f"dq_acc.shape = {dq_acc.shape}") + print(f"dq_acc.layout = {dq_acc.layout}") + print(f"tRS_rAcc size = {cute.size(tRS_rAcc)}, shape = {tRS_rAcc.shape}") + print(f"size_tRS_rD (per epi_tile 64x32) = {size_tRS_rD}") + print(f"num_epi_tiles = {self.num_epi_tiles}") + print("--- coordinate mapping (compile-time) ---") + for j in cutlass.range_constexpr(cute.size(tRS_rAcc)): + print(f" j={j} epi={j // size_tRS_rD} sub={j % 8 // 4} row={tRS_cD[j][0]} col={tRS_cD[j][1]}") + print("--- q16 ldmatrix (64x16 chunked) ---") + print(f"tQ16_sQ.shape = {tQ16_sQ.shape}, layout = {tQ16_sQ.layout}") + print(f"tQ16_rQ.shape = {tQ16_rQ.shape}") + print(f"num_q16_tiles = {num_q16_tiles}") + print("=== end layout analysis ===") + + # Merged loop: kg, dq gate + cache exp2(g)/exp2(gn-g) + # All done in 64×16 chunks: k ldmatrix + kg stmatrix per chunk + k_exp_gn_g_regs = cute.make_rmem_tensor(tRS_rAcc.layout, self.acc_dtype) + tRS_rDbk_cache = tiled_copy_r2s_fp32.retile(dbk_prod) + # Ragged: chunk-last token is sub_seq_len-1, not BT-1. + gn_row = sub_seq_len - Int32(1) + + for tile16 in cutlass.range_constexpr(num_q16_tiles): + # Load 8 k values for this 64×16 tile. + cute.copy( + tiled_copy_s2r_q16, + tKG16_sK[(None, None, tile16)], + rK16, + ) + + for local_j in cutlass.range_constexpr(size_per_q16): + j = tile16 * size_per_q16 + local_j + c = tRS_cD[j][1] + g_val = cutlass.Float32(tRS_rG[j]) + k_val = cutlass.Float32(rK16[local_j]) + gn_val = cutlass.Float32(sG[(gn_row, c)]) + + exp_g = cute.math.exp2(g_val) + exp_gn_g = cute.math.exp2(gn_val - g_val) + + tRS_rG[j] = exp_g + tRS_rDbk_cache[j] = exp_gn_g + + rKg16[local_j] = cutlass.BFloat16(k_val * exp_g) + k_exp_gn_g_regs[j] = k_val * exp_gn_g + tRS_rAcc[j] = tRS_rAcc[j] * exp_g * self.scale + + # Store 8 kg values to sKg + cute.copy( + tiled_copy_r2s_kg16, + rKg16, + tKG16_sKg[(None, None, tile16)], + ) + + # k is no longer needed; q is consumed below for dg_part1. + pipeline_load_k.consumer_release(load_k_cs) + load_k_cs.advance() + + # Epilogue: write dq[BT, BK] — epi-tiles per k_iter + # tRS_rAcc already retiled and gated above + + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + pipeline_epi_done.consumer_wait(epi_done_cs) + pipeline_epi_done.consumer_release(epi_done_cs) + epi_done_cs.advance() + + self._write_epi_tile( + epi_idx, + tiled_copy_r2s_fp32, + tRS_rAcc, + tRS_sEpi, + size_tRS_rD, + tRS_rD, + pipeline_epi_ready, + epi_ready_ps, + ) + epi_ready_ps.advance() + + # ── Write dg_part1 → sDg staged, signal warp 3 for TMA store ── + # dg_part1 = q * dq. Compute it after dq is gated and write + # directly to the epilogue staging fragment, avoiding a full + # extra fp32 register tensor for dg_part1. + pipeline_load_q.consumer_wait(load_q_cs) + pipeline_dg_ready.producer_acquire(dg_ready_ps) + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + for tile16_sub in cutlass.range_constexpr(2): + tile16 = epi_idx * 2 + tile16_sub + cute.copy( + tiled_copy_s2r_q16, + tQ16_sQ[(None, None, tile16)], + tQ16_rQ, + ) + for local_j in cutlass.range_constexpr(size_per_q16): + epi_v = tile16_sub * size_per_q16 + local_j + j = epi_idx * size_tRS_rD + epi_v + q_val = cutlass.Float32(tQ16_rQ[local_j]) + tRS_rD[epi_v] = q_val * tRS_rAcc[j] + cute.copy( + tiled_copy_r2s_fp32, + tRS_rD, + tRS_sDg[(None, None, None, epi_idx, dg_ready_ps.index)], + ) + pipeline_load_q.consumer_release(load_q_cs) + load_q_cs.advance() + cute.arch.fence_view_async_shared() + pipeline_dg_ready.producer_commit(dg_ready_ps) + dg_ready_ps.advance() + + # ── dw = dv @ h (no g-scaling, result → sDw SMEM) ── + # h is still valid in SMEM from dq loop (not released) + dw_acc.fill(0.0) + vloop_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + + for v_iter in cutlass.range(self.num_v_tiles): + pipeline_load_dv.consumer_wait(load_dv_cs) + + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_coord_dv = (None, None, k_block_idx, load_dv_cs.index) + k_coord_h = (None, None, k_block_idx, load_h_release_cs.index) + cute.gemm( + vloop_tiled_mma, + dw_acc, + tCrDv[k_coord_dv], + tCrH[k_coord_h], + dw_acc, + ) + warpgroup.commit_group() + + if v_iter > 0: + warpgroup.wait_group(1) + + pipeline_load_dv.consumer_release(load_dv_cs) + pipeline_load_h.consumer_release(load_h_release_cs) + load_dv_cs.advance() + load_h_release_cs.advance() + + warpgroup.wait_group(0) + + # ── Write -dw (fp32→bf16, transposed) to sDw via stmatrix.trans ── + tRS_rAcc_dw = tiled_copy_r2s_dw.retile(dw_acc) + rDw_shape = cute.shape(thr_copy_r2s_dw.partition_S(sDw_write)) + tRS_rDw = cute.make_rmem_tensor_like(cute.make_layout(rDw_shape[:3]), self.io_dtype) + for idx in cutlass.range_constexpr(cute.size(tRS_rDw)): + tRS_rDw[idx] = cutlass.BFloat16(-tRS_rAcc_dw[idx]) + cute.copy(tiled_copy_r2s_dw, tRS_rDw, tRS_sDw) + cute.arch.fence_view_async_shared() + pipeline.NamedBarrier(barrier_id=BARRIER_DW_READY, num_threads=128).sync() + + # Convert dw_acc (C layout, fp32) → A operand (A layout, bf16, negated) + dw_as_a = self.make_acc_into_op(dw_acc, dwkg_tiled_mma, negate=True) + + # ── dA -= dw @ kg: accumulate into dA_acc ── + dwkg_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for k_block_idx in cutlass.range_constexpr(num_k_blocks_dwkg): + cute.gemm( + dwkg_tiled_mma, + dA_acc, + dw_as_a[(None, None, k_block_idx)], + tCrKg[(None, None, k_block_idx)], + dA_acc, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + + # ── dkgb = A^T @ (-dw): A loaded transposed via COL_MAJOR SMEM ── + dkgb_acc.fill(0.0) + dkgb_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for k_block_idx in cutlass.range_constexpr(num_k_blocks_dkgb): + cute.gemm( + dkgb_tiled_mma, + dkgb_acc, + tCrA_dkgb[(None, None, k_block_idx)], + tCrDw_dkgb[(None, None, k_block_idx)], + dkgb_acc, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + + # ── dk_inter = vnew @ dh ── + dk_acc.fill(0.0) + vloop_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + + for v_iter in cutlass.range(self.num_v_tiles): + pipeline_load_tv.consumer_wait(load_tv_cs) + pipeline_load_dh.consumer_wait(load_dh_cs) + + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_coord = (None, None, k_block_idx, load_tv_cs.index) + cute.gemm( + vloop_tiled_mma, + dk_acc, + tCrVnew[k_coord], + tCrDh[k_coord], + dk_acc, + ) + warpgroup.commit_group() + + if v_iter > 0: + warpgroup.wait_group(1) + + pipeline_load_tv.consumer_release(load_tv_cs) + pipeline_load_dh.consumer_release(load_dh_cs) + load_tv_cs.advance() + load_dh_cs.advance() + + cute.copy(tiled_copy_s2r_kg, tKG_sKg_load, tKG_rKg_load) + + warpgroup.wait_group(0) + + # dk = exp2(gn - g) * dk_inter + dkgb * exp2(g) * beta + tRS_rDk = tiled_copy_r2s_fp32.retile(dk_acc) + tRS_rDkgb = tiled_copy_r2s_fp32.retile(dkgb_acc) + cD_dk = cute.make_identity_tensor((BT, BK)) + tRS_cD_dk = thr_copy_r2s.partition_D(cD_dk) + + # Register arrays for dgk and db_k reductions + kdk_regs = cute.make_rmem_tensor(tRS_rDk.layout, self.acc_dtype) + dg_part2_regs = cute.make_rmem_tensor(tRS_rDk.layout, self.acc_dtype) + tRS_rDbk = tiled_copy_r2s_fp32.retile(dbk_prod) + tRS_rKg_ld = tiled_copy_r2s_fp32.retile(tKG_rKg_load) + + for j in cutlass.range_constexpr(cute.size(tRS_rDk)): + r = tRS_cD_dk[j][0] + c = tRS_cD_dk[j][1] + beta_val = cutlass.Float32(sBeta[(r,)]) + kg_val = cutlass.Float32(tRS_rKg_ld[j]) + dk_inter_j = tRS_rDk[j] + dkgb_j = tRS_rDkgb[j] + + exp_gn_g_j = tRS_rDbk[j] + exp_g_j = tRS_rG[j] + + kdk_regs[j] = k_exp_gn_g_regs[j] * dk_inter_j + tRS_rDbk[j] = dkgb_j * kg_val + + dg_part2_regs[j] = kg_val * dkgb_j * beta_val - k_exp_gn_g_regs[j] * dk_inter_j + + tRS_rDk[j] = exp_gn_g_j * dk_inter_j + dkgb_j * exp_g_j * beta_val + + # Ragged: zero OOB rows before any SMEM epilogue write. + # kdk needs this for the column reduction over BT rows; + # dg_part2 needs it so in-bounds tail chunks can still + # use full 64-row TMA reduce_add without polluting the + # next sequence. + if r >= sub_seq_len: + kdk_regs[j] = cutlass.Float32(0.0) + dg_part2_regs[j] = cutlass.Float32(0.0) + # Epilogue: write dk[BT, BK] + tRS_rAcc = tRS_rDk + + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + pipeline_epi_done.consumer_wait(epi_done_cs) + pipeline_epi_done.consumer_release(epi_done_cs) + epi_done_cs.advance() + + self._write_epi_tile( + epi_idx, + tiled_copy_r2s_fp32, + tRS_rAcc, + tRS_sEpi, + size_tRS_rD, + tRS_rD, + pipeline_epi_ready, + epi_ready_ps, + ) + epi_ready_ps.advance() + + # ── db_k: register row reduction (accumulate across k_iters) ── + for i in cutlass.range_constexpr(n_rows_dk): + partial_db_k_regs[i] = partial_db_k_regs[i] + dbk_prod_mn[i, None].load().reduce( + cute.ReductionOp.ADD, cutlass.Float32.zero, 0 + ) + + # Save gn before sG is overwritten (needed by dgk) + # ld128: each of 16 threads reads 4 consecutive gn values + my_gn = cute.make_rmem_tensor((4,), self.acc_dtype) + my_gn.fill(cutlass.Float32(0.0)) + if compute_tidx < Int32(BK // 4): + gn_col_base = compute_tidx * Int32(4) + my_gn.store(smem_load_f32x4_sw128(sG_raw_ptr, sub_seq_len - Int32(1), gn_col_base)) + + # ── dgk: m_last * (exp2(gn)*sum_v(h*dh) + sum_t(kdk)) ── + + # Write kdk_regs → sKdk via stmatrix (separate buffer) + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + for epi_v in cutlass.range_constexpr(size_tRS_rD): + tRS_rD[epi_v] = kdk_regs[epi_idx * size_tRS_rD + epi_v] + cute.copy( + tiled_copy_r2s_fp32, + tRS_rD, + tRS_sKdk_write[(None, None, None, epi_idx)], + ) + + cute.arch.fence_view_async_shared() + pipeline.NamedBarrier(barrier_id=BARRIER_DG_COMPUTE, num_threads=128).sync() + pipeline_dgk_hdh_ready.consumer_wait(dgk_hdh_ready_cs) + + # Column reduction via ld128: 16 threads × 4 cols each + # Compute dgk and write to sDgkHdh for broadcast + if compute_tidx < Int32(BK // 4): + col_base = compute_tidx * Int32(4) + dgk = cute.make_rmem_tensor((4,), self.acc_dtype) + dgk.fill(cutlass.Float32(0.0)) + for row in cutlass.range(BT, unroll_full=True): + vals = smem_load_f32x4_sw128(sKdk_raw_ptr, Int32(row), col_base) + for ci in cutlass.range_constexpr(4): + dgk[ci] = dgk[ci] + vals[ci] + hdh_off = k_iter * Int32(BK) + col_base + for ci in cutlass.range_constexpr(4): + hdh_val = cutlass.Float32(sDgkHdh[(hdh_off + Int32(ci),)]) + dgk[ci] = cute.math.exp2(my_gn[ci]) * hdh_val + dgk[ci] + sDgkHdh[(hdh_off + Int32(ci),)] = dgk[ci] + + cute.arch.fence_view_async_shared() + pipeline.NamedBarrier(barrier_id=BARRIER_DG_COMPUTE, num_threads=128).sync() + + # Acquire sDg stage BEFORE writing — ensures warp 3 + # has finished TMA-reading this stage from a prior round. + pipeline_dg_ready.producer_acquire(dg_ready_ps) + + # Write dg_part2 → sDg staged via stmatrix (without dgk) + tRS_rDgP2 = tiled_copy_r2s_fp32.retile(dg_part2_regs) + for epi_idx in cutlass.range_constexpr(self.num_epi_tiles): + for epi_v in cutlass.range_constexpr(size_tRS_rD): + tRS_rD[epi_v] = tRS_rDgP2[epi_idx * size_tRS_rD + epi_v] + cute.copy( + tiled_copy_r2s_fp32, + tRS_rD, + tRS_sDg[(None, None, None, epi_idx, dg_ready_ps.index)], + ) + + # Fence + barrier: ensure stmatrix of dg_part2 is visible + cute.arch.fence_view_async_shared() + pipeline.NamedBarrier(barrier_id=BARRIER_DG_COMPUTE, num_threads=128).sync() + + # 16 threads add dgk directly to sDg SMEM last row + # Ragged: last row = sub_seq_len - 1 (== BT-1 for full chunk). + dgk_row = sub_seq_len - Int32(1) + if compute_tidx < Int32(BK // 4): + col_base = compute_tidx * Int32(4) + for ci in cutlass.range_constexpr(4): + col = col_base + Int32(ci) + dgk_off = k_iter * Int32(BK) + col + old_val = cutlass.Float32(sDg_staged[(dgk_row, col, dg_ready_ps.index)]) + sDg_staged[(dgk_row, col, dg_ready_ps.index)] = old_val + cutlass.Float32(sDgkHdh[(dgk_off,)]) + + pipeline_dgk_hdh_ready.consumer_release(dgk_hdh_ready_cs) + dgk_hdh_ready_cs.advance() + + # ── dg_part2+dgk epilogue: signal warp 3 for TMA reduce_add ── + cute.arch.fence_view_async_shared() + pipeline_dg_ready.producer_commit(dg_ready_ps) + dg_ready_ps.advance() + + # Release g after dk/dg are done with it + pipeline_load_g.consumer_release(load_g_cs) + load_g_cs.advance() + + # ═══════════════════════════════════════════════ + # dA post-processing via WGMMA: + # dA_final = -lower_tri(sA @ (lower_tri(dA_raw * beta) @ sA)) + # GEMM 1: M @ sA (dwkg pattern: RMEM A, sA_row as B) + # GEMM 2: sA @ temp (dkgb pattern: sA as A, temp in sDw as B) + # ═══════════════════════════════════════════════ + + # Step 1: mask + beta in dA_acc + tRS_rAcc_dA = tiled_copy_r2s_dA_fp32.retile(dA_acc) + cD_dA = cute.make_identity_tensor((BT, BT)) + tRS_cD_dA = thr_copy_r2s_dA.partition_D(cD_dA) + + for j in cutlass.range_constexpr(cute.size(tRS_rAcc_dA)): + r = tRS_cD_dA[j][0] + c = tRS_cD_dA[j][1] + if r > c: + tRS_rAcc_dA[j] = tRS_rAcc_dA[j] * cutlass.Float32(sBeta[(c,)]) + else: + tRS_rAcc_dA[j] = cutlass.Float32(0.0) + + # Step 2: GEMM 1 — M @ sA_row via dwkg_tiled_mma + # Convert M (dA_acc, C layout) → A operand for dwkg. + m_as_a = self.make_acc_into_op(dA_acc, dwkg_tiled_mma) + dA_acc.fill(0.0) + dwkg_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for k_block_idx in cutlass.range_constexpr(num_k_blocks_post2): + cute.gemm( + dwkg_tiled_mma, + dA_acc, + m_as_a[(None, None, k_block_idx)], + tCrA_row_post[(None, None, k_block_idx)], + dA_acc, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + + # db_k writeback + db GMEM write — overlap with WGMMA GEMM 1 + reduction_target_dk = self._reduction_target_n(vloop_tiled_mma) + red_rank_dk = cute.rank(reduction_target_dk) + for r_idx in cutlass.range_constexpr(red_rank_dk): + for i in cutlass.range_constexpr(n_rows_dk): + partial_db_k_regs[i] = cute.arch.warp_reduction_sum( + partial_db_k_regs[i], + threads_in_group=reduction_target_dk.shape[r_idx], + ) + for i in cutlass.range_constexpr(n_rows_dk): + if coord_mn_dk[i, 0][1] == 0: + row = coord_mn_dk[i, 0][0] + sDb[(row,)] = cutlass.Float32(sDb[(row,)]) + partial_db_k_regs[i] + + # Ensure all warps' sDb writes are visible before reading + pipeline.NamedBarrier(barrier_id=BARRIER_DB_SYNC, num_threads=128).sync() + + # Ragged: only write db rows that belong to this sequence's chunk. + if compute_tidx < sub_seq_len: + db_gmem[(chunk_tok_offset + compute_tidx, (head_idx, Int32(0)))] = sDb[(compute_tidx,)] + + # Step 3: write temp (in dA_acc) → bf16 → sDw_wide via stmatrix.trans. + # temp becomes the K-major B operand for GEMM 2. + tRS_rAcc_dA_dw = tiled_copy_r2s_dA_dw.retile(dA_acc) + rM_shape = cute.shape(thr_copy_r2s_dA_dw.partition_S(sDw_write_wide)) + tRS_rM = cute.make_rmem_tensor_like(cute.make_layout(rM_shape[:3]), self.io_dtype) + for idx in cutlass.range_constexpr(cute.size(tRS_rM)): + tRS_rM[idx] = cutlass.BFloat16(tRS_rAcc_dA_dw[idx]) + cute.copy(tiled_copy_r2s_dA_dw, tRS_rM, tRS_sDw_wide) + cute.arch.fence_view_async_shared() + pipeline.NamedBarrier(barrier_id=BARRIER_DB_SYNC, num_threads=128).sync() + + # Step 4: GEMM 2 — sA @ temp → dA_acc (always m64n64) + dA_acc.fill(0.0) + dA_post1_tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for k_block_idx in cutlass.range_constexpr(num_k_blocks_post1): + cute.gemm( + dA_post1_tiled_mma, + dA_acc, + tCrA_post1[(None, None, k_block_idx)], + tCrDw_post1[(None, None, k_block_idx)], + dA_acc, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + + # Step 5: mask + negate + tRS_rAcc_dA = tiled_copy_r2s_dA_fp32.retile(dA_acc) + cD_dA = cute.make_identity_tensor((BT, BT)) + tRS_cD_dA = thr_copy_r2s_dA.partition_D(cD_dA) + + for j in cutlass.range_constexpr(cute.size(tRS_rAcc_dA)): + r = tRS_cD_dA[j][0] + c = tRS_cD_dA[j][1] + if r > c: + tRS_rAcc_dA[j] = -tRS_rAcc_dA[j] + else: + tRS_rAcc_dA[j] = cutlass.Float32(0.0) + + # Release A and beta after post-processing + pipeline_load_A.consumer_release(load_A_cs) + load_A_cs.advance() + pipeline_load_beta.consumer_release(load_beta_cs) + load_beta_cs.advance() + + # ═══════════════════════════════════════════════ + # dA epilogue: write dA (2 epi-tiles) — after all k_iters + # ═══════════════════════════════════════════════ + tRS_rAcc = tiled_copy_r2s_dA_fp32.retile(dA_acc) + + for epi_idx in cutlass.range_constexpr(self.num_dA_epi_tiles): + pipeline_epi_done.consumer_wait(epi_done_cs) + pipeline_epi_done.consumer_release(epi_done_cs) + epi_done_cs.advance() + + self._write_epi_tile( + epi_idx, + tiled_copy_r2s_dA_fp32, + tRS_rAcc, + tRS_sEpi_dA, + size_tRS_rD_dA, + tRS_rD_dA, + pipeline_epi_ready, + epi_ready_ps, + ) + epi_ready_ps.advance() + + return + + +# ===================================================================== +# Compilation cache +# ===================================================================== + +_bwd_wy_kernel_cache: dict = {} + + +def _compile_bwd_wy_variant( + H: int, + K: int, + V: int, + scale: float, + chunk_size: int, + beta_dtype: type[cutlass.Numeric], + use_fast_math: bool, + bk: int = 32, + bv: int = 64, + min_occupancy: int = 2, +): + kernel_obj = ChunkKdaBwdWyDqkgFusedSM90( + chunk_size=chunk_size, + head_dim_k=K, + head_dim_v=V, + scale=scale, + use_fast_math=use_fast_math, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + + sym_b = cute.sym_int() + sym_nt = cute.sym_int() + sym_cu = cute.sym_int() + sym_ci = cute.sym_int() + + do_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + h_fake = make_fake_compact_tensor( + cutlass.BFloat16, + (1, sym_nt, H, K, V), + stride_order=(4, 3, 2, 1, 0), + assumed_align=128, + ) + vnew_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + dh_fake = make_fake_compact_tensor( + cutlass.BFloat16, + (1, sym_nt, H, K, V), + stride_order=(4, 3, 2, 1, 0), + assumed_align=128, + ) + dq_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + dk_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + cu_fake = make_fake_compact_tensor(cutlass.Int32, (sym_cu,), assumed_align=128) + ci_fake = make_fake_compact_tensor(cutlass.Int32, (sym_ci, 2), stride_order=(1, 0), assumed_align=128) + stream_fake = make_fake_stream(use_tvm_ffi_env_stream=True) + + g_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + q_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + k_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + dv_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + v_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + A_fake = make_fake_compact_tensor( + cutlass.BFloat16, (1, sym_b, H, chunk_size), stride_order=(3, 2, 1, 0), assumed_align=128 + ) + dA_fake = make_fake_compact_tensor( + cutlass.Float32, (1, sym_b, H, chunk_size), stride_order=(3, 2, 1, 0), assumed_align=128 + ) + dv2_fake = make_fake_compact_tensor(cutlass.BFloat16, (1, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + db_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, H), stride_order=(2, 1, 0), assumed_align=128) + beta_fake = make_fake_compact_tensor(beta_dtype, (1, sym_b, H), stride_order=(2, 1, 0), assumed_align=128) + + dg_fake = make_fake_compact_tensor(cutlass.Float32, (1, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + compiled_fn = cute.compile( + kernel_obj, + do_fake, + h_fake, + vnew_fake, + dh_fake, + g_fake, + q_fake, + k_fake, + dq_fake, + dk_fake, + dg_fake, + dv_fake, + v_fake, + A_fake, + dA_fake, + dv2_fake, + db_fake, + beta_fake, + cu_fake, + ci_fake, + (Int32(1), Int32(1), Int32(H), Int32(K), Int32(V)), + Int32(1), + stream_fake, + options=COMPILE_OPTIONS, + ) + return compiled_fn + + +def _get_compiled_bwd_wy( + H: int, + K: int, + V: int, + scale: float, + chunk_size: int, + beta_dtype: torch.dtype, + bk: int = 32, + bv: int = 64, + min_occupancy: int = 2, +): + key = (H, K, V, scale, chunk_size, beta_dtype, USE_FAST_MATH, bk, bv, min_occupancy) + if key not in _bwd_wy_kernel_cache: + _bwd_wy_kernel_cache[key] = _compile_bwd_wy_variant( + H, + K, + V, + scale, + chunk_size, + _torch_to_cutlass_dtype[beta_dtype], + USE_FAST_MATH, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + return _bwd_wy_kernel_cache[key] + + +# ===================================================================== +# Public Python wrapper — FLA-compatible +# ===================================================================== + + +def chunk_kda_bwd_wy_dqkg_fused( + 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: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.Tensor | None = None, + *, + bk: int = 32, + bv: int = 64, + min_occupancy: int = 2, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """SM90 wrapper for the WY dq+kg fused backward kernel. + + Returns: + (dq, dk, dv, db, dg, dA), matching FLA's output order. + """ + from fla.ops.utils.index import prepare_chunk_indices + + from cula.utils import prepare_uniform_cu_seqlens + + B, T, H, K = q.shape + V = v.shape[-1] + BT = chunk_size + device = q.device + + if scale is None: + scale = K**-0.5 + + if cu_seqlens is None: + cu_seqlens = prepare_uniform_cu_seqlens(B, T, device, torch.int32) + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + + assert cu_seqlens.dtype == torch.int32 + assert do.dtype == torch.bfloat16 + assert h.dtype == torch.bfloat16 + assert g.dtype == torch.float32 + assert q.dtype == torch.bfloat16 + assert k.dtype == torch.bfloat16 + assert A.dtype == torch.bfloat16 + assert beta.dtype in _torch_to_cutlass_dtype, f"SM90 kernel only supports fp32/bf16 beta, got {beta.dtype}" + + T_total = B * T + num_seqs = cu_seqlens.shape[0] - 1 + total_nt_val = chunk_indices.shape[0] + ps = (Int32(num_seqs), Int32(T_total), Int32(H), Int32(K), Int32(V)) + + dq = torch.empty(1, T_total, H, K, dtype=torch.float32, device=device) + dk = torch.empty(1, T_total, H, K, dtype=torch.float32, device=device) + dv_out = torch.empty(1, T_total, H, V, dtype=torch.bfloat16, device=device) + db = torch.empty(1, T_total, H, dtype=torch.float32, device=device) + dg = torch.empty(1, T_total, H, K, dtype=torch.float32, device=device) + dA = torch.empty(1, T_total, H, BT, dtype=torch.float32, device=device) + + if B != 1: + do = do.reshape(1, T_total, H, V) + h = h.reshape(1, total_nt_val, H, K, V) + g = g.reshape(1, T_total, H, K) + q = q.reshape(1, T_total, H, K) + k = k.reshape(1, T_total, H, K) + v_new = v_new.reshape(1, T_total, H, V) + dh = dh.reshape(1, total_nt_val, H, K, V) + dv = dv.reshape(1, T_total, H, V) + v = v.reshape(1, T_total, H, V) + A = A.reshape(1, T_total, H, BT) + beta = beta.reshape(1, T_total, H) + + compiled_fn = _get_compiled_bwd_wy(H, K, V, scale, chunk_size, beta.dtype, bk=bk, bv=bv, min_occupancy=min_occupancy) + + compiled_fn( + do, + h, + v_new, + dh, + g, + q, + k, + dq, + dk, + dg, + dv, + v, + A, + dA, + dv_out, + db, + beta, + cu_seqlens, + chunk_indices, + ps, + Int32(total_nt_val), + ) + + if B != 1: + dq = dq.reshape(B, T, H, K) + dk = dk.reshape(B, T, H, K) + dv_out = dv_out.reshape(B, T, H, V) + db = db.reshape(B, T, H) + dg = dg.reshape(B, T, H, K) + dA = dA.reshape(B, T, H, BT) + + return dq, dk, dv_out, db, dg, dA diff --git a/tests/conftest.py b/tests/conftest.py index a9338aca..cd90e7af 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,9 +8,15 @@ def _is_sm100() -> bool: return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 +def _is_sm90() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9 + + def pytest_configure(config): config.addinivalue_line("markers", "sm100_only: only run on SM100 devices") - config.addinivalue_line("markers", "sm90_only: skip on SM100 devices") + config.addinivalue_line("markers", "sm90_only: only run on SM90 devices") + config.addinivalue_line("markers", "benchmark: long-running benchmark-shaped coverage") + config.addinivalue_line("markers", "sanitizer: tests intended to run under compute-sanitizer") config.addinivalue_line( "markers", "kda_fast: KDA test case included in fast (default) mode", @@ -33,8 +39,9 @@ def pytest_configure(config): def pytest_collection_modifyitems(config, items): is_sm100 = _is_sm100() + is_sm90 = _is_sm90() skip_non_sm100 = pytest.mark.skip(reason="SM100-only test: skip on non-SM100 devices") - skip_on_sm100 = pytest.mark.skip(reason="SM90-only test: skip on SM100") + skip_non_sm90 = pytest.mark.skip(reason="SM90-only test: skip on non-SM90 devices") marker_expr = config.option.markexpr or "" include_slow = "kda_slow" in marker_expr @@ -49,8 +56,8 @@ def pytest_collection_modifyitems(config, items): for item in items: if "sm100_only" in item.keywords and not is_sm100: item.add_marker(skip_non_sm100) - if "sm90_only" in item.keywords and is_sm100: - item.add_marker(skip_on_sm100) + if "sm90_only" in item.keywords and not is_sm90: + item.add_marker(skip_non_sm90) if include_slow: continue if "kda_slow" in item.keywords: diff --git a/tests/test_chunk_wy_dqkg_sm90.py b/tests/test_chunk_wy_dqkg_sm90.py new file mode 100644 index 00000000..a7b97c9f --- /dev/null +++ b/tests/test_chunk_wy_dqkg_sm90.py @@ -0,0 +1,689 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical tests for the SM90 fused WY-DqKG backward kernel. + +The tests compare SM90 fused outputs against the FLA Triton fused kernel. +Both fixed-length and ragged varlen partial-chunk paths are covered. +""" + +import os + +import torch +from fla.ops.kda.chunk_bwd import chunk_kda_bwd_wy_dqkg_fused as chunk_kda_bwd_triton + +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + BK as BENCHMARK_BK, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + BT as BENCHMARK_BT, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + BV as BENCHMARK_BV, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + MIN_OCC as BENCHMARK_MIN_OCC, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + SEED as BENCHMARK_SEED, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + K as BENCHMARK_K, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + V as BENCHMARK_V, +) +from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( + benchmark_fixed_configs, + benchmark_varlen_configs, + prepare_bwd_wy_dqkg_fused_inputs, +) +from benchmarks.utils import exclusive_cumsum +from cula.ops.chunk_wy_dqkg_sm90 import chunk_kda_bwd_wy_dqkg_fused + +# pytest is optional — when absent (e.g. minimal CI venv) we still allow +# the module to be executed directly via ``python tests/test_*.py``. +try: + import pytest + + _HAS_PYTEST = True +except ImportError: + _HAS_PYTEST = False + + class _DummyMark: + def __getattr__(self, _name): + return lambda *a, **kw: lambda f: f + + class _DummyPytest: + mark = _DummyMark() + + @staticmethod + def main(*_a, **_kw): + raise SystemExit("pytest not installed; run via __main__ instead") + + pytest = _DummyPytest() # type: ignore[assignment] + + +OUT_NAMES = ("dq", "dk", "dv", "db", "dg", "dA") + + +def _env_int(name, default): + value = os.environ.get(name) + return default if value is None else int(value) + + +def _benchmark_test_heads(): + heads = os.environ.get("CULA_BENCHMARK_TEST_HEADS") + if heads is None: + return (32,) + return tuple(int(h.strip()) for h in heads.split(",") if h.strip()) + + +def _benchmark_fixed_test_cases(): + return [(H, B, T) for H in _benchmark_test_heads() for B, T in benchmark_fixed_configs()] + + +def _benchmark_varlen_test_cases(): + return [ + (H, seq_lens, total_len, dist) + for H in _benchmark_test_heads() + for seq_lens, total_len, dist in benchmark_varlen_configs() + ] + + +def _determinism_fixed_test_cases(): + min_t = _env_int("CULA_DETERMINISM_FIXED_MIN_T", 16384) + return [(H, B, T) for H, B, T in _benchmark_fixed_test_cases() if min_t <= T] + + +def _determinism_varlen_test_cases(): + min_total_len = _env_int("CULA_DETERMINISM_VARLEN_MIN_TOTAL_LEN", 16384) + return [ + (H, seq_lens, total_len, dist) + for H, seq_lens, total_len, dist in _benchmark_varlen_test_cases() + if total_len >= min_total_len + ] + + +def _fixed_case_id(case): + H, B, T = case + return f"H{H}-B{B}-T{T}" + + +def _varlen_case_id(case): + H, seq_lens, total_len, dist = case + return f"H{H}-{dist}-{len(seq_lens)}seqs-T{total_len}-min{min(seq_lens)}-max{max(seq_lens)}" + + +BENCHMARK_FIXED_TEST_CASES = _benchmark_fixed_test_cases() +BENCHMARK_VARLEN_TEST_CASES = _benchmark_varlen_test_cases() +DETERMINISM_FIXED_TEST_CASES = _determinism_fixed_test_cases() +DETERMINISM_VARLEN_TEST_CASES = _determinism_varlen_test_cases() + + +def accuracy_stats(ref, out): + """Compute err_ratio, relative max diff, and mean absolute difference.""" + ref_f = ref.float() + out_f = out.float() + diff = (ref_f - out_f).abs() + err = diff.flatten().pow(2).mean().sqrt().item() + base = ref_f.flatten().pow(2).mean().sqrt().item() + err_ratio = err / (base + 1e-8) + max_diff = diff.max().item() + denom = ref_f.abs().max().item() + rel_max = max_diff / denom if denom > 0 else 0.0 + mean_diff = diff.mean().item() + return err_ratio, rel_max, mean_diff + + +def _print_tensor_stats(name, tensor, ref_tensor=None): + """Print diagnostics for a tensor, optionally vs a reference.""" + f = tensor.float() + print(f"{name} norm={f.norm().item():.6e} min={f.min().item():.6e} max={f.max().item():.6e}", flush=True) + if ref_tensor is not None: + err_ratio, rel_max, mean_diff = accuracy_stats(ref_tensor, tensor) + print(f" vs ref: err_ratio={err_ratio:.6f} rel_max={rel_max:.6f} mean_diff={mean_diff:.6e}", flush=True) + + +def _assert_outputs_match_fla(sm90_outputs, fla_outputs, case_id, *, max_err=0.05, verbose=False): + for name, sm90, fla in zip(OUT_NAMES, sm90_outputs, fla_outputs): + assert sm90.shape == fla.shape, f"{case_id}: {name} shape sm90={tuple(sm90.shape)} fla={tuple(fla.shape)}" + err_ratio, rel_max, mean_diff = accuracy_stats(fla, sm90) + if verbose: + print( + f"{case_id} {name}: err_ratio={err_ratio:.6f} rel_max={rel_max:.6f} mean_diff={mean_diff:.6e}", + flush=True, + ) + _print_tensor_stats(f"{name}_sm90", sm90, fla) + assert err_ratio < max_err, f"{case_id}: {name} vs FLA err_ratio={err_ratio:.6f} too high" + + +def _run_matches_fla_fixed( + B=1, + T=64, + H=4, + K=128, + V=128, + BT=64, + verbose=False, + bk=32, + bv=64, + min_occupancy=2, + beta_dtype=torch.float32, +): + """Verify SM90 fused outputs on fixed-length inputs against FLA.""" + device = "cuda" + dtype = torch.bfloat16 + scale = K**-0.5 + NT = T // BT + + torch.manual_seed(0) + do_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + h_tensor = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 + g_tensor = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 + q_tensor = torch.randn(B, T, H, K, dtype=dtype, device=device) + k_tensor = torch.randn(B, T, H, K, dtype=dtype, device=device) + vnew_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + dh_tensor = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 + dv_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + v_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + A_tensor = torch.randn(B, T, H, BT, dtype=dtype, device=device) + beta_tensor = (torch.rand(B, T, H, dtype=torch.float32, device=device) * 0.5 + 0.5).to(beta_dtype) + + if verbose: + print(f"=== inputs: B={B} T={T} H={H} K={K} V={V} BT={BT} NT={NT} beta={beta_dtype} ===", flush=True) + + if verbose: + print("=== invoking FLA Triton baseline ===", flush=True) + fla_outputs = chunk_kda_bwd_triton( + q=q_tensor, + k=k_tensor, + v=v_tensor, + v_new=vnew_tensor, + g=g_tensor, + beta=beta_tensor, + A=A_tensor, + h=h_tensor, + do=do_tensor, + dh=dh_tensor, + dv=dv_tensor, + scale=scale, + chunk_size=BT, + ) + torch.cuda.synchronize() + + if verbose: + print(f"A[0,0,0,:4]={A_tensor[0, 0, 0, :4].tolist()}", flush=True) + print(f"A[0,1,0,:4]={A_tensor[0, 1, 0, :4].tolist()}", flush=True) + print(f"A[0,:4,0,0]={A_tensor[0, :4, 0, 0].tolist()}", flush=True) + print(f"A[0,:4,0,1]={A_tensor[0, :4, 0, 1].tolist()}", flush=True) + + if verbose: + print("\n=== invoking SM90 fused wrapper ===", flush=True) + sm90_outputs = chunk_kda_bwd_wy_dqkg_fused( + q=q_tensor, + k=k_tensor, + v=v_tensor, + v_new=vnew_tensor, + g=g_tensor, + beta=beta_tensor, + A=A_tensor, + h=h_tensor, + do=do_tensor, + dh=dh_tensor, + dv=dv_tensor, + scale=scale, + chunk_size=BT, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + torch.cuda.synchronize() + + _assert_outputs_match_fla(sm90_outputs, fla_outputs, f"fixed B={B} T={T} H={H}", verbose=verbose) + return sm90_outputs + + +@pytest.mark.sm90_only +@pytest.mark.parametrize("B, T, H, K, V, BT", [(1, 64, 4, 128, 128, 64)]) +@pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16], ids=["beta_fp32", "beta_bf16"]) +def test_matches_fla_fixed(B, T, H, K, V, BT, beta_dtype): + _run_matches_fla_fixed(B, T, H, K, V, BT, beta_dtype=beta_dtype) + + +def _run_matches_fla_uniform_varlen( + B=2, + T=64, + H=4, + K=128, + V=128, + BT=64, + verbose=False, + bk=32, + bv=64, + min_occupancy=2, +): + """Verify SM90 fused varlen path on uniform cu_seqlens (= prepare_uniform_cu_seqlens). + + Uses the explicit cu_seqlens / chunk_indices code path through the wrapper: + feeds reshape-to-[1, B*T, ...] tensors plus cu_seqlens=[0, T, 2T, ..., B*T]. + """ + from fla.ops.utils.index import prepare_chunk_indices + + from cula.utils import prepare_uniform_cu_seqlens + + device = "cuda" + dtype = torch.bfloat16 + scale = K**-0.5 + NT = T // BT + + torch.manual_seed(0) + do_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + h_tensor = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 + g_tensor = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 + q_tensor = torch.randn(B, T, H, K, dtype=dtype, device=device) + k_tensor = torch.randn(B, T, H, K, dtype=dtype, device=device) + vnew_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + dh_tensor = torch.randn(B, NT, H, K, V, dtype=dtype, device=device) * 0.01 + dv_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + v_tensor = torch.randn(B, T, H, V, dtype=dtype, device=device) + A_tensor = torch.randn(B, T, H, BT, dtype=dtype, device=device) + beta_tensor = torch.rand(B, T, H, dtype=torch.float32, device=device) * 0.5 + 0.5 + + cu_seqlens = prepare_uniform_cu_seqlens(B, T, device, torch.int32) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + + fla_outputs = chunk_kda_bwd_triton( + q=q_tensor, + k=k_tensor, + v=v_tensor, + v_new=vnew_tensor, + g=g_tensor, + beta=beta_tensor, + A=A_tensor, + h=h_tensor, + do=do_tensor, + dh=dh_tensor, + dv=dv_tensor, + scale=scale, + chunk_size=BT, + ) + + sm90_outputs = chunk_kda_bwd_wy_dqkg_fused( + q=q_tensor, + k=k_tensor, + v=v_tensor, + v_new=vnew_tensor, + g=g_tensor, + beta=beta_tensor, + A=A_tensor, + h=h_tensor, + do=do_tensor, + dh=dh_tensor, + dv=dv_tensor, + scale=scale, + chunk_size=BT, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + torch.cuda.synchronize() + + if verbose: + print(f"=== uniform varlen B={B} T={T} (T_total={B * T}) ===", flush=True) + _assert_outputs_match_fla(sm90_outputs, fla_outputs, f"uniform varlen B={B} T={T} H={H}", verbose=verbose) + return sm90_outputs + + +@pytest.mark.sm90_only +@pytest.mark.parametrize("B, T, H, K, V, BT", [(2, 64, 4, 128, 128, 64)]) +def test_matches_fla_uniform_varlen(B, T, H, K, V, BT): + _run_matches_fla_uniform_varlen(B, T, H, K, V, BT) + + +def _run_matches_fla_ragged_varlen( + seq_lens, + H=4, + K=128, + V=128, + BT=64, + verbose=False, + bk=32, + bv=64, + min_occupancy=2, +): + """Verify SM90 partial-chunk row mask on ragged cu_seqlens. + + All 6 outputs (dq, dk, dv, db, dg, dA) asserted per-sequence against + the Triton FLA reference (which supports varlen via cu_seqlens). + """ + import itertools + + from fla.ops.utils.index import prepare_chunk_indices + + device = "cuda" + dtype = torch.bfloat16 + scale = K**-0.5 + + T_total = sum(seq_lens) + cu_seqlens = torch.tensor( + [0] + list(itertools.accumulate(seq_lens)), + dtype=torch.int32, + device=device, + ) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT_total = chunk_indices.shape[0] + + torch.manual_seed(0) + do_tensor = torch.randn(1, T_total, H, V, dtype=dtype, device=device) + h_tensor = torch.randn(1, NT_total, H, K, V, dtype=dtype, device=device) * 0.01 + g_tensor = torch.randn(1, T_total, H, K, dtype=torch.float32, device=device) * 0.1 + q_tensor = torch.randn(1, T_total, H, K, dtype=dtype, device=device) + k_tensor = torch.randn(1, T_total, H, K, dtype=dtype, device=device) + vnew_tensor = torch.randn(1, T_total, H, V, dtype=dtype, device=device) + dh_tensor = torch.randn(1, NT_total, H, K, V, dtype=dtype, device=device) * 0.01 + dv_tensor = torch.randn(1, T_total, H, V, dtype=dtype, device=device) + v_tensor = torch.randn(1, T_total, H, V, dtype=dtype, device=device) + A_tensor = torch.randn(1, T_total, H, BT, dtype=dtype, device=device) + beta_tensor = torch.rand(1, T_total, H, dtype=torch.float32, device=device) * 0.5 + 0.5 + + # FLA Triton reference (supports varlen via cu_seqlens) + dq_fla, dk_fla, dv_fla, db_fla, dg_fla, dA_fla = chunk_kda_bwd_triton( + q=q_tensor, + k=k_tensor, + v=v_tensor, + v_new=vnew_tensor, + g=g_tensor, + beta=beta_tensor, + A=A_tensor, + h=h_tensor, + do=do_tensor, + dh=dh_tensor, + dv=dv_tensor, + scale=scale, + chunk_size=BT, + cu_seqlens=cu_seqlens, + ) + + dq_sm90, dk_sm90, dv_sm90, db_sm90, dg_sm90, dA_sm90 = chunk_kda_bwd_wy_dqkg_fused( + q=q_tensor, + k=k_tensor, + v=v_tensor, + v_new=vnew_tensor, + g=g_tensor, + beta=beta_tensor, + A=A_tensor, + h=h_tensor, + do=do_tensor, + dh=dh_tensor, + dv=dv_tensor, + scale=scale, + chunk_size=BT, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + torch.cuda.synchronize() + + if verbose: + print(f"=== ragged seq_lens={seq_lens} T_total={T_total} NT_total={NT_total} ===", flush=True) + print(f"chunk_indices = {chunk_indices.tolist()}", flush=True) + + names_pairs = [ + ("dq", dq_sm90, dq_fla), + ("dk", dk_sm90, dk_fla), + ("dv", dv_sm90, dv_fla), + ("db", db_sm90, db_fla), + ("dg", dg_sm90, dg_fla), + ("dA", dA_sm90, dA_fla), + ] + # Per-seq comparison for each of the 6 outputs + for b in range(len(seq_lens)): + start = cu_seqlens[b].item() + end = cu_seqlens[b + 1].item() + for name, sm90, ref in names_pairs: + sm90_slice = sm90[0, start:end] + ref_slice = ref[0, start:end] + err_ratio, rel_max, mean_diff = accuracy_stats(ref_slice, sm90_slice) + if verbose: + print( + f" seq[{b}] [{start},{end}) {name:>3}: " + f"err_ratio={err_ratio:.6f} rel_max={rel_max:.6f} mean_diff={mean_diff:.6e}", + flush=True, + ) + assert err_ratio < 0.003, f"seq[{b}] {name} err_ratio={err_ratio:.6f} too high" + + return dq_sm90 + + +@pytest.mark.sm90_only +@pytest.mark.parametrize("seq_lens", [[64, 96], [48, 112]]) +def test_matches_fla_ragged_varlen(seq_lens): + _run_matches_fla_ragged_varlen(seq_lens) + + +def _prepare_benchmark_fixed_inputs(B, T, H, K=BENCHMARK_K, V=BENCHMARK_V, BT=BENCHMARK_BT): + device = torch.device("cuda") + cu_seqlens = torch.tensor(exclusive_cumsum([T] * B), dtype=torch.int32, device=device) + return prepare_bwd_wy_dqkg_fused_inputs( + B=B, + T=T, + H=H, + K=K, + V=V, + chunk_size=BT, + device=device, + seed=BENCHMARK_SEED, + cu_seqlens=cu_seqlens, + ) + + +def _prepare_benchmark_varlen_inputs(seq_lens, total_len, H, K=BENCHMARK_K, V=BENCHMARK_V, BT=BENCHMARK_BT): + device = torch.device("cuda") + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + return prepare_bwd_wy_dqkg_fused_inputs( + B=1, + T=total_len, + H=H, + K=K, + V=V, + chunk_size=BT, + device=device, + seed=BENCHMARK_SEED, + cu_seqlens=cu_seqlens, + ) + + +def _run_sm90_from_benchmark_inputs( + inputs, + *, + chunk_size=BENCHMARK_BT, + bk=BENCHMARK_BK, + bv=BENCHMARK_BV, + min_occupancy=BENCHMARK_MIN_OCC, +): + outputs = chunk_kda_bwd_wy_dqkg_fused( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + v_new=inputs["v_new"], + g=inputs["g"], + beta=inputs["beta"], + A=inputs["A"], + h=inputs["h"], + do=inputs["do"], + dh=inputs["dh"], + dv=inputs["dv"], + scale=inputs["scale"], + cu_seqlens=inputs["cu_seqlens"], + chunk_size=chunk_size, + chunk_indices=inputs["chunk_indices"], + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + torch.cuda.synchronize() + return outputs + + +def _run_benchmark_input_determinism( + inputs, + case_id, + *, + iters=None, + chunk_size=BENCHMARK_BT, + bk=BENCHMARK_BK, + bv=BENCHMARK_BV, + min_occupancy=BENCHMARK_MIN_OCC, +): + """Multiple SM90 calls on one benchmark-shaped input must be bitwise identical.""" + if iters is None: + iters = _env_int("CULA_DETERMINISM_ITERS", 2) + + ref = tuple( + out.clone() + for out in _run_sm90_from_benchmark_inputs( + inputs, + chunk_size=chunk_size, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + ) + for i in range(iters): + actual = _run_sm90_from_benchmark_inputs( + inputs, + chunk_size=chunk_size, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + for name, got, expected in zip(OUT_NAMES, actual, ref): + assert torch.isfinite(got.float()).all(), f"{case_id}: {name} has non-finite values at iter {i}" + assert torch.equal(got, expected), f"{case_id}: non-deterministic {name} at iter {i}" + + +def _run_determinism(B=1, T=64, H=4, K=128, V=128, BT=64, iters=2, bk=32, bv=64, min_occupancy=2): + """Compatibility entry point for focused sanitizer/debug commands.""" + inputs = _prepare_benchmark_fixed_inputs(B, T, H, K=K, V=V, BT=BT) + _run_benchmark_input_determinism( + inputs, + f"fixed-H{H}-B{B}-T{T}", + iters=iters, + chunk_size=BT, + bk=bk, + bv=bv, + min_occupancy=min_occupancy, + ) + + +@pytest.mark.benchmark +@pytest.mark.sm90_only +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + "H, B, T", + DETERMINISM_FIXED_TEST_CASES, + ids=[_fixed_case_id(case) for case in DETERMINISM_FIXED_TEST_CASES], +) +def test_determinism_benchmark_fixed_cases(H, B, T): + inputs = _prepare_benchmark_fixed_inputs(B, T, H) + _run_benchmark_input_determinism(inputs, f"fixed-H{H}-B{B}-T{T}") + torch.cuda.empty_cache() + + +@pytest.mark.benchmark +@pytest.mark.sm90_only +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + "H, seq_lens, total_len, dist", + DETERMINISM_VARLEN_TEST_CASES, + ids=[_varlen_case_id(case) for case in DETERMINISM_VARLEN_TEST_CASES], +) +def test_determinism_benchmark_varlen_cases(H, seq_lens, total_len, dist): + inputs = _prepare_benchmark_varlen_inputs(seq_lens, total_len, H) + case_id = _varlen_case_id((H, seq_lens, total_len, dist)) + _run_benchmark_input_determinism(inputs, case_id) + torch.cuda.empty_cache() + + +def _run_benchmark_sanitizer_case(inputs, chunk_size=BENCHMARK_BT): + outputs = _run_sm90_from_benchmark_inputs(inputs, chunk_size=chunk_size) + _, total_len, H, K = inputs["q"].shape + V = inputs["v"].shape[-1] + expected_shapes = ( + (1, total_len, H, K), + (1, total_len, H, K), + (1, total_len, H, V), + (1, total_len, H), + (1, total_len, H, K), + (1, total_len, H, chunk_size), + ) + for out, expected_shape in zip(outputs, expected_shapes): + assert tuple(out.shape) == expected_shape + + +@pytest.mark.sanitizer +@pytest.mark.benchmark +@pytest.mark.sm90_only +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.skipif( + os.environ.get("CULA_RUN_SANITIZER_TESTS") != "1", + reason="set CULA_RUN_SANITIZER_TESTS=1 and run under compute-sanitizer", +) +@pytest.mark.parametrize( + "H, B, T", + BENCHMARK_FIXED_TEST_CASES, + ids=[_fixed_case_id(case) for case in BENCHMARK_FIXED_TEST_CASES], +) +def test_sanitizer_benchmark_fixed_cases(H, B, T): + inputs = _prepare_benchmark_fixed_inputs(B, T, H) + _run_benchmark_sanitizer_case(inputs) + torch.cuda.empty_cache() + + +@pytest.mark.sanitizer +@pytest.mark.benchmark +@pytest.mark.sm90_only +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.skipif( + os.environ.get("CULA_RUN_SANITIZER_TESTS") != "1", + reason="set CULA_RUN_SANITIZER_TESTS=1 and run under compute-sanitizer", +) +@pytest.mark.parametrize( + "H, seq_lens, total_len, dist", + BENCHMARK_VARLEN_TEST_CASES, + ids=[_varlen_case_id(case) for case in BENCHMARK_VARLEN_TEST_CASES], +) +def test_sanitizer_benchmark_varlen_cases(H, seq_lens, total_len, dist): + inputs = _prepare_benchmark_varlen_inputs(seq_lens, total_len, H) + _run_benchmark_sanitizer_case(inputs) + torch.cuda.empty_cache() + + +if __name__ == "__main__": + import sys + + bk = int(sys.argv[1]) if len(sys.argv) > 1 else 32 + bv = int(sys.argv[2]) if len(sys.argv) > 2 else 32 + occ = int(sys.argv[3]) if len(sys.argv) > 3 else 1 + print(f"Testing with bk={bk} bv={bv} min_occupancy={occ}") + _run_matches_fla_fixed(verbose=True, bk=bk, bv=bv, min_occupancy=occ) + print("\n✅ test_matches_fla_fixed PASSED") + print("\n=== uniform varlen (B=2) ===") + _run_matches_fla_uniform_varlen(verbose=True, bk=bk, bv=bv, min_occupancy=occ) + print("\n✅ test_matches_fla_uniform_varlen PASSED") diff --git a/tests/test_kda_chunk_sm90.py b/tests/test_kda_chunk_sm90.py new file mode 100644 index 00000000..a1ac933e --- /dev/null +++ b/tests/test_kda_chunk_sm90.py @@ -0,0 +1,96 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SM90 correctness tests for the public ``cula.kda.chunk_kda`` entrypoint.""" + +import pytest +import torch + +from benchmarks import bench_kda_fwd_bwd_e2e as bench +from benchmarks.utils import ( + SEED, + exclusive_cumsum, + prepare_safe_gate_inputs, + relative_rms_error_rel_max_mean_abs, + set_seed, +) + +pytestmark = pytest.mark.sm90_only + +D = 128 +OUT_NAMES = ("o", "ht", "dq", "dk", "dv", "dg", "dbeta", "dh0") + +SM90_CHUNK_CASES = ( + ("fixed_recompute_beta_fp32", 2, 64, (64, 64), torch.float32, False), + ("fixed_saved_intermediates_beta_fp32", 1, 64, (64,), torch.float32, True), + ("varlen_recompute_beta_bf16", 1, 96, (31, 65), torch.bfloat16, False), +) + + +def _make_inputs(batch_size, length, seq_lens, beta_dtype): + device = torch.device("cuda") + cu_seqlens = torch.tensor(exclusive_cumsum(list(seq_lens)), dtype=torch.int32, device=device) + inputs = prepare_safe_gate_inputs( + batch_size, + length, + 2, + D, + device, + cu_seqlens=cu_seqlens, + has_init_state=True, + num_v_heads=2, + ) + inputs["beta"] = inputs["beta"].to(beta_dtype) + set_seed(SEED + 1) + return { + "q": inputs["q"], + "k": inputs["k"], + "v": inputs["v"], + "g": inputs["g"], + "beta": inputs["beta"], + "scale": inputs["scale"], + "A_log": inputs["A_log"], + "dt_bias": inputs["dt_bias"], + "init_state": inputs["init_state"], + "cu_seqlens": cu_seqlens, + "lower_bound": inputs["lower_bound"], + "do": torch.randn_like(inputs["v"]), + "dht": torch.randn_like(inputs["init_state"]), + } + + +def _run_case(case): + case_id, batch_size, length, seq_lens, beta_dtype, disable_recompute = case + inputs = _make_inputs(batch_size, length, seq_lens, beta_dtype) + previous = bench.DISABLE_RECOMPUTE + bench.DISABLE_RECOMPUTE = disable_recompute + try: + reference = bench.run_kda_e2e_with_grads(**inputs, fn=bench.fla_chunk_kda) + actual = bench.run_kda_e2e_with_grads(**inputs, fn=bench.cula_chunk_kda) + torch.cuda.synchronize() + finally: + bench.DISABLE_RECOMPUTE = previous + + for name in OUT_NAMES: + assert reference[name].shape == actual[name].shape + assert torch.isfinite(actual[name]).all(), f"{case_id}: {name} contains non-finite values" + rel_rms, rel_max, mean_abs = relative_rms_error_rel_max_mean_abs(reference[name], actual[name]) + assert rel_rms < 0.05, f"{case_id}: {name} rel_rms={rel_rms:.6f}, mean_abs={mean_abs:.6e}" + assert rel_max < 0.25, f"{case_id}: {name} rel_max={rel_max:.6f}, mean_abs={mean_abs:.6e}" + + +@pytest.mark.parametrize("case", SM90_CHUNK_CASES, ids=[case[0] for case in SM90_CHUNK_CASES]) +def test_chunk_kda_sm90_entry_matches_fla(case): + _run_case(case) + torch.cuda.empty_cache()