diff --git a/benchmark/csa/bench_csa_compressor.py b/benchmark/csa/bench_csa_compressor.py new file mode 100644 index 000000000..0277828c8 --- /dev/null +++ b/benchmark/csa/bench_csa_compressor.py @@ -0,0 +1,503 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CSA fused-compressor benchmark with CUDA-graph variants (PR #427 fairness follow-up). + +This harness measures the wall-clock cost of the gated-softmax pooling region that the +fused ``cudnn.csa`` kernels replace, over the PR's production shapes (ratio=4, coff=2), +in a **single run** that emits two tables: + + * per-call -- each variant timed with a pair of CUDA events and a synchronize, + launch/host overhead included (reproduces the methodology behind the + per-call wall-clock table in docs/fe-oss-apis/csa.md). + * graph -- the SAME region captured into a CUDA graph and timed as a replay. + +Graph variants are captured **symmetrically** for eager and fused, as two directly +measured columns (no subtraction): + + * fwd-graph -- forward only. Eager forward is captured under ``torch.no_grad`` (no + autograd graph); fused forward is the wrapper (no autograd by design). + * total-graph -- forward + backward captured together. Eager backward is the autograd + backward of the captured forward; it is captured with **stable, + pre-allocated zero ``.grad`` buffers** (the supported torch graph + pattern): the captured region zeros those buffers in place, then runs + forward + backward, so every replay accumulates into a zeroed buffer -- + numerically identical to a single fresh backward (verified once per + shape via a graph-vs-fresh-backward cross-check printed at the end). + Fused backward is the explicit backward wrapper (kernel launches, no + autograd), captured right after the fused forward. + +The backward replay alone is NOT a separately captured quantity (the eager total graph +captures fwd+bwd as one unit and cannot be split into independent replays); it is +approximately ``total - fwd`` and is reported only as that reference, never as a column. + +Motivation: the eager region is ~39 forward and ~51 backward kernel launches per call, so +its per-call wall clock is dominated by launch/host overhead the fused path (1 + 1 +kernels) does not pay. Capturing each side into a graph collapses its per-operation +launches into a single replay, which is the fairest wall-clock basis for launch-bound +shapes. Numbers are reported as-measured (median of ``--iters`` after ``--warmup``); no +variant is favored. + +Not collected by pytest. Run, e.g.:: + + CUDA_VISIBLE_DEVICES=0 python benchmark/csa/bench_csa_compressor.py --iters 50 --warmup 20 +""" + +import argparse +import contextlib +import json +import math +import statistics + +import torch + +import cudnn.csa + +# --------------------------------------------------------------------------- +# Eager reference (verbatim upstream numerics; mirrors test_CSA_compressor.py) +# --------------------------------------------------------------------------- + + +def _batch_of_row(cu_seqlens, total): + """Segment index owning each packed row (mirror of Megatron-LM ``batch_of_row``).""" + n_seg = cu_seqlens.shape[0] - 1 + row_idx = torch.arange(total, device=cu_seqlens.device, dtype=torch.int64) + return torch.bucketize(row_idx, cu_seqlens[1:], right=True).clamp(max=max(n_seg - 1, 0)) + + +def _overlap_transform_thd(tensor, is_first_in_seg, head_dim, fill_value=0): + """Mirror of Megatron-LM Compressor._overlap_transform_thd (coff == 2).""" + n, ratio, b_dim, _ = tensor.size() + d = head_dim + new_tensor = tensor.new_full((n, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + prev_data = torch.roll(tensor[:, :, :, :d], shifts=1, dims=0) + prev_data[is_first_in_seg] = fill_value + new_tensor[:, :ratio] = prev_data + return new_tensor + + +def eager_pool(kv, score, ape, cu_seqlens, cu_seqlens_comp, total_comp, ratio, d, coff): + """Verbatim upstream eager pooling region (softmax weights rounded to bf16).""" + device = kv.device + row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_comp.dtype) + batch_ids = _batch_of_row(cu_seqlens_comp, total_comp) + valid_comp = row_idx < cu_seqlens_comp[-1] + local_pos = row_idx - cu_seqlens_comp[batch_ids] + local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) + base = cu_seqlens[batch_ids].unsqueeze(1) + local_pos.unsqueeze(1) * ratio + base = torch.where(valid_comp.unsqueeze(1), base, torch.zeros_like(base)) + offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) + gather_idx = base + offsets # (total_comp, ratio) + + kv_grouped = kv[gather_idx] + score_grouped = score[gather_idx] + score_grouped = score_grouped + ape.view(1, ratio, 1, -1) + + if coff == 2: + is_first = local_pos == 0 + kv_grouped = _overlap_transform_thd(kv_grouped, is_first, d, fill_value=0) + score_grouped = _overlap_transform_thd(score_grouped, is_first, d, fill_value=float("-inf")) + + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32).to(kv_grouped.dtype) + out = (kv_grouped * weights).sum(dim=1) + return out # (total_comp, 1, d) + + +# --------------------------------------------------------------------------- +# Input construction +# --------------------------------------------------------------------------- + + +def make_inputs(lens, d, ratio, coff, seed=1234, device="cuda"): + """Build a random THD pack (kv, score, ape, cu, cuc, total_comp, grad_out) for ``lens``.""" + total = sum(lens) + w = coff * d + gen = torch.Generator(device="cpu").manual_seed(seed) + kv = torch.randn(total, 1, w, generator=gen, dtype=torch.float32).to(torch.bfloat16) + score = (torch.randn(total, 1, w, generator=gen, dtype=torch.float32).mul_(1.5)).to(torch.bfloat16) + ape = torch.randn(ratio, w, generator=gen, dtype=torch.float32).mul_(0.25) + cu = torch.tensor([0, *torch.tensor(lens).cumsum(0)], dtype=torch.int32, device=device) + seg_comp = torch.tensor([seg_len // ratio for seg_len in lens]) + cuc = torch.tensor([0, *seg_comp.cumsum(0)], dtype=torch.int32, device=device) + total_comp = int(cuc[-1].item()) + go = torch.randn(total_comp, 1, d, generator=gen, dtype=torch.float32).to(torch.bfloat16) + return kv.to(device), score.to(device), ape.to(device), cu, cuc, total_comp, go.to(device) + + +# --------------------------------------------------------------------------- +# Variant callables (static-buffer friendly) +# --------------------------------------------------------------------------- + + +def fused_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff): + """One fused forward wrapper call; returns the pooled output.""" + total = kv.shape[0] + return cudnn.csa.csa_compressor_forward_wrapper( + kv.view(total, -1), + score.view(total, -1), + ape, + cu, + cuc, + ratio=ratio, + head_dim=d, + coff=coff, + total_comp=total_comp, + )["out"] + + +def fused_backward(kv, score, ape, cu, cuc, go, ratio, d, coff): + """One fused backward wrapper call; returns the gradient TupleDict.""" + total = kv.shape[0] + return cudnn.csa.csa_compressor_backward_wrapper( + kv.view(total, -1), + score.view(total, -1), + ape, + cu, + cuc, + go.view(go.shape[0], d), + ratio=ratio, + head_dim=d, + coff=coff, + ) + + +def fused_forward_backward(kv, score, ape, cu, cuc, go, total_comp, ratio, d, coff): + """Fused forward immediately followed by the fused backward (the total region).""" + fused_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff) + fused_backward(kv, score, ape, cu, cuc, go, ratio, d, coff) + + +def eager_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff): + """One eager forward over the replaced region (verbatim upstream numerics).""" + return eager_pool(kv, score, ape, cu, cuc, total_comp, ratio, d, coff) + + +# --------------------------------------------------------------------------- +# Timing helpers +# --------------------------------------------------------------------------- + + +def _noop(): + """Do nothing (placeholder ``pre`` hook for graph replay timing).""" + pass + + +def _median_event_ms(fn, warmup, iters): + """Median CUDA-event wall clock (ms) of per-call fn() (launch overhead included).""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + ts = [] + for _ in range(iters): + s = torch.cuda.Event(True) + e = torch.cuda.Event(True) + s.record() + fn() + e.record() + torch.cuda.synchronize() + ts.append(s.elapsed_time(e)) + return statistics.median(ts) + + +def _median_eager_bwd_event_ms(make_fwd, go, warmup, iters): + """Median CUDA-event wall clock (ms) of ONLY the autograd backward. + + The forward that builds the grad graph runs OUTSIDE the timed region (matching the PR + harness and the docs' "eager backward goes through torch autograd" basis). ``make_fwd`` + returns a fresh leaf set plus the forward output, so backward always starts from an + unallocated ``.grad`` (the fresh-gradient path). + """ + for _ in range(warmup): + _kvl, _scl, _apl, o = make_fwd() + o.backward(go) + torch.cuda.synchronize() + ts = [] + for _ in range(iters): + _kvl, _scl, _apl, o = make_fwd() # forward outside timing + torch.cuda.synchronize() + s = torch.cuda.Event(True) + e = torch.cuda.Event(True) + s.record() + o.backward(go) + e.record() + torch.cuda.synchronize() + ts.append(s.elapsed_time(e)) + return statistics.median(ts) + + +def _median_graph_ms(graph, pre, warmup, iters): + """Median CUDA-event wall clock (ms) of graph replays. ``pre`` runs before each replay.""" + for _ in range(warmup): + pre() + graph.replay() + torch.cuda.synchronize() + ts = [] + for _ in range(iters): + pre() + s = torch.cuda.Event(True) + e = torch.cuda.Event(True) + s.record() + graph.replay() + e.record() + torch.cuda.synchronize() + ts.append(s.elapsed_time(e)) + return statistics.median(ts) + + +def _capture(fn, warmup, no_grad=False): + """Warmup fn() on a side stream, then capture fn() into a CUDAGraph (capture-safe).""" + cm = torch.no_grad() if no_grad else contextlib.nullcontext() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + with cm: + for _ in range(warmup): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + with cm: + fn() + return g + + +def _capture_eager_total(kv, score, ape, cu, cuc, total_comp, go, ratio, d, coff, warmup): + """Capture eager fwd+bwd with STABLE pre-allocated zero ``.grad`` buffers. + + The leaf ``.grad`` tensors are created as zeros BEFORE warmup and kept across + warmup/capture/replay (never released). The captured region zeros them IN PLACE, then + runs forward + autograd backward, so every replay accumulates into a zeroed buffer -- + numerically identical to a single fresh backward. The in-place zero cost is inside the + graph (honestly counted in the total-graph time). + + Returns ``(graph, check)`` where ``check`` compares one graph replay's grads against a + fresh non-graph backward (graph-vs-fresh numerical-consistency evidence). + """ + kvl, scl, apl = _make_leaves(kv, score, ape) + # Stable zero grad buffers: fixed addresses the captured graph can replay against. + kvl.grad = torch.zeros_like(kvl) + scl.grad = torch.zeros_like(scl) + apl.grad = torch.zeros_like(apl) + + def _fwd_bwd(): + """The captured region: zero the kept grad buffers, then eager fwd + bwd.""" + # In-place zero of the SAME buffers each replay (kept, not released). + kvl.grad.zero_() + scl.grad.zero_() + apl.grad.zero_() + o = eager_forward(kvl, scl, apl, cu, cuc, total_comp, ratio, d, coff) + o.backward(go) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(warmup): + _fwd_bwd() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + _fwd_bwd() + + # ---- numerics cross-check: one graph replay vs a fresh non-graph backward ---- + g.replay() + torch.cuda.synchronize() + gk, gs, ga = kvl.grad.clone(), scl.grad.clone(), apl.grad.clone() + + ref_k, ref_s, ref_a = _make_leaves(kv, score, ape) + o_ref = eager_forward(ref_k, ref_s, ref_a, cu, cuc, total_comp, ratio, d, coff) + o_ref.backward(go) + torch.cuda.synchronize() + + check = { + "equal_dKV": torch.equal(gk, ref_k.grad), + "equal_dScore": torch.equal(gs, ref_s.grad), + "max_abs_dKV": (gk.float() - ref_k.grad.float()).abs().max().item(), + "max_abs_dScore": (gs.float() - ref_s.grad.float()).abs().max().item(), + "max_abs_dAPE": (ga - ref_a.grad).abs().max().item(), + "allclose": ( + torch.allclose(gk, ref_k.grad, rtol=2e-3, atol=2e-3) + and torch.allclose(gs, ref_s.grad, rtol=2e-3, atol=2e-3) + and torch.allclose(ga, ref_a.grad, rtol=1e-4, atol=1e-4) + ), + } + return g, check + + +# --------------------------------------------------------------------------- +# Per-shape measurement +# --------------------------------------------------------------------------- + + +def measure(lens, d, ratio, coff, warmup, iters, seed=1234): + """Measure per-call and graph-replay wall clock for one shape (eager vs fused). + + Returns a dict of median CUDA-event timings (ms) plus a numerics ``check`` comparing + one eager total-graph replay against a fresh non-graph backward. + """ + kv, score, ape, cu, cuc, total_comp, go = make_inputs(lens, d, ratio, coff, seed) + out = {} + + # ---- per-call (non-graph): reproduces the PR's event methodology ---- + out["eager_fwd"] = _median_event_ms(lambda: eager_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff), warmup, iters) + out["fused_fwd"] = _median_event_ms(lambda: fused_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff), warmup, iters) + + # eager backward per-call: forward (builds grad graph) outside timing, backward timed. + def _fwd_for_bwd(): + """Fresh leaves + eager forward (outside timing) feeding one timed backward.""" + kvl, scl, apl = _make_leaves(kv, score, ape) + o = eager_forward(kvl, scl, apl, cu, cuc, total_comp, ratio, d, coff) + return kvl, scl, apl, o + + out["eager_bwd"] = _median_eager_bwd_event_ms(_fwd_for_bwd, go, warmup, iters) + + out["fused_bwd"] = _median_event_ms(lambda: fused_backward(kv, score, ape, cu, cuc, go, ratio, d, coff), warmup, iters) + + # ---- graph variants ---- + # forward-only graphs (eager under no_grad so no autograd graph is built). + g_eager_fwd = _capture( + lambda: eager_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff), + warmup, + no_grad=True, + ) + out["eager_fwd_graph"] = _median_graph_ms(g_eager_fwd, pre=_noop, warmup=warmup, iters=iters) + del g_eager_fwd + + g_fused_fwd = _capture(lambda: fused_forward(kv, score, ape, cu, cuc, total_comp, ratio, d, coff), warmup) + out["fused_fwd_graph"] = _median_graph_ms(g_fused_fwd, pre=_noop, warmup=warmup, iters=iters) + del g_fused_fwd + + # total graphs: forward + backward captured together (static grad buffers for eager). + g_eager_total, out["check"] = _capture_eager_total(kv, score, ape, cu, cuc, total_comp, go, ratio, d, coff, warmup) + out["eager_total_graph"] = _median_graph_ms(g_eager_total, pre=_noop, warmup=warmup, iters=iters) + del g_eager_total + + g_fused_total = _capture( + lambda: fused_forward_backward(kv, score, ape, cu, cuc, go, total_comp, ratio, d, coff), + warmup, + ) + out["fused_total_graph"] = _median_graph_ms(g_fused_total, pre=_noop, warmup=warmup, iters=iters) + del g_fused_total + + return out + + +def _fmt(ms): + """Format a millisecond value as a fixed-width microsecond column.""" + return f"{ms * 1000:7.1f}" + + +def _make_leaves(kv, score, ape): + """Fresh grad-enabled leaf clones of the three inputs (one construction site).""" + return kv.clone().requires_grad_(True), score.clone().requires_grad_(True), ape.clone().requires_grad_(True) + + +def _x(num, den): + """Speedup helper (callers pass displayed 1-decimal values).""" + return float("inf") if den == 0 else num / den + + +def _xtrunc1(num, den): + """Speedup from displayed 1-decimal values, truncated to 1 decimal (never rounded up).""" + return float("inf") if den == 0 else math.floor((num / den) * 10) / 10 + + +def main(): + """Run all shapes and print the per-call/graph tables, checks, and JSON payload.""" + p = argparse.ArgumentParser() + p.add_argument("--iters", type=int, default=50) + p.add_argument("--warmup", type=int, default=20) + p.add_argument("--seed", type=int, default=1234) + args = p.parse_args() + + shapes = [ + ([8192], 128), + ([8192, 8192, 8192], 128), + ([8192], 512), + ([8192, 8192, 8192], 512), + ] + ratio, coff = 4, 2 + + print(f"# CSA compressor benchmark -- ratio={ratio} coff={coff} " f"warmup={args.warmup} iters={args.iters} (median, CUDA events, us)") + print(f"# GPU: {torch.cuda.get_device_name()} CC={torch.cuda.get_device_capability()}") + + rows = [] + for lens, d in shapes: + r = measure(lens, d, ratio, coff, args.warmup, args.iters, args.seed) + tag = f"{len(lens)}x{lens[0] // 1000}k/d{d}" + rows.append((tag, d, r)) + torch.cuda.empty_cache() + + # ---- per-call table ---- + print("\n# per-call wall clock (CUDA events, launch overhead included):") + h1 = f"{'shape':>14} | {'eager_fwd':>9} {'fused_fwd':>9} | {'eager_bwd':>9} {'fused_bwd':>9}" + print(h1) + print("-" * len(h1)) + for tag, _, r in rows: + print(f"{tag:>14} | {_fmt(r['eager_fwd'])} {_fmt(r['fused_fwd'])} | " f"{_fmt(r['eager_bwd'])} {_fmt(r['fused_bwd'])}") + + # ---- graph table (fwd-graph and total-graph, directly measured; no bwd-graph) ---- + print("\n# graph replay (fwd-only graph; fwd+bwd total graph):") + h2 = f"{'shape':>14} | {'e_fgraph':>9} {'f_fgraph':>9} {'fwd':>6} | " f"{'e_tgraph':>9} {'f_tgraph':>9} {'total':>6}" + print(h2) + print("-" * len(h2)) + for tag, _, r in rows: + efg = round(r["eager_fwd_graph"] * 1000, 1) + ffg = round(r["fused_fwd_graph"] * 1000, 1) + etg = round(r["eager_total_graph"] * 1000, 1) + ftg = round(r["fused_total_graph"] * 1000, 1) + print( + f"{tag:>14} | {_fmt(r['eager_fwd_graph'])} {_fmt(r['fused_fwd_graph'])} " + f"{_xtrunc1(efg, ffg):5.1f}x | {_fmt(r['eager_total_graph'])} {_fmt(r['fused_total_graph'])} " + f"{_xtrunc1(etg, ftg):5.1f}x" + ) + + # ---- speedup summary (eager / fused), per-call and graph, from displayed values ---- + print("\n# speedup (eager / fused), from displayed values:") + for tag, _, r in rows: + ef = round(r["eager_fwd"] * 1000, 1) + ff = round(r["fused_fwd"] * 1000, 1) + eb = round(r["eager_bwd"] * 1000, 1) + fb = round(r["fused_bwd"] * 1000, 1) + efg = round(r["eager_fwd_graph"] * 1000, 1) + ffg = round(r["fused_fwd_graph"] * 1000, 1) + etg = round(r["eager_total_graph"] * 1000, 1) + ftg = round(r["fused_total_graph"] * 1000, 1) + print(f" {tag:>14}: fwd {_x(ef, ff):5.2f}x (fwd-graph {_x(efg, ffg):5.2f}x) | " f"bwd {_x(eb, fb):5.2f}x | total-graph {_x(etg, ftg):5.2f}x") + + # ---- eager graph-vs-fresh-backward numerics cross-check ---- + print("\n# eager total-graph vs fresh non-graph backward (numerical consistency):") + for tag, _, r in rows: + c = r["check"] + flag = "PASS" if c["allclose"] else "FAIL" + print( + f" {tag:>14}: dKV eq={c['equal_dKV']} max={c['max_abs_dKV']:.3e} | " + f"dScore eq={c['equal_dScore']} max={c['max_abs_dScore']:.3e} | " + f"dAPE max={c['max_abs_dAPE']:.3e} -> {flag}" + ) + + # ---- machine-readable JSON ---- + payload = { + "gpu": torch.cuda.get_device_name(), + "cc": list(torch.cuda.get_device_capability()), + "ratio": ratio, + "coff": coff, + "warmup": args.warmup, + "iters": args.iters, + "rows": [ + { + "shape": t, + "head_dim": d, + "check": r["check"], + **{k: round(v * 1000.0, 1) for k, v in r.items() if k != "check"}, + } + for t, d, r in rows + ], + } + print("\n# JSON\n" + json.dumps(payload, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/csa/gate_csa_compressor_r128.py b/benchmark/csa/gate_csa_compressor_r128.py new file mode 100644 index 000000000..7bb94a5ec --- /dev/null +++ b/benchmark/csa/gate_csa_compressor_r128.py @@ -0,0 +1,613 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""ratio=128 CSA fused-compressor numerics contract gate (see docs/fe-oss-apis/csa.md). + +Runs the contract over the full validation envelope — 15 union cases (coff {1, 2} x +head_dim {128, 512} x {8k / 3x8k / ragged / edge / 64k / 128k token packs}; the 128k +d=512 packs select the large-bucket schedules), 1 scaled-input case and 4 padding +cases (these 20 finite-intermediate cases each run all three gates), plus 1 +overflow-intermediate case (fp32 score+ape == +Inf: determinism + eager NaN-pattern +faithfulness, fp64 comparator explicitly skipped) — 21 cases — and one ratio=4/coff=2 +reference cross-check. The case list is asserted to select every shipped +(config, schedule) kernel, forward and backward, through the dispatch tables +themselves. + +Per case, three gates: + + GATE 1 (deterministic; all supported inputs): 3 runs per direction (zero-initialized + + 2 NaN-prefilled output buffers): outputs bitwise-identical across runs — an + unwritten slot would disagree between the zero-init and prefilled runs, so the + pair proves kernel-side writes cover every slot — and no NaN survives on the 20 + finite-intermediate cases (their fp32 reference is NaN-free; the overflow case + instead replays its NaN pattern bit-stable on integer views). dAPE is exempt + (fp32 atomics); its replay delta is recorded. + GATE 2 (faithful to the fp32-intermediate eager reference; tolerances calibrated on + the DOCUMENTED INPUT DISTRIBUTION below): forward out and backward dKV/dScore: + n_diff <= max(1, 0.1% numel) AND max_abs <= 1.6e-2; dAPE max_abs <= 1e-3. On the + overflow-intermediate case the same faithfulness is checked as non-finite + propagation instead: the fused NaN mask must EQUAL the eager reference's on all + four outputs (both sides compute in fp32, so the fp32 overflow poisons both + alike). + GATE 3 (fp64-oracle parity; inputs whose fp32 intermediates stay finite): for the + deterministic outputs (out/dKV/dScore), per tensor, + max|fused - fp64| <= max|eager_fp32 - fp64| * (1 + 1e-6) + 1e-4. For dAPE the + same check is gated on the unit-scale cases only: dAPE is an fp32 atomic + accumulation whose reduction-order error fluctuates run to run AROUND the eager + error (observed within ~1.5x either way), so on scaled inputs — where the 1e-4 + absolute term no longer dominates — its parity outcome is a coin flip and is + recorded instead of gated. On the overflow-intermediate case the comparator is + SKIPPED with the reason recorded: the fp64 oracle stays finite where fp32 + overflows (asserted by the case), so a NaN-vs-finite distance gates nothing. + +GATE 2's absolute thresholds and every published worst-observed number are properties +of the gate input distribution, NOT of all supported inputs: bf16's grid is relative, +so absolute deviations scale with the input magnitudes (multiplying kv/grad_out by +2^k scales every out/dKV/dScore deviation exactly 2^k). The gate distribution is + kv ~ N(0, 1) bf16, score ~ N(0, 1.5^2) bf16, ape ~ N(0, 0.25^2) fp32, + grad_out ~ N(0, 1) bf16 (seeds 1234 / 7; padding cases 11 / 13). +The scaled-input case (kv and grad_out x64, an exact bf16 exponent shift) turns the +scale dependence into committed evidence: differing-element counts and fp64 parity +of the deterministic outputs are invariant, absolute deviations scale exactly (its +recorded max_abs values exceed the unit-scale thresholds by design, so GATE 2 checks +scale-adjusted bounds there: out x64, dKV x64, dScore/dAPE x4096). + +The overflow-intermediate case commits the fp32-boundary side of the contract: +kv = 1, score = bf16 max, ape = fp32 max on one 128-token c1d128 segment drives the +fp32 score + ape add to +Inf, so every output of the kernels AND of the fp32 eager +reference is NaN while the fp64 oracle stays finite (out == 1.0, asserted). The case +gates bit-stable replay of the NaN pattern and NaN-mask equality vs eager on all four +outputs; the non-finite clause of the contract is THIS committed case, not an +order-independence theorem — near fp32 max the fused evaluation order can saturate +earlier than the eager one (un-normalized chunk partials; see the Numerics caveat in +docs/fe-oss-apis/csa.md). + +Padding cases put static row-capacity padding (`total_comp > cu_seqlens_comp[-1]`) +AND token-capacity padding (`kv/score` rows beyond `cu_seqlens[-1]`) on a ragged +pack, and run the same three gates on the full padded shapes: the eager reference +replicates the row-0 window on capacity rows exactly like the kernels, its autograd +returns exact zeros on all never-consumed token rows, and the fused side runs with +NONZERO grad_out on the padding rows while the reference zeroes them (proving the +kernel ignores them). The token-padding gradient rows are additionally asserted to +come back exactly zero. + +Also RECORDS (not gates) whether out/dKV/dScore happen to be bitwise-equal to eager +per case, so docs can state honestly which parts remain bitwise. The eager reference +is the fp32-intermediate mirror of the Megatron-LM eager pooling region, identical to +the one in test/python/fe_api/csa/test_CSA_compressor.py; it is cross-checked here +against the production ratio=4 backward (bitwise dKV/dScore) before being trusted. + +Requires a CC 10.0 GPU and the ``cudnn[cutedsl]`` install. Not collected by pytest. +Run, e.g.:: + + CUDA_VISIBLE_DEVICES=0 python benchmark/csa/gate_csa_compressor_r128.py --json gate.json +""" + +import argparse +import json +import math +import sys + +import torch + +RATIO = 128 +CONFIGS = {"c1d128": (1, 128), "c2d128": (2, 128), "c1d512": (1, 512), "c2d512": (2, 512)} +EDGE_PACK = [127, 8192, 0, 129, 128, 3, 515, 1024] +CASES = [ + ("c1d128", "1x8192", [8192]), + ("c1d128", "3x8192", [8192] * 3), + ("c1d128", "ragged3", [1023, 2048, 509]), + ("c1d128", "edgepack", EDGE_PACK), + ("c1d128", "1x65536", [65536]), + ("c1d128", "1x131072", [131072]), + ("c2d128", "1x8192", [8192]), + ("c2d128", "edgepack", EDGE_PACK), + ("c2d128", "1x65536", [65536]), + ("c1d512", "1x8192", [8192]), + ("c1d512", "1x65536", [65536]), + ("c1d512", "1x131072", [131072]), + ("c2d512", "1x8192", [8192]), + ("c2d512", "1x65536", [65536]), + ("c2d512", "1x131072", [131072]), +] +# (config, shape_name, lens, scale): kv and grad_out multiplied by `scale` (a power of +# two, so the bf16 inputs shift exponents without re-rounding). See the module +# docstring: proves determinism + fp64 parity of the deterministic outputs are +# invariant under exact exponent shifts (while the fp32 intermediates stay finite) +# and absolute deviations scale with magnitude. +SCALED_CASES = [("c2d128", "1x8192", [8192], 64.0)] +PADDING_CASES = [("c1d512", 8, 37), ("c2d512", 8, 21), ("c1d128", 5, 37), ("c2d128", 5, 21)] +ABS_TOL = 1.6e-2 +DIFF_FRAC = 0.001 +DAPE_TOL = 1e-3 + + +# --------------------------------------------------------------------------- +# Eager reference (fp32-intermediate mirror of the Megatron-LM eager region; +# identical to the reference in test_CSA_compressor.py) +# --------------------------------------------------------------------------- + + +def _batch_of_row(cu_seqlens, total): + n_seg = cu_seqlens.shape[0] - 1 + row_idx = torch.arange(total, device=cu_seqlens.device, dtype=torch.int64) + return torch.bucketize(row_idx, cu_seqlens[1:], right=True).clamp(max=max(n_seg - 1, 0)) + + +def _overlap_transform_thd(tensor, is_first_in_seg, head_dim, fill_value=0): + n, ratio, b_dim, _ = tensor.size() + d = head_dim + new_tensor = tensor.new_full((n, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + prev_data = torch.roll(tensor[:, :, :, :d], shifts=1, dims=0) + prev_data[is_first_in_seg] = fill_value + new_tensor[:, :ratio] = prev_data + return new_tensor + + +def _eager_pool(kv, score, ape, cu_seqlens, cu_seqlens_comp, total_comp, ratio, d, coff, mode): + device = kv.device + row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_comp.dtype) + batch_ids = _batch_of_row(cu_seqlens_comp, total_comp) + valid_comp = row_idx < cu_seqlens_comp[-1] + local_pos = row_idx - cu_seqlens_comp[batch_ids] + local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) + base = cu_seqlens[batch_ids].unsqueeze(1) + local_pos.unsqueeze(1) * ratio + base = torch.where(valid_comp.unsqueeze(1), base, torch.zeros_like(base)) + offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) + gather_idx = base + offsets # (total_comp, ratio) + + if mode == "fp32": + kv = kv.float() + score = score.float() + elif mode == "fp64": + kv = kv.double() + score = score.double() + ape = ape.double() + + kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) + score_grouped = score[gather_idx] + score_grouped = score_grouped + ape.view(1, ratio, 1, -1) + + if coff == 2: + is_first = local_pos == 0 + kv_grouped = _overlap_transform_thd(kv_grouped, is_first, d, fill_value=0) + score_grouped = _overlap_transform_thd(score_grouped, is_first, d, fill_value=float("-inf")) + + if mode == "fp32": + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32) + out = (kv_grouped * weights).sum(dim=1).to(torch.bfloat16) + else: # fp64 oracle + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float64) + out = (kv_grouped * weights).sum(dim=1) + return out # (total_comp, 1, d) + + +def run_eager_bwd(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode): + """Forward + backward through the eager reference; returns (out, dKV, dScore, dAPE).""" + dtype = torch.float64 if mode == "fp64" else None + kv_l = (kv.to(dtype) if dtype else kv.clone()).requires_grad_(True) + score_l = (score.to(dtype) if dtype else score.clone()).requires_grad_(True) + ape_l = (ape.to(dtype) if dtype else ape.clone()).requires_grad_(True) + out = _eager_pool(kv_l, score_l, ape_l, cu, cuc, total_comp, ratio, d, coff, mode) + out.backward(go.to(out.dtype)) + torch.cuda.synchronize() + return out.detach(), kv_l.grad.detach(), score_l.grad.detach(), ape_l.grad.detach() + + +def _make_inputs(lens, d, ratio, coff, seed=1234): + total = sum(lens) + w = coff * d + gen = torch.Generator(device="cpu").manual_seed(seed) + kv = torch.randn(total, 1, w, generator=gen, dtype=torch.float32).to(torch.bfloat16) + score = (torch.randn(total, 1, w, generator=gen, dtype=torch.float32).mul_(1.5)).to(torch.bfloat16) + ape = torch.randn(ratio, w, generator=gen, dtype=torch.float32).mul_(0.25) + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda") + seg_comp = torch.tensor([seg_len // ratio for seg_len in lens]) + cuc = torch.tensor([0] + list(seg_comp.cumsum(0)), dtype=torch.int32, device="cuda") + total_comp = int(cuc[-1].item()) + return kv.cuda(), score.cuda(), ape.cuda(), cu, cuc, total_comp + + +def _make_go(total_comp, d, seed=7): + gen = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(total_comp, 1, d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda() + + +# --------------------------------------------------------------------------- +# Gates +# --------------------------------------------------------------------------- + + +def tol_check(fused, eager, tol=ABS_TOL): + diff = (fused.float() - eager.float()).abs() + n_diff = int((diff > 0).sum().item()) + max_abs = float(diff.max().item()) if diff.numel() else 0.0 + ok = n_diff <= max(1, int(DIFF_FRAC * fused.numel())) and max_abs <= tol + return ok, n_diff, max_abs + + +def oracle_check(fused, eager, oracle): + err_f = float((fused.double() - oracle.double()).abs().max().item()) + err_e = float((eager.double() - oracle.double()).abs().max().item()) + return err_f <= err_e * (1 + 1e-6) + 1e-4, err_f, err_e + + +def _case_schedules(M, d, coff, nb_total): + """The (fwd, bwd) schedules the dispatch tables select for this case.""" + return M._fwd_schedule_r128(RATIO, d, coff, nb_total), M._bwd_schedule_r128(RATIO, d, coff, nb_total) + + +def _shipped_schedules(M): + """Every (config, schedule) pair the dispatch tables can select, per direction — + derived from the tables' own boundaries so the coverage assertion tracks edits.""" + fwd, bwd = set(), set() + for coff, d in CONFIGS.values(): + for nb in (1, M._SMALL_NB_MAX, M._SMALL_NB_MAX + 1, M._BWD_SMALL_NB_MAX, M._BWD_SMALL_NB_MAX + 1, M._LARGE_NB_MIN - 1, M._LARGE_NB_MIN): + fwd.add(((coff, d), M._fwd_schedule_r128(RATIO, d, coff, nb))) + bwd.add(((coff, d), M._bwd_schedule_r128(RATIO, d, coff, nb))) + return fwd, bwd + + +def gate_case(M, config, shape_name, lens, rec_list, covered, scale=1.0): + coff, d = CONFIGS[config] + kv, score, ape, cu, cuc, tc = _make_inputs(lens, d, RATIO, coff) + go = _make_go(tc, d) + scaled = scale != 1.0 + if scaled: + assert math.log2(scale).is_integer(), "scale must be a power of two (exact bf16 exponent shift)" + kv = (kv.float() * scale).to(torch.bfloat16) + go = (go.float() * scale).to(torch.bfloat16) + total = kv.shape[0] + w = coff * d + kvf, scf, gof = kv.view(total, w), score.view(total, w), go.view(tc, d) + sched_fwd, sched_bwd = _case_schedules(M, d, coff, tc) + covered["fwd"].add(((coff, d), sched_fwd)) + covered["bwd"].add(((coff, d), sched_bwd)) + + # ---------- forward: 3 runs (zero-init, then 2 NaN-prefilled) ---------- + out = torch.empty(tc, d, dtype=torch.bfloat16, device="cuda") + outs = [] + for rep in range(3): + if rep == 0: + out.zero_() + else: + out.fill_(float("nan")) + M.run_fwd_r128(kvf, scf, ape, cu, cuc, out, tc, RATIO, d, coff) + torch.cuda.synchronize() + outs.append(out.clone()) + fwd_stable = all(torch.equal(outs[0], o) for o in outs[1:]) + fwd_nan_ok = not any(torch.isnan(o).any().item() for o in outs) + ref32 = _eager_pool(kv, score, ape, cu, cuc, tc, RATIO, d, coff, "fp32").view(tc, d) + ref64 = _eager_pool(kv, score, ape, cu, cuc, tc, RATIO, d, coff, "fp64").view(tc, d) + fwd_tol_ok, fwd_nd, fwd_ma = tol_check(outs[0], ref32, ABS_TOL * scale) + fwd_or_ok, fwd_errf, fwd_erre = oracle_check(outs[0], ref32.double(), ref64) + fwd_bitwise = int((outs[0].float() - ref32.float()).abs().max().item() == 0) + + # ---------- backward: 3 runs (zero-init, then 2 NaN-prefilled) ---------- + runs = [] + for rep in range(3): + gkv = torch.empty(total, w, dtype=torch.bfloat16, device="cuda") + gs = torch.empty(total, w, dtype=torch.bfloat16, device="cuda") + if rep == 0: + gkv.zero_() + gs.zero_() + else: + gkv.fill_(float("nan")) + gs.fill_(float("nan")) + gape = torch.zeros_like(ape) + M.run_bwd_r128(kvf, scf, ape, cu, cuc, gof, gkv, gs, gape, tc, RATIO, d, coff) + torch.cuda.synchronize() + runs.append((gkv, gs, gape)) + bwd_stable = all(torch.equal(runs[0][0], r[0]) and torch.equal(runs[0][1], r[1]) for r in runs[1:]) + bwd_nan_ok = not any(torch.isnan(r[0]).any().item() or torch.isnan(r[1]).any().item() for r in runs) + ape_replay = max(float((runs[0][2] - r[2]).abs().max().item()) for r in runs[1:]) + gkv, gs, gape = runs[0] + + _, ekv, es, eape = run_eager_bwd(kv, score, ape, cu, cuc, tc, RATIO, d, coff, go, mode="fp32") + _, okv, os_, oape = run_eager_bwd(kv, score, ape, cu, cuc, tc, RATIO, d, coff, go, mode="fp64") + # Scaled inputs scale the absolute deviations of each output by its input factors + # (kv enters out linearly; grad_out enters dKV linearly; both enter dScore/dAPE). + kv_tol_ok, kv_nd, kv_ma = tol_check(gkv.view_as(ekv), ekv, ABS_TOL * scale) + s_tol_ok, s_nd, s_ma = tol_check(gs.view_as(es), es, ABS_TOL * scale * scale) + ape_ma = float((gape - eape).abs().max().item()) + ape_tol_ok = ape_ma <= DAPE_TOL * scale * scale + kv_or_ok, kv_errf, kv_erre = oracle_check(gkv.view_as(ekv), ekv, okv) + s_or_ok, s_errf, s_erre = oracle_check(gs.view_as(es), es, os_) + ape_or_ok, ape_errf, ape_erre = oracle_check(gape, eape, oape) + kv_bitwise = int(torch.equal(gkv.view_as(ekv), ekv)) + s_bitwise = int(torch.equal(gs.view_as(es), es)) + + g1 = fwd_stable and fwd_nan_ok and bwd_stable and bwd_nan_ok + g2 = fwd_tol_ok and kv_tol_ok and s_tol_ok and ape_tol_ok + # dAPE parity is gated at unit scale only (recorded when scaled; see module docstring). + g3 = fwd_or_ok and kv_or_ok and s_or_ok and (ape_or_ok or scaled) + ok = g1 and g2 and g3 + rec_list.append( + dict( + config=config, + shape=shape_name, + scale=scale, + sched_fwd=list(sched_fwd), + sched_bwd=list(sched_bwd), + ok=ok, + g1_deterministic=g1, + g2_tolerance=g2, + g3_oracle=g3, + fwd=dict( + stable=fwd_stable, + nan_ok=fwd_nan_ok, + n_diff=fwd_nd, + numel=outs[0].numel(), + max_abs=fwd_ma, + bitwise_vs_eager=fwd_bitwise, + err_fused=fwd_errf, + err_eager=fwd_erre, + ), + dkv=dict(n_diff=kv_nd, numel=int(ekv.numel()), max_abs=kv_ma, bitwise_vs_eager=kv_bitwise, err_fused=kv_errf, err_eager=kv_erre), + dscore=dict(n_diff=s_nd, max_abs=s_ma, bitwise_vs_eager=s_bitwise, err_fused=s_errf, err_eager=s_erre), + dape=dict(max_abs=ape_ma, replay_delta=ape_replay, parity_gated=not scaled, parity_ok=bool(ape_or_ok), err_fused=ape_errf, err_eager=ape_erre), + bwd=dict(stable=bwd_stable, nan_ok=bwd_nan_ok), + ) + ) + label = f"{shape_name}*{scale:g}" if scaled else shape_name + print( + f" {config:8s} {label:10s} {'PASS' if ok else 'FAIL'} " + f"g1={'Y' if g1 else 'N'} g2={'Y' if g2 else 'N'} g3={'Y' if g3 else 'N'} | " + f"fwd nd={fwd_nd}/{outs[0].numel()} ma={fwd_ma:.2e}{' BITW' if fwd_bitwise else ''} | " + f"dKV nd={kv_nd} ma={kv_ma:.2e}{' BITW' if kv_bitwise else ''} | " + f"dS nd={s_nd} ma={s_ma:.2e}{' BITW' if s_bitwise else ''} | dAPE {ape_ma:.1e}", + flush=True, + ) + return ok + + +def gate_padding(M, config, pad, tok_pad, rec_list, covered): + """Static row-capacity padding (+pad rows, incoming grads ignored) + token-capacity + padding (+tok_pad kv/score rows, gradients exactly zero), all three gates on the + full padded shapes. The eager reference models the capacity rows (row-0 window + replication) and runs with the padding-row grad_out zeroed — the fused side keeps + it NONZERO, so matching it proves the kernel ignores padding-row gradients.""" + coff, d = CONFIGS[config] + lens = [1023, 2048, 509] + kv, score, ape, cu, cuc, total_true = _make_inputs(lens, d, RATIO, coff) + total_comp = total_true + pad + go = _make_go(total_comp, d, seed=11) + go_zero_pad = go.clone() + go_zero_pad[total_true:] = 0 + gen = torch.Generator(device="cpu").manual_seed(13) + w = coff * d + kv2 = torch.cat([kv.view(kv.shape[0], -1), torch.randn(tok_pad, w, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda()]) + score2 = torch.cat([score.view(score.shape[0], -1), torch.randn(tok_pad, w, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda()]) + total = kv2.shape[0] + gof = go.view(total_comp, d) + sched_fwd, sched_bwd = _case_schedules(M, d, coff, total_comp) + covered["fwd"].add(((coff, d), sched_fwd)) + covered["bwd"].add(((coff, d), sched_bwd)) + + def run_dir(poison): + out = torch.empty(total_comp, d, dtype=torch.bfloat16, device="cuda") + gkv = torch.empty(total, w, dtype=torch.bfloat16, device="cuda") + gs = torch.empty(total, w, dtype=torch.bfloat16, device="cuda") + for t in (out, gkv, gs): + if poison: + t.fill_(float("nan")) + else: + t.zero_() + gape = torch.zeros_like(ape) + M.run_fwd_r128(kv2, score2, ape, cu, cuc, out, total_comp, RATIO, d, coff) + M.run_bwd_r128(kv2, score2, ape, cu, cuc, gof, gkv, gs, gape, total_comp, RATIO, d, coff) + torch.cuda.synchronize() + return out, gkv, gs, gape + + a = run_dir(False) + b = run_dir(True) + c = run_dir(True) + stable = all(torch.equal(a[i], b[i]) and torch.equal(a[i], c[i]) for i in range(3)) + nan_ok = not any(torch.isnan(t).any().item() for r in (a, b, c) for t in r[:3]) + ape_replay = max(float((a[3] - r[3]).abs().max().item()) for r in (b, c)) + out, gkv, gs, gape = a + + # Eager reference at full padded shape (kv/score with the token-padding rows, + # total_comp with the capacity rows, padding-row grad_out zeroed). + kv_full = kv2.view(total, 1, w) + score_full = score2.view(total, 1, w) + ref32 = _eager_pool(kv_full, score_full, ape, cu, cuc, total_comp, RATIO, d, coff, "fp32").view(total_comp, d) + ref64 = _eager_pool(kv_full, score_full, ape, cu, cuc, total_comp, RATIO, d, coff, "fp64").view(total_comp, d) + _, ekv, es, eape = run_eager_bwd(kv_full, score_full, ape, cu, cuc, total_comp, RATIO, d, coff, go_zero_pad, mode="fp32") + _, okv, os_, oape = run_eager_bwd(kv_full, score_full, ape, cu, cuc, total_comp, RATIO, d, coff, go_zero_pad, mode="fp64") + + fwd_tol_ok, fwd_nd, fwd_ma = tol_check(out, ref32) + kv_tol_ok, kv_nd, kv_ma = tol_check(gkv.view_as(ekv), ekv) + s_tol_ok, s_nd, s_ma = tol_check(gs.view_as(es), es) + ape_ma = float((gape - eape).abs().max().item()) + fwd_or_ok, fwd_errf, fwd_erre = oracle_check(out, ref32.double(), ref64) + kv_or_ok, kv_errf, kv_erre = oracle_check(gkv.view_as(ekv), ekv, okv) + s_or_ok, s_errf, s_erre = oracle_check(gs.view_as(es), es, os_) + ape_or_ok, ape_errf, ape_erre = oracle_check(gape, eape, oape) + pad_zero_ok = bool((gkv[kv.shape[0] :] == 0).all().item() and (gs[kv.shape[0] :] == 0).all().item()) + + g1 = stable and nan_ok + g2 = fwd_tol_ok and kv_tol_ok and s_tol_ok and ape_ma <= DAPE_TOL + g3 = fwd_or_ok and kv_or_ok and s_or_ok and ape_or_ok + ok = g1 and g2 and g3 and pad_zero_ok + rec_list.append( + dict( + kind="padding", + config=config, + pad=pad, + tok_pad=tok_pad, + sched_fwd=list(sched_fwd), + sched_bwd=list(sched_bwd), + ok=ok, + g1_deterministic=g1, + g2_tolerance=g2, + g3_oracle=g3, + fwd=dict(stable=stable, nan_ok=nan_ok, n_diff=fwd_nd, numel=out.numel(), max_abs=fwd_ma, err_fused=fwd_errf, err_eager=fwd_erre), + dkv=dict(n_diff=kv_nd, numel=int(ekv.numel()), max_abs=kv_ma, err_fused=kv_errf, err_eager=kv_erre), + dscore=dict(n_diff=s_nd, max_abs=s_ma, err_fused=s_errf, err_eager=s_erre), + dape=dict(max_abs=ape_ma, replay_delta=ape_replay, parity_gated=True, parity_ok=bool(ape_or_ok), err_fused=ape_errf, err_eager=ape_erre), + pad_zero_ok=pad_zero_ok, + ) + ) + print( + f" padding {config} +{pad} rows +{tok_pad} tokens: {'PASS' if ok else 'FAIL'} " + f"g1={'Y' if g1 else 'N'} g2={'Y' if g2 else 'N'} g3={'Y' if g3 else 'N'} | " + f"fwd nd={fwd_nd} ma={fwd_ma:.1e} | dKV nd={kv_nd} ma={kv_ma:.1e} | dS nd={s_nd} ma={s_ma:.1e} | " + f"dAPE {ape_ma:.1e} padzero={pad_zero_ok}", + flush=True, + ) + return ok + + +def gate_overflow(M, rec_list, covered): + """Overflow-intermediate case (see the module docstring): kv = 1, score = bf16 max, + ape = fp32 max on one 128-token c1d128 segment. The fp32 ``score + ape`` add is + +Inf, so the kernels and the fp32 eager reference return all-NaN outputs alike + (both compute in fp32) while the fp64 oracle stays finite (out == 1.0): GATE 2 + becomes NaN-mask equality vs eager on all four outputs, GATE 3 is skipped with the + reason recorded (oracle finiteness asserted so the reason stays honest), and + GATE 1's bitwise replay compares int16 bit views (torch.equal is false on NaN); + the zero-init/NaN-prefill pair still proves kernel-side writes cover every slot, + because an unwritten slot would disagree between the zero-init and prefilled + runs.""" + config = "c1d128" + coff, d = CONFIGS[config] + total, w = 128, coff * d + tc = total // RATIO + kv = torch.full((total, 1, w), 1.0, dtype=torch.float32).to(torch.bfloat16).cuda() + score = torch.full((total, 1, w), torch.finfo(torch.bfloat16).max, dtype=torch.float32).to(torch.bfloat16).cuda() + ape = torch.full((RATIO, w), torch.finfo(torch.float32).max, dtype=torch.float32, device="cuda") + cu = torch.tensor([0, total], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, tc], dtype=torch.int32, device="cuda") + go = torch.full((tc, 1, d), 1.0, dtype=torch.float32).to(torch.bfloat16).cuda() + kvf, scf, gof = kv.view(total, w), score.view(total, w), go.view(tc, d) + sched_fwd, sched_bwd = _case_schedules(M, d, coff, tc) + covered["fwd"].add(((coff, d), sched_fwd)) + covered["bwd"].add(((coff, d), sched_bwd)) + + runs = [] + for rep in range(3): + out = torch.empty(tc, d, dtype=torch.bfloat16, device="cuda") + gkv = torch.empty(total, w, dtype=torch.bfloat16, device="cuda") + gs = torch.empty(total, w, dtype=torch.bfloat16, device="cuda") + for t in (out, gkv, gs): + if rep == 0: + t.zero_() + else: + t.fill_(float("nan")) + gape = torch.zeros_like(ape) + M.run_fwd_r128(kvf, scf, ape, cu, cuc, out, tc, RATIO, d, coff) + M.run_bwd_r128(kvf, scf, ape, cu, cuc, gof, gkv, gs, gape, tc, RATIO, d, coff) + torch.cuda.synchronize() + runs.append((out, gkv, gs, gape)) + out, gkv, gs, gape = runs[0] + # dAPE (fp32 atomics) is exempt from the bitwise replay as everywhere else; its + # NaN mask is order-independent (NaN is absorbing under addition) and checked below. + stable = all(all(torch.equal(runs[0][i].view(torch.int16), r[i].view(torch.int16)) for i in range(3)) for r in runs[1:]) + + ref32 = _eager_pool(kv, score, ape, cu, cuc, tc, RATIO, d, coff, "fp32").view(tc, d) + ref64 = _eager_pool(kv, score, ape, cu, cuc, tc, RATIO, d, coff, "fp64").view(tc, d) + _, ekv, es, eape = run_eager_bwd(kv, score, ape, cu, cuc, tc, RATIO, d, coff, go, mode="fp32") + _, okv, os_, oape = run_eager_bwd(kv, score, ape, cu, cuc, tc, RATIO, d, coff, go, mode="fp64") + + pairs = (("out", out, ref32), ("dkv", gkv.view_as(ekv), ekv), ("dscore", gs.view_as(es), es), ("dape", gape, eape)) + masks_equal = {name: bool((torch.isnan(f) == torch.isnan(e)).all().item()) for name, f, e in pairs} + # This input poisons every slot (every window sees the +Inf add); asserting the + # eager masks are FULL pins the case so it cannot silently hollow out. + masks_full = {name: bool(torch.isnan(e).all().item()) for name, _, e in pairs} + # GATE 3 skip: assert the reason (fp64 does not overflow on these inputs), then + # record it instead of running the meaningless NaN-vs-finite comparison. + oracle_finite = all(bool(torch.isfinite(t).all().item()) for t in (ref64, okv, os_, oape)) + skip_reason = "fp32 intermediates overflow while the fp64 oracle stays finite (asserted): a NaN-vs-finite distance gates nothing" + + g1 = stable + g2 = all(masks_equal.values()) and all(masks_full.values()) + ok = g1 and g2 and oracle_finite + rec_list.append( + dict( + kind="overflow", + config=config, + shape="1x128ovf", + sched_fwd=list(sched_fwd), + sched_bwd=list(sched_bwd), + ok=ok, + g1_deterministic=g1, + g2_eager_nan_pattern=g2, + g3_oracle_skipped=skip_reason, + oracle_finite=oracle_finite, + nan_numel={name: [int(torch.isnan(f).sum().item()), int(f.numel())] for name, f, _ in pairs}, + ) + ) + print( + f" {config:8s} {'1x128ovf':10s} {'PASS' if ok else 'FAIL'} " + f"g1={'Y' if g1 else 'N'} g2={'Y' if g2 else 'N'} g3=SKIP(oracle finite: {'Y' if oracle_finite else 'N'}) | " + f"NaN out {int(torch.isnan(out).sum())}/{out.numel()} dKV {int(torch.isnan(gkv).sum())}/{gkv.numel()} " + f"dS {int(torch.isnan(gs).sum())}/{gs.numel()} dAPE {int(torch.isnan(gape).sum())}/{gape.numel()} == eager", + flush=True, + ) + return ok + + +def validate_reference(): + """One ratio=4/coff=2 cross-check: the production backward is bitwise against the + eager reference above, proving the reference before the r128 cases trust it.""" + from cudnn.csa.compressor import compressor_sm100 as K4 + + kv, score, ape, cu, cuc, tc = _make_inputs([2048], 128, 4, 2) + go = _make_go(tc, 128) + total = kv.shape[0] + kvf, scf, gof = kv.view(total, 256), score.view(total, 256), go.view(tc, 128) + gkv = torch.zeros_like(kvf) + gs = torch.zeros_like(scf) + gape = torch.zeros_like(ape) + K4.run_bwd(kvf, scf, ape, cu, cuc, gof, gkv, gs, gape, tc, 4, 128, 2) + torch.cuda.synchronize() + _, ekv, es, eape = run_eager_bwd(kv, score, ape, cu, cuc, tc, 4, 128, 2, go, mode="fp32") + okv = torch.equal(gkv.view_as(ekv), ekv) + os_ = torch.equal(gs.view_as(es), es) + oape = float((gape - eape).abs().max().item()) + print(f" reference validation @ r4/c2: dKV bitwise={okv} dScore bitwise={os_} dAPE {oape:.1e}", flush=True) + return okv and os_ and oape <= DAPE_TOL + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--quick", action="store_true", help="run only the first two union cases (smoke)") + ap.add_argument("--json", default=None, help="write the per-case records to this path") + args = ap.parse_args() + from cudnn.csa.compressor import compressor_sm100_r128 as M + + assert torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0), "requires a CC 10.0 GPU" + print(f"r128 contract gate on {torch.cuda.get_device_name()}", flush=True) + recs = [] + covered = {"fwd": set(), "bwd": set()} + ok = validate_reference() + for config, shape_name, lens in CASES if not args.quick else CASES[:2]: + ok = gate_case(M, config, shape_name, lens, recs, covered) and ok + if not args.quick: + for config, shape_name, lens, scale in SCALED_CASES: + ok = gate_case(M, config, shape_name, lens, recs, covered, scale=scale) and ok + for config, pad, tok_pad in PADDING_CASES: + ok = gate_padding(M, config, pad, tok_pad, recs, covered) and ok + ok = gate_overflow(M, recs, covered) and ok + # Every shipped (config, schedule) kernel must have been selected by >= 1 case + # (derived from the dispatch tables, so adding a bucket fails the gate until a + # case covers it). + ship_fwd, ship_bwd = _shipped_schedules(M) + missing = [("fwd", *m) for m in sorted(ship_fwd - covered["fwd"])] + [("bwd", *m) for m in sorted(ship_bwd - covered["bwd"])] + print( + f" schedule coverage: fwd {len(covered['fwd'])}/{len(ship_fwd)} bwd {len(covered['bwd'])}/{len(ship_bwd)}" + + (f" MISSING {missing}" if missing else ""), + flush=True, + ) + ok = ok and not missing + n_pass = sum(1 for r in recs if r["ok"]) + print(f"GATE {'PASS' if ok else 'FAIL'} {n_pass}/{len(recs)}", flush=True) + if args.json: + with open(args.json, "w") as f: + json.dump(dict(all_pass=ok, cases=recs), f, indent=1) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmark/csa/reg_probe_csa_compressor_r128.py b/benchmark/csa/reg_probe_csa_compressor_r128.py new file mode 100644 index 000000000..d73825650 --- /dev/null +++ b/benchmark/csa/reg_probe_csa_compressor_r128.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""ptxas register/spill probe for every shipped ratio=128 CSA compressor kernel. + +JIT-compiles the full shipped dispatch envelope — every (config, schedule) pair the +``nb_total`` bucket tables can select, forward and backward, over coff {1, 2} x +head_dim {128, 512} (16 kernels) — then runs ``ptxas -v`` on each kernel's PTX and +prints a table of registers / spill bytes / stack bytes / ex2.approx count. This +reproduces the register table published in docs/fe-oss-apis/csa.md. + +Exits nonzero if any kernel spills, uses stack, or fails ptxas. + +Requires a CC 10.0 GPU (the JIT needs a device), ``ptxas`` on PATH, and the +``cudnn[cutedsl]`` install. Not collected by pytest. Run, e.g.:: + + CUDA_VISIBLE_DEVICES=0 python benchmark/csa/reg_probe_csa_compressor_r128.py +""" + +import argparse +import os +import re +import shutil +import subprocess +import sys +import tempfile + +os.environ.setdefault("CUTE_DSL_KEEP", "ptx") # keep PTX artifacts on the compiled handles + +import torch # noqa: E402 + + +def ptxas_one(tag, ptx_text, arch, out_dir): + path = os.path.join(out_dir, f"{tag}.ptx") + with open(path, "w") as f: + f.write(ptx_text) + r = subprocess.run(["ptxas", "-v", f"-arch={arch}", "-o", os.devnull, path], capture_output=True, text=True) + out = r.stderr + r.stdout + regs = re.search(r"Used (\d+) registers", out) + spill = re.search(r"(\d+) bytes spill stores, (\d+) bytes spill loads", out) + stack = re.search(r"(\d+) bytes stack frame", out) + n_ex2 = ptx_text.count("ex2.approx") + print( + f"{tag:36} regs={regs.group(1) if regs else '?':>3} " + f"spill={(spill.group(1) + '/' + spill.group(2)) if spill else '?'} " + f"stack={stack.group(1) if stack else '?'} ex2.approx={n_ex2:3d} rc={r.returncode}", + flush=True, + ) + clean = r.returncode == 0 and spill is not None and stack is not None and spill.group(1) == spill.group(2) == "0" and stack.group(1) == "0" + return clean + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--arch", default="sm_100a", help="ptxas target architecture (default: sm_100a)") + ap.add_argument("--keep-ptx", default=None, help="directory to keep the per-kernel .ptx files in (default: temporary)") + args = ap.parse_args() + # CUTE_DSL_KEEP=ptx makes the DSL drop each kernel's PTX into the current + # directory; run the compiles from the (possibly temporary) output directory so + # the repository tree stays clean. + out_dir = os.path.abspath(args.keep_ptx) if args.keep_ptx else tempfile.mkdtemp(prefix="csa_r128_ptx_") + os.makedirs(out_dir, exist_ok=True) + os.chdir(out_dir) + from cudnn.csa.compressor import compressor_sm100_r128 as M + + assert torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0), "requires a CC 10.0 GPU" + dev = torch.device("cuda", torch.cuda.current_device()) + # Compile every schedule bucket each shipped config can select at runtime + # (precompile with nb_total=None walks the small/default/large tables). + for coff, d in [(1, 128), (2, 128), (1, 512), (2, 512)]: + M.precompile_fwd_r128(128, d, coff, dev) + M.precompile_bwd_r128(128, d, coff, dev) + + all_clean = True + n = 0 + for key, fn in sorted(M._COMPILED.items(), key=str): + kind, _ratio, d, coff, sched, _dev = key + if kind == "r128fwd": + vec, tchunks, threads_x, twophase, fastexp = sched + tag = f"fwd_c{coff}d{d}_v{vec}t{tchunks}x{threads_x}" + ("_2ph" if twophase else "") + ("_fexp" if fastexp else "") + else: + vec, tchunks, threads_x, fastexp = sched + tag = f"bwd_c{coff}d{d}_v{vec}t{tchunks}x{threads_x}" + ("_fexp" if fastexp else "") + all_clean = ptxas_one(tag, fn.artifacts.PTX, args.arch, out_dir) and all_clean + n += 1 + print(f"{n} kernels probed ({args.arch}); {'ALL 0 spill / 0 stack' if all_clean else 'SPILL/STACK OR PTXAS FAILURE DETECTED'}", flush=True) + if args.keep_ptx: + print(f"PTX kept in {out_dir}", flush=True) + else: + os.chdir(tempfile.gettempdir()) + shutil.rmtree(out_dir, ignore_errors=True) + sys.exit(0 if all_clean else 1) + + +if __name__ == "__main__": + main() diff --git a/docs/fe-oss-apis/csa.md b/docs/fe-oss-apis/csa.md new file mode 100644 index 000000000..c9985d877 --- /dev/null +++ b/docs/fe-oss-apis/csa.md @@ -0,0 +1,419 @@ +# CSA Fused Compressor + +**This is an experimental API and subject to change.** + +## Overview + +The CSA module hosts CuTe-DSL kernels for the CSA/HCA experimental attention variants +(the components that are not shared with the [DSA module](dsa.md)). Its first operation +is the **fused Compressor**: one forward and one backward kernel for the `Compressor` +gated-softmax pooling region (THD packed layout) used by CSA/HCA in Megatron-LM. + +The kernels were ported from Megatron-LM at the maintainers' request +([Megatron-LM PR #5984](https://github.com/NVIDIA/Megatron-LM/pull/5984); measurements +and numerics in +[Megatron-LM issue #5968](https://github.com/NVIDIA/Megatron-LM/issues/5968)). The eager +region they replace decomposes into ~39 forward and ~51 backward kernel launches per +call (at `compress_ratio = 4`, `coff = 2`) and materializes `(total_comp, 2*ratio, 1, head_dim)` +window intermediates; the fused path is 1 + 1 kernels (plus one fp32 `dAPE`-buffer +zero-fill in backward — the backward kernel writes `dKV`/`dScore` in full, including +exact zeros to never-consumed positions, so those buffers need no fills). + +### Semantics + +For each THD segment `s` (`cu_seqlens[s]..cu_seqlens[s+1]`) and each output block `b` of +`ratio` tokens, with the overlapping window form (`coff == 2`, window size `2 * ratio`): + +- `k in [0, ratio)`: previous block's token, first-half projection column, APE row `k` + — invalid for the segment's first block (score `-inf`, kv `0`); +- `k in [ratio, 2*ratio)`: own block's token, second-half projection column, APE row + `k - ratio`. + +The own-block window form (`coff == 1`, window size `ratio`) drops the overlap: every +`k in [0, ratio)` is the block's own token on projection column `j` with APE row `k`, and +every window is fully valid (no first-block exception). + +```text +out[b, j] = sum_k kv[w(b,k), c(k,j)] * softmax_k(score[w(b,k), c(k,j)] + ape[k % ratio, c(k,j)]) +``` + +computed in fp32 with a single final bf16 rounding. Per-segment tail tokens +(`seqlen % ratio`) are dropped, as in the eager code. Output rows beyond +`cu_seqlens_comp[-1]` (a static CUDA-graph capacity) are computed with first-in-segment +semantics from token 0, exactly like the eager gather; the backward ignores incoming +gradients on such padding rows. + +### Numerics + +All arithmetic is fp32 with one final bf16 rounding; `mul.rn.f32` / `fma.rn.f32` are +pinned in PTX so results do not depend on compiler FMA contraction. The numerics +contract is **per ratio family** (the two families intentionally differ — do not +assume the ratio=4 guarantees at ratio=128): + +- **`ratio == 4`** (production, unchanged): against an fp32-intermediate eager + reference (same op order, fp32 throughout), `dKV`/`dScore` are **bit-identical** + and the forward matches within one bf16 rounding step on a tiny fraction of + elements. +- **`ratio == 128`** (deterministic tolerance contract): the kernels are + **deterministic and faithful to the fp32-intermediate eager reference** (the + same eager region computed with fp32 intermediates and one final bf16 rounding — + the comparator every number below is measured against): + + 1. **Determinism** — forward `out` and backward `dKV`/`dScore` are bitwise + run-to-run deterministic on all supported inputs (see below; NaN-prefill + replay tested, including bit-stable replay of the NaN pattern on the gate's + overflow-intermediate case), + 2. **Same values within final-bf16 rounding on the gate tolerances** — + `out`/`dKV`/`dScore` match the eager reference within differing elements + `<= max(1, 0.1%)` of the tensor and `max_abs <= 1.6e-2`, thresholds + **calibrated on the gate's documented input distribution** (`kv`, `grad_out` + ~ N(0, 1) bf16, `score` ~ N(0, 1.5²) bf16, `ape` ~ N(0, 0.25²) fp32). + Absolute bf16 deviations are NOT magnitude-free — bf16's grid is relative, so + scaling `kv`/`grad_out` by `2^k` scales every deviation exactly `2^k`, with + differing-element counts and fp64 parity unchanged for as long as the scaled + inputs keep every fp32 intermediate finite (see 3.). The gate's ×64 + scaled-input case commits that as evidence (deviations ×64 forward/`dKV`, + ×4096 `dScore`; counts and parity identical to its unit-scale twin), and + 3. **The eager reference's non-finite propagation** — both sides compute the + window math in fp32, so finite inputs that overflow the eager reference's + fp32 intermediates (a `score + ape` sum or a backward `kv · grad_out` product + beyond fp32 range — elementwise ops both sides evaluate identically) poison + the kernels' outputs as they poison the reference's, instead of the kernels + returning clean-looking finite values. The gate's overflow-intermediate case + commits this: `score` = bf16 max with `ape` = fp32 max drives fp32 + `score + ape` to +Inf, every fused output is NaN with the NaN mask equal to + the eager reference's on all four outputs, and the pattern replays + bit-stable; the fp64 comparator is explicitly skipped there because the fp64 + oracle stays finite (a NaN-vs-finite distance gates nothing). The mirror + caveat: the fused evaluation order saturates EARLIER than the eager one near + fp32 max — its forward chunk partials are un-normalized (up to `2 · ratio ×` + the reference's normalized weighted sums), so e.g. uniform-score inputs with + `|kv|` within a factor `~2 · ratio` of fp32 max return ±Inf where the eager + reference is still finite. + + On inputs whose fp32 intermediates stay finite in BOTH evaluation orders — the + eager reference's and the fused kernels' (which saturate earlier, see 3.) — + `out`/`dKV`/`dScore` additionally carry **fp64-oracle parity**: per case and per + tensor they are **at least as close to an fp64 oracle as the fp32-intermediate + eager reference itself** (`err_fused <= err_eager * (1 + 1e-6) + 1e-4`, asserted + in the tests and on every finite-intermediate case of the committed gate + script): the fused reorders and fast-exp buckets do not lose accuracy against + the eager reference on the tested envelope. + + Worst observed on the gate distribution (19 unit-scale cases = 15 union — the + 128k-token d=512 packs select the large-bucket schedules — + 4 padding, B200; + reproduce with `benchmark/csa/gate_csa_compressor_r128.py`, which also asserts + the case list selects every shipped schedule): forward 0.0275% differing elements + (9/32,768), `dKV` 0.0041% (2,770/67,108,864), `dScore` 0.0113% (104/921,856); + worst `max_abs` 3.91e-3 for forward/`dKV` (the 131k-token d=512 cases) and + 1.95e-3 for `dScore` (bf16-rounding magnitudes). The deviations come from + measured-win reduction reorders (forward chunked-softmax merges; the backward's + fused `den`/`S` partial merge with a hoisted `1/den`) and an `ex2.approx`-based + fast exp enabled on specific schedule buckets in BOTH directions (forward: + per-bucket table entries; backward: the default outside the d=128 vec=1 + small-pack buckets). `dAPE` (fp32, atomic accumulation) is gated at + `max_abs <= 1e-3` vs the fp32 eager reference on the gate distribution — observed + worst 2.0e-6, run-to-run replay deltas of the same magnitude — and scales the + same way (observed 2.0e-3 under the ×64 case). `dAPE`'s distance from the fp64 + oracle fluctuates run to run around the eager reference's own (within ~1.5× + either way — atomic reduction-order noise), so its parity check is asserted on + the unit-scale cases, where the formula's absolute term dominates, and recorded + on scaled inputs. + +In BOTH families the forward, `dKV` and `dScore` are **bitwise run-to-run +deterministic** (fixed chunk boundaries and merge orders, no atomics in those +outputs; NaN-prefill replay is part of the test suite). `dAPE` is reduced with one +fp32 atomic per `(k, dim)` per CTA in both families and is **not** bitwise +run-to-run deterministic; the backward APIs raise under +`torch.use_deterministic_algorithms(True)` (warn-only mode warns and runs) — callers +that need a fully deterministic backward must use an eager implementation. + +### Support surface (`check_support`) + +- Compute capability **10.0** (the only validated architecture so far; the kernels use + no arch-specific features, wider enablement is possible after validation) +- `ratio == 4`, `coff in {1, 2}` (`coff == 2` is the production CSA/HCA configuration, + `coff == 1` the own-block window form) — served by the generic kernels, which are + generic over `(ratio, head_dim, coff in {1, 2})` but keep the whole pooling window in + registers (register-bound beyond `ratio = 32`) +- `ratio == 128`, `coff in {1, 2}`, `head_dim in {128, 512}` — served by dedicated + kernels (see [ratio=128](#ratio128) below); the wrappers and class APIs + route by `ratio` transparently, and the gate can be widened per configuration once + validated. **The ratio=128 numerics contract differs** (see [Numerics](#numerics)) +- BF16 `kv` / `score` / `out`, FP32 `ape`, int32 `cu_seqlens` / `cu_seqlens_comp` +- int32 flat offsets: `total_tokens * coff * head_dim < 2**31` and + `total_comp * head_dim < 2**31` +- `total_comp > 0` requires `total_tokens >= ratio` (each compressed row gathers a + window of `ratio` tokens) +- `head_dim <= 8388480` (forward launch `gridDim.y` bound) +- contiguous tensors on one CUDA device, with 16-byte-aligned base pointers (4-byte + for the int32 cu_seqlens) — contiguity does not imply base alignment for + storage-offset views, so this is checked per call + +## Installation + +```bash +pip install nvidia-cudnn-frontend[cutedsl] +``` + +## API Usage + +### High-level wrappers + +```python +from cudnn import CSA + +# forward: (total_tokens, coff*head_dim) BF16 kv/score, (ratio, coff*head_dim) FP32 ape, +# (B+1,) int32 cu_seqlens / cu_seqlens_comp +result = CSA.csa_compressor_forward_wrapper( + kv, score, ape, cu_seqlens, cu_seqlens_comp, + ratio=4, head_dim=128, coff=2, + total_comp=None, # defaults to cu_seqlens_comp[-1] (synchronizes); pass a static + # capacity explicitly to stay CUDA-graph capture-safe + stream=None, +) +out = result["out"] # (total_comp, head_dim) BF16 + +grads = CSA.csa_compressor_backward_wrapper( + kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, + ratio=4, head_dim=128, coff=2, stream=None, +) +grad_kv, grad_score, grad_ape = grads # BF16, BF16, FP32 +``` + +Set `coff=1` for the `ratio`-token own-block window (and use `kv` / `score` tensors whose +packed width is `head_dim`); `coff=2` selects the `2 * ratio` overlapping window shown +above. + +The wrappers cache compiled API instances; the underlying JIT is shared per +`(ratio, head_dim, coff, device)`, so runtime shape changes never recompile. + +### Class API + +```python +from cudnn import CSACompressorForward, CSACompressorBackward + +op = CSACompressorForward( + sample_kv, sample_score, sample_ape, sample_cu_seqlens, sample_cu_seqlens_comp, + sample_out, ratio=4, coff=2, +) +op.check_support() +op.compile() +op.execute(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, current_stream=None) +``` + +`CSACompressorBackward.execute` additionally takes `grad_out` and the +`grad_kv` / `grad_score` / `grad_ape` buffers. `grad_kv` / `grad_score` may be +**uninitialized**: the kernel writes every position (disjoint, atomic-free stores; +never-consumed positions — segment tails, the last block's first-half columns +(`coff == 2` only; `coff == 1` has no first-half columns), segments shorter than `ratio`, +token-capacity padding beyond `cu_seqlens[-1]` — get exact zeros +from their unique owning CTA, matching autograd). When `total_comp == 0` the kernel is +not launched and the buffers are left untouched (zero them yourself if you need +autograd's exact zeros; the high-level wrapper does). `grad_ape` must be +**zero-initialized before every `execute` call and before every CUDA-graph replay +that reuses the buffer** — the kernel only accumulates into it (fp32 atomics) and +never clears it. (The high-level wrapper allocates a fresh zeroed `grad_ape` per +call; the zero-fill is captured together with the kernel, so wrapper graph replays +re-zero automatically.) + +### CUDA graphs + +The launch path is capture-compatible once the kernels for a `(ratio, head_dim, coff)` +configuration are compiled: run one warmup call (or `compile()`) per configuration +before capturing, and pass `total_comp` explicitly. A call that would JIT under capture +raises a `RuntimeError` instead of corrupting the capture. + +### Environment variables + +- `CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH=0` — disable the cached-launch host optimization + (a per-config snapshot of the CuTe-DSL launch state, replayed with in-place argument + mutation; it removes tens of microseconds of per-call host overhead for these + microsecond-scale kernels). The snapshot construction introspects + private-but-stable DSL internals; on any structural mismatch (e.g. a future + `nvidia-cutlass-dsl` upgrade) it falls back to the regular launch path automatically. + +## Performance + +Measured on 1x B200 (CC 10.0, driver 590.48.01); BF16 `kv`/`score`, FP32 `ape`; `ratio = 4`, `coff = 2`; +THD packs of 8192-token sequences; eager baseline = the exact replaced region of +Megatron-LM `Compressor._forward_thd` on identical inputs. + +*Isolated GPU kernel time* (nsys, sum of kernel durations per iteration, 50 iterations +after 20 warmup; no launch/host overhead; backward includes its `dAPE` zero-fill): + +| THD pack | head_dim | eager fwd | fused fwd | fwd | eager bwd | fused bwd | bwd | +|---|---|---|---|---|---|---|---| +| 1 x 8192 | 128 | 117.8 us | 4.5 us | **26.5x** | 187.2 us | 12.8 us | **14.6x** | +| 3 x 8192 | 128 | 229.8 us | 10.0 us | **23.0x** | 352.7 us | 22.2 us | **15.9x** | +| 1 x 8192 | 512 | 263.3 us | 12.4 us | **21.2x** | 425.0 us | 22.8 us | **18.6x** | +| 3 x 8192 | 512 | 664.3 us | 35.0 us | **19.0x** | 1155.8 us | 66.0 us | **17.5x** | + +*End-to-end wall clock of the same region* (CUDA events, median of 100 after 30 warmup; +includes launch overhead; eager backward goes through torch autograd with the forward +outside the timed region, fused backward is the explicit backward wrapper call — not +comparable to the kernel-time numbers above): + +| THD pack | head_dim | eager fwd | fused fwd | fwd | eager bwd | fused bwd | bwd | +|---|---|---|---|---|---|---|---| +| 1 x 8192 | 128 | 343.7 us | 37.4 us | **9.2x** | 558.7 us | 51.9 us | **10.8x** | +| 3 x 8192 | 128 | 389.3 us | 38.2 us | **10.2x** | 585.8 us | 62.4 us | **9.4x** | +| 1 x 8192 | 512 | 423.3 us | 39.9 us | **10.6x** | 666.5 us | 62.8 us | **10.6x** | +| 3 x 8192 | 512 | 831.7 us | 62.7 us | **13.3x** | 1503.6 us | 108.3 us | **13.9x** | + +The previously published per-call wall clock (commit `b3ceb7c`) was `333.7 / 37.8 / 506.7 +/ 47.2 us` at `1 x 8192 / 128` (and analogously for the other packs). The re-run above +reproduces the forward columns within ~±4% and the fused/backward columns within +2-14%: +the fused per-call wall clock carries the cached-launcher host optimization, whose +snapshot replay jitters run-to-run (the published fused-forward is even non-monotonic in +pack size, `37.8 -> 35.1 us`), and eager backward drifts with GPU boost-clock state. The +CUDA-graph replay numbers below collapse that host jitter into a single replay. + +*CUDA-graph replay of the same region, both implementations captured symmetrically as a +forward-only graph and a forward+backward graph* (median of 100 after 30 warmup, replay +timing via `benchmark/csa/bench_csa_compressor.py`; capturing each side into a graph +collapses its per-operation launches into a single replay, so this is the fairest +wall-clock basis for launch-bound shapes): + +| THD pack | head_dim | eager fwd | fused fwd | fwd | eager total | fused total | total | +|---|---|---|---|---|---|---|---| +| 1 x 8192 | 128 | 119.2 us | 10.8 us | **11.0x** | 364.6 us | 22.8 us | **15.9x** | +| 3 x 8192 | 128 | 235.3 us | 14.8 us | **15.8x** | 668.4 us | 38.6 us | **17.3x** | +| 1 x 8192 | 512 | 269.8 us | 17.1 us | **15.7x** | 802.5 us | 40.0 us | **20.0x** | +| 3 x 8192 | 512 | 673.1 us | 39.2 us | **17.1x** | 2057.6 us | 107.8 us | **19.0x** | + +`eager total` / `fused total` capture forward + backward together. The eager total graph +captures the autograd backward of the captured forward against stable, pre-allocated zero +`.grad` buffers (the captured region zeros them in place, then runs forward + backward, so +every replay accumulates into a zeroed buffer — numerically identical to a single fresh +backward, verified per shape in the harness); the fused total graph captures the forward +wrapper immediately followed by the backward wrapper. The backward replay alone is not a +separately captured quantity — it is approximately `total - fwd` and is reported only as +that reference, never as a measured column. Graph speedups are `eager / fused` of the +displayed µs, truncated to one decimal. + +Capturing each side into a graph narrows the absolute eager-vs-fused forward gap most on +the smallest, launch-bound shape (about 2.8x at `1 x 8192 / 128` forward) and barely on +the largest, less launch-bound shape (1.2x at `3 x 8192 / 512`); the speedup ratio is +similar to or larger than per-call because a graph replay collapses each side's +per-operation launches into one. + +Environment: driver 590.48.01, PyTorch 2.13.0 (CUDA 13.3), `nvidia-cutlass-dsl` 4.6.1. +Measurement basis: identical inputs over exactly the replaced region for both +implementations; per-call eager backward = the torch autograd backward with the forward +outside the timed region; graph eager backward = the autograd backward of the captured +forward, captured together with it against stable zero `.grad` buffers; fused backward = +the backward wrapper (kernel + the fp32 `dAPE` zero-fill + host validation, no autograd +engine). Both wall-clock tables above (per-call and graph) are from a single run of +`benchmark/csa/bench_csa_compressor.py`; the kernel-time table is from nsys. + +An ncu hardware-ceiling audit of the **ported kernels (prior to the two optimization +commits)** — cache-flushed, `--set full`, all four benchmark shapes — showed measured +DRAM read volume matching the algorithmically necessary bytes within 1% (the THD gather +adds no over-read; stores fully coalesced at 32/32 bytes per sector, loads 29-30/32), +with neither L2 (<27% of peak) nor DRAM (<33%) close to saturation at these +microsecond-scale sizes: the gap to a pure DRAM-floor time was a mix of sub-wave grid +width / occupancy, memory latency, and (at the largest shape) issue pressure — not +wasted traffic. The two optimizations that audit identified are folded in here: 32-bit +vectorized forward accesses and backward kernel-side zero-writes replacing the two bf16 +grad-buffer fills. They do not change the bytes the kernels must read, so the +traffic-optimality conclusion carries over; the utilization percentages above predate +them. + +## ratio=128 + +`ratio = 128` is served by dedicated kernels (`compressor_sm100_r128.py`): the generic +kernels keep the whole `coff * ratio` window in per-thread registers, which spills +kilobytes of local memory per thread at `ratio = 128` (255-register cap). The dedicated +forward streams the window with a chunked softmax (one CTA per output row, the window +split over `threadIdx.y` chunk-rows, per-column `(max, denom, acc)` triples, one +fixed-order smem merge); its launch schedule is **bucketed by the output-row count** +(small / default / large, all precompiled for CUDA-graph safety) and per bucket selects +the online-rescale or two-phase accumulation form and, where it measured faster, an +`ex2.approx`-based fast exp. The dedicated backward stages each row's window into +shared memory chunk-parallel, computes `exp` and per-chunk partial `den`/`S` sums in +the same pass, merges the partials per column in a fixed chunk order, and stores the +gradients with a hoisted `1/den` multiply — plus kernel-side zero-writes to +never-consumed slots and fp32-atomic `dAPE`, exactly as at ratio=4. The backward +defaults to the same fast exp outside the d=128 small-pack (vec=1) buckets. + +**The ratio=128 numerics contract is the deterministic tolerance contract described in +[Numerics](#numerics), NOT the ratio=4 bitwise-`dKV`/`dScore` contract.** The +reduction orders (forward chunk merge, backward `den`/`S` partial merge), the +backward's hoisted reciprocal, and the fast-exp buckets all differ from the eager op +order by design, each adopted on a measured same-GPU win and gated on tolerance + +fp64-oracle parity (worst observed deviations in [Numerics](#numerics)). Everything +stays bitwise run-to-run: fixed chunk boundaries and merge orders, no atomics in +forward/`dKV`/`dScore`. + +Both kernels are register-flat in the window length (ptxas sm_100a: forward 32-51 +registers, backward 48-128 registers, 0 spill / 0 stack across every shipped +(config, schedule) kernel — 16 kernels total; reproduce the per-kernel table with +`benchmark/csa/reg_probe_csa_compressor_r128.py`) and JIT in ~0.4-1.0 s per +configuration. +The backward picks its `rows_per_cta` at launch time (a runtime argument — no +recompile) so the grid fits one resident wave; the static row capacity fixes it under +CUDA-graph capture. Small packs (`nb_total <= 192`, d=128) switch the backward to a +vec=1 schedule bucket (also precompiled) for grid-fill. + +Measured on 1x B200 (CC 10.0, driver 590.48.01), torch 2.13.0 / CUDA 13.3 / +`nvidia-cutlass-dsl` 4.6.1; eager baseline = the fp32-intermediate reference region of +the test suite on identical inputs (the numerics-contract reference; the upstream eager +region differs only by a bf16 weight-rounding cast at ~equal cost). Packs are +single-sequence THD unless noted. + +*Isolated GPU kernel time* (nsys, sum of kernel durations inside a +`cudaProfilerApi`-gated 100-iteration unsynced loop / 100; forward and backward each +measured in their own loop; the fwd+bwd total is one combined loop per training step — +forward + backward + the ~1 us fp32 `dAPE` zero-fill — and is not the sum of the +standalone columns, because the interleaved backward changes L2 residency for the +forward at long context): + +| config | tokens | eager fwd | fused fwd | fwd | eager bwd | fused bwd | bwd | fused fwd+bwd total | +|---|---|---|---|---|---|---|---|---| +| coff 1, d 128 | 8192 | 84.7 us | 5.7 us | **15.0x** | 127.3 us | 8.4 us | **15.2x** | 15.1 us | +| coff 1, d 128 | 131072 | 451.8 us | 14.9 us | **30.2x** | 471.9 us | 63.7 us | **7.4x** | 82.5 us | +| coff 2, d 128 | 8192 | 187.1 us | 6.4 us | **29.1x** | 246.2 us | 11.2 us | **21.9x** | 18.9 us | +| coff 2, d 128 | 65536 | 726.3 us | 16.4 us | **44.4x** | 960.0 us | 61.2 us | **15.7x** | 76.7 us | +| coff 1, d 512 | 65536 | 641.1 us | 48.4 us | **13.2x** | 795.0 us | 104.2 us | **7.6x** | 150.0 us | +| coff 2, d 512 | 65536 | 2315.2 us | 78.8 us | **29.4x** | 3261.9 us | 205.8 us | **15.9x** | 287.7 us | +| coff 2, d 512 | 131072 | 4526.5 us | 117.6 us | **38.5x** | 6304.3 us | 457.6 us | **13.8x** | 580.7 us | + +*End-to-end wall clock per call* (CUDA events, median of 100 after 30 warmup; fused +backward includes a `grad_ape.zero_()` per call; not comparable to the kernel-time +numbers above — at small packs the ~9 us python launch path dominates the 5.7-11 us +kernels): + +| config | tokens | fused fwd | fused bwd | +|---|---|---|---| +| coff 1, d 128 | 8192 | 14.1 us | 24.3 us | +| coff 1, d 128 | 131072 | 24.4 us | 82.1 us | +| coff 2, d 128 | 8192 | 15.2 us | 26.3 us | +| coff 2, d 128 | 65536 | 25.5 us | 75.6 us | +| coff 1, d 512 | 65536 | 56.3 us | 118.9 us | +| coff 2, d 512 | 65536 | 88.1 us | 220.8 us | + +The ratio=128 numerics contract (21-case gate: 15 union + 1 scaled-input + 4 padding +cases, each running the determinism, tolerance and fp64-parity gates, plus 1 +overflow-intermediate case that gates bit-stable NaN replay and NaN-pattern equality +vs the eager fp32 reference with the fp64 comparator explicitly skipped, with shipped +schedule coverage asserted) and the per-kernel ptxas table are reproducible from the +committed scripts `benchmark/csa/gate_csa_compressor_r128.py` and +`benchmark/csa/reg_probe_csa_compressor_r128.py`. + +## Testing + +```bash +(cd test/python && pytest fe_api/csa/test_CSA_compressor.py) +``` + +The tests validate numerics against an fp32-intermediate eager reference (bitwise +`dKV`/`dScore` at ratio=4; the deterministic tolerance contract at ratio=128, with +fp64-oracle parity on finite-intermediate inputs), the upstream eager numerics, plus ragged packs, static-capacity +padding, kernel-side zero-writes into uninitialized gradient buffers (NaN-canary with +exact-zero assertions on every never-consumed slot class, and the `total_comp == 0` +zeros fallback), run-to-run determinism, `grad_ape` zeroing ownership, the ratio=128 +dispatch envelope (schedule selection at every `nb_total` bucket boundary at L0, and +one execution of every shipped (config, schedule) kernel against the full contract — +tolerance, fp64-oracle parity, determinism — at L1; run +`pytest -m "L0 or L1"`), CUDA-graph capture/replay, and `check_support` boundaries. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 2bc3ccfb2..3c25c11cd 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -22,6 +22,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [Grouped GEMM + Wgrad](gemm_fusions/grouped_gemm_wgrad.md) - [Block Sparse Attention (BSA)](bsa.md) - [Native Sparse Attention (NSA)](nsa.md) +- [CSA Fused Compressor](csa.md) - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md) - [SDPA Forward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-forward-fe-oss-sm100-d256) - [SDPA Backward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-backward-fe-oss-sm100-d256) diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 6e8b99829..a5256c866 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -282,6 +282,11 @@ def _dlopen_cudnn(): "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), "block_sparse_attention_backward": (".block_sparse_attention", "block_sparse_attention_backward"), "DSA": (".deepseek_sparse_attention", "DSA"), + "CSA": (".csa", "CSA"), + "CSACompressorForward": (".csa", "CSACompressorForward"), + "CSACompressorBackward": (".csa", "CSACompressorBackward"), + "csa_compressor_forward_wrapper": (".csa", "csa_compressor_forward_wrapper"), + "csa_compressor_backward_wrapper": (".csa", "csa_compressor_backward_wrapper"), "NSA": (".native_sparse_attention", "NSA"), "GemmSwigluSm100": (".gemm_swiglu", "GemmSwigluSm100"), "gemm_swiglu_wrapper_sm100": (".gemm_swiglu", "gemm_swiglu_wrapper_sm100"), diff --git a/python/cudnn/csa/README.md b/python/cudnn/csa/README.md new file mode 100644 index 000000000..47ab06523 --- /dev/null +++ b/python/cudnn/csa/README.md @@ -0,0 +1,19 @@ +# CSA module + +Fused CuTe-DSL kernels for the CSA/HCA experimental attention variants (the components +that are not shared with the DSA module, which lives in +`python/cudnn/deepseek_sparse_attention/`). + +- **Compressor**: fused forward+backward kernels for the `Compressor` gated-softmax + pooling (THD packed layout): gather -> `+ APE` -> optional overlap-window transform + (`coff == 2`) -> fp32 softmax -> gated weighted sum -> bf16 cast, as one kernel per + direction. Ported from + Megatron-LM ([PR #5984](https://github.com/NVIDIA/Megatron-LM/pull/5984), measurements + in [issue #5968](https://github.com/NVIDIA/Megatron-LM/issues/5968)). See + [docs/fe-oss-apis/csa.md](../../../docs/fe-oss-apis/csa.md). + +## Acknowledgements + +The fused Compressor kernels were contributed by the GLM training-performance team +(Zhipu AI). The CSA/HCA attention variants and the surrounding DSA/CSA kernel family are +by Hongxiao Bai, Jiayu Sun and Jie Fang. diff --git a/python/cudnn/csa/__init__.py b/python/cudnn/csa/__init__.py new file mode 100644 index 000000000..e133f27c2 --- /dev/null +++ b/python/cudnn/csa/__init__.py @@ -0,0 +1,53 @@ +"""``cudnn.csa`` — CuTe-DSL kernels for the CSA/HCA experimental attention variants. + +Symbols (the fused ``Compressor`` APIs) resolve lazily on first attribute access, so +importing ``cudnn`` never pulls in the optional ``[cutedsl]`` dependency stack. +""" + +from importlib import import_module + +_SYMBOLS = { + "CSACompressorForward": (".compressor", "CSACompressorForward"), + "CSACompressorBackward": (".compressor", "CSACompressorBackward"), + "csa_compressor_forward_wrapper": (".compressor", "csa_compressor_forward_wrapper"), + "csa_compressor_backward_wrapper": (".compressor", "csa_compressor_backward_wrapper"), +} + + +def _load_symbol(name): + """Import the symbol behind lazy attribute ``name`` and cache it in module globals.""" + module_name, symbol_name = _SYMBOLS[name] + module = import_module(module_name, package=__name__) + symbol = getattr(module, symbol_name) + globals()[name] = symbol + return symbol + + +def __getattr__(name): + """Resolve the lazily exported symbols and the ``CSA`` namespace (PEP 562).""" + if name == "CSA": + return CSA + if name in _SYMBOLS: + return _load_symbol(name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +class CSANamespace: + """Namespace object mirroring the package's lazy symbols (``cudnn.CSA.``).""" + + def __getattr__(self, name): + """Lazily resolve ``CSA.`` through the package's symbol table.""" + if name in _SYMBOLS: + return _load_symbol(name) + raise AttributeError(f"CSA has no attribute {name!r}") + + +CSA = CSANamespace() + +__all__ = [ + "CSA", + "CSACompressorBackward", + "CSACompressorForward", + "csa_compressor_backward_wrapper", + "csa_compressor_forward_wrapper", +] diff --git a/python/cudnn/csa/compressor/__init__.py b/python/cudnn/csa/compressor/__init__.py new file mode 100644 index 000000000..6788c0f17 --- /dev/null +++ b/python/cudnn/csa/compressor/__init__.py @@ -0,0 +1,15 @@ +"""Public surface for the fused CSA/HCA Compressor kernels (re-exports from ``.api``).""" + +from .api import ( + CSACompressorForward, + CSACompressorBackward, + csa_compressor_forward_wrapper, + csa_compressor_backward_wrapper, +) + +__all__ = [ + "CSACompressorBackward", + "CSACompressorForward", + "csa_compressor_backward_wrapper", + "csa_compressor_forward_wrapper", +] diff --git a/python/cudnn/csa/compressor/api.py b/python/cudnn/csa/compressor/api.py new file mode 100644 index 000000000..fe3bd29ff --- /dev/null +++ b/python/cudnn/csa/compressor/api.py @@ -0,0 +1,670 @@ +"""APIBase wrappers for the fused CSA/HCA Compressor gated-pooling CuTe-DSL kernels. + +``CSACompressorForward`` / ``CSACompressorBackward`` wrap the forward and backward +kernels in ``compressor_sm100.py`` (ported from Megatron-LM, see +https://github.com/NVIDIA/Megatron-LM/pull/5984 and +https://github.com/NVIDIA/Megatron-LM/issues/5968). The kernels fuse the gated-softmax +pooling region of the CSA/HCA ``Compressor`` for the THD packed layout: + + out[b, j] = sum_k kv[w(b, k), c(k, j)] * softmax_k(score[w(b, k), c(k, j)] + ape[k % ratio, c(k, j)]) + +over the per-block window ``w``: ``2 * ratio`` entries for the overlapping ``coff == 2`` +form (the previous block's half-window is invalid for each segment's first block), or +``ratio`` entries for the own-block ``coff == 1`` form (no overlap, every window valid). +The framework-side autograd wiring stays in the caller (e.g. a ``torch.autograd.Function`` +that calls the forward wrapper in ``forward()`` and the backward wrapper in +``backward()``); these APIs are pure kernels-plus-validation. + +Validated envelope (``check_support``): compute capability 10.0, BF16 ``kv``/``score``/ +``out``, FP32 ``ape``, int32 ``cu_seqlens``/``cu_seqlens_comp``, int32 flat offsets +(``total_tokens * coff * head_dim < 2**31``), and per ratio: + +- ``ratio == 4``, ``coff in {1, 2}`` (``coff == 2`` is the production CSA/HCA + configuration, ``coff == 1`` the own-block window form) — served by the generic + kernels in ``compressor_sm100.py`` (whole window in registers; optimal at small + ratios, register-bound beyond ``ratio = 32``); +- ``ratio == 128``, ``coff in {1, 2}``, ``head_dim in {128, 512}`` — served by the + dedicated kernels in ``compressor_sm100_r128.py`` (bucketed-schedule chunked-softmax + forward; staged smem backward with fused per-chunk reductions). The wrappers route + by ``ratio`` transparently. NOTE: the numerics contracts differ per family — see + below. + +Numerics contract (see the kernel modules and docs/fe-oss-apis/csa.md for details): +fp32 arithmetic with one final bf16 rounding, ``mul.rn``/``fma.rn`` pinned in PTX. +Forward, ``dKV`` and ``dScore`` are bitwise run-to-run deterministic in BOTH families. +At ``ratio == 4`` dKV/dScore are additionally bit-identical to the fp32-intermediate +eager autograd; at ``ratio == 128`` the contract is faithfulness to that +fp32-intermediate eager reference: out/dKV/dScore match it within final-bf16 rounding +at the gate tolerances (differing elements <= max(1, 0.1%), max_abs <= 1.6e-2, +calibrated on the documented gate input distribution — absolute bf16 deviations scale +with the input magnitudes), inputs that overflow the reference's fp32 intermediates +reproduce its NaN/Inf propagation (both sides compute in fp32; gate-tested), and on +inputs whose fp32 intermediates stay finite those outputs additionally carry +fp64-oracle parity (at least as close to an fp64 oracle as the eager reference). That +is the approved deterministic tolerance contract (reduction reorders + fast-exp +buckets). +``dAPE`` uses one fp32 atomic per ``(k, dim)`` per CTA in both families and is not +run-to-run deterministic (the backward APIs refuse to run under +``torch.use_deterministic_algorithms(True)``). +""" + +from __future__ import annotations + +import threading +import warnings +from collections import OrderedDict +from contextlib import contextmanager +from typing import Iterator, Optional + +import torch +import cuda.bindings.driver as cuda + +from cudnn.api_base import APIBase, TupleDict + +from .compressor_sm100 import ( + CU_ALIGN_BYTES, + PTR_ALIGN_BYTES, + SUPPORTED_COMPUTE_CAPABILITY, + precompile_bwd, + precompile_fwd, + run_bwd, + run_fwd, +) +from .compressor_sm100_r128 import ( + precompile_bwd_r128, + precompile_fwd_r128, + run_bwd_r128, + run_fwd_r128, +) + +# int32 flat offsets: every element offset the kernels compute must fit in int32. +_INT32_LIMIT = 2**31 +# Forward launch schedule gridDim.y bound: at 128 threads per column CTA the largest +# launchable head_dim is 128 * 65535 (identical for the 64-thread vec == 2 path, which +# halves the column count). +_MAX_HEAD_DIM = 128 * 65535 +# Bound + eviction policy follow python/cudnn/graph.py's graph_cache precedent. +_API_CACHE_MAXSIZE = 256 + + +class _LruDict: + """Bounded thread-safe LRU mapping (per ``graph.py``'s ``graph_cache`` precedent).""" + + def __init__(self, maxsize: int = _API_CACHE_MAXSIZE): + """Create an empty LRU mapping evicting past ``maxsize`` entries.""" + self._data: OrderedDict = OrderedDict() + self._lock = threading.Lock() + self._maxsize = maxsize + + def get(self, key, default=None): + """Return the value for ``key`` (refreshing its recency) or ``default``.""" + with self._lock: + if key not in self._data: + return default + self._data.move_to_end(key) + return self._data[key] + + def put(self, key, value) -> None: + """Insert ``key -> value`` as most recent, evicting past ``maxsize``.""" + with self._lock: + self._data[key] = value + self._data.move_to_end(key) + while len(self._data) > self._maxsize: + self._data.popitem(last=False) + + +def _resolve_stream_handle(current_stream: Optional[cuda.CUstream]) -> Optional[int]: + """Integer stream handle for the launch path (None -> torch current stream).""" + if current_stream is None: + return None + return int(current_stream) + + +@contextmanager +def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch.device) -> Iterator[None]: + """Run torch work on ``current_stream`` (device-tagged) when one is given.""" + if current_stream is None: + yield + return + with torch.cuda.stream(torch.cuda.get_stream_from_external(int(current_stream), device)): + yield + + +def _reject_deterministic_backward() -> None: + """The backward accumulates ``dAPE`` with fp32 atomics and is not deterministic. + + Mirrors torch's deterministic-mode semantics: strict mode raises, warn-only mode + warns and runs. + """ + if torch.are_deterministic_algorithms_enabled(): + message = ( + "CSA compressor backward accumulates dAPE with fp32 atomics and is not " + "deterministic; torch.use_deterministic_algorithms(True) is set. Use a " + "deterministic (eager) implementation instead." + ) + if torch.is_deterministic_algorithms_warn_only_enabled(): + warnings.warn(message, RuntimeWarning, stacklevel=2) + else: + raise RuntimeError(message) + + +class _CSACompressorBase(APIBase): + """Shared descriptor plumbing and ``check_support`` for forward and backward.""" + + def __init__( + self, + sample_kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16 + sample_score: torch.Tensor, # (total_tokens, coff * head_dim) BF16 + sample_ape: torch.Tensor, # (ratio, coff * head_dim) FP32 + sample_cu_seqlens: torch.Tensor, # (B + 1,) INT32 token offsets + sample_cu_seqlens_comp: torch.Tensor, # (B + 1,) INT32 compressed-block offsets + sample_out: torch.Tensor, # (total_comp, head_dim) BF16 (forward output / backward grad_out) + ratio: int = 4, + coff: int = 2, + ): + """Capture tensor descriptors and the ``(ratio, coff)`` configuration. + + The ``sample_*`` tensors provide shape/dtype/stride/device metadata only (meta + tensors are accepted); validation runs later in ``check_support`` and nothing + is read or launched until ``execute``. + """ + super().__init__() + self._warn_experimental_api() + + self.kv_desc = self._make_tensor_desc(sample_kv, name="sample_kv") + self.score_desc = self._make_tensor_desc(sample_score, name="sample_score") + self.ape_desc = self._make_tensor_desc(sample_ape, name="sample_ape") + self.cu_desc = self._make_tensor_desc(sample_cu_seqlens, name="sample_cu_seqlens") + self.cuc_desc = self._make_tensor_desc(sample_cu_seqlens_comp, name="sample_cu_seqlens_comp") + self.out_desc = self._make_tensor_desc(sample_out, name="sample_out") + + self.ratio = int(ratio) + self.coff = int(coff) + + self.total_tokens = None + self.total_comp = None + self.head_dim = None + self.n_seg = None + self.target_device: Optional[torch.device] = None + + def check_support(self) -> bool: + """Validate the configuration against the kernels' validated envelope. + + Malformed inputs and configurations outside the envelope raise ``ValueError`` + (device-capability failures raise ``RuntimeError``), mirroring the other FE-OSS + APIs; there is no soft fallback path inside this API. + """ + self._logger.debug("Entering check_support") + if self.ratio == 4: + self._value_error_if( + self.coff not in (1, 2), + f"CSA compressor at ratio=4 is validated for coff in {{1, 2}} (coff=2 is the production CSA/HCA form), got coff={self.coff}", + ) + elif self.ratio == 128: + self._value_error_if( + self.coff not in (1, 2), + f"CSA compressor at ratio=128 supports coff in {{1, 2}}, got coff={self.coff}", + ) + else: + self._value_error_if( + True, + f"CSA compressor is validated for ratio in {{4, 128}} only, got ratio={self.ratio}, coff={self.coff}", + ) + self._value_error_if( + self.kv_desc.ndim != 2, + f"kv must be 2-D (total_tokens, coff * head_dim), got {self.kv_desc.shape}", + ) + self._value_error_if( + self.out_desc.ndim != 2, + f"out/grad_out must be 2-D (total_comp, head_dim), got {self.out_desc.shape}", + ) + total_tokens, width = self.kv_desc.shape + total_comp, head_dim = self.out_desc.shape + self._value_error_if( + head_dim < 1 or width != self.coff * head_dim, + f"kv width must equal coff * head_dim = {self.coff} * {head_dim}, got {width}", + ) + self._value_error_if( + self.score_desc.shape != self.kv_desc.shape, + f"score shape {self.score_desc.shape} != kv shape {self.kv_desc.shape}", + ) + self._value_error_if( + self.ape_desc.shape != (self.ratio, width), + f"ape must be (ratio, coff * head_dim) = ({self.ratio}, {width}), got {self.ape_desc.shape}", + ) + self._value_error_if( + self.cu_desc.ndim != 1 or self.cuc_desc.ndim != 1, + "cu_seqlens and cu_seqlens_comp must be 1-D", + ) + self._value_error_if( + self.cu_desc.shape != self.cuc_desc.shape or self.cu_desc.shape[0] < 2, + f"cu_seqlens and cu_seqlens_comp must both have B + 1 >= 2 entries, got {self.cu_desc.shape} and {self.cuc_desc.shape}", + ) + + self._check_dtype(self.kv_desc, torch.bfloat16, name="kv") + self._check_dtype(self.score_desc, torch.bfloat16, name="score") + self._check_dtype(self.ape_desc, torch.float32, name="ape") + self._check_dtype(self.cu_desc, torch.int32, name="cu_seqlens") + self._check_dtype(self.cuc_desc, torch.int32, name="cu_seqlens_comp") + self._check_dtype(self.out_desc, torch.bfloat16, name="out/grad_out") + + # int32 flat offsets: the kernels index flat views with int32 arithmetic. + self._value_error_if( + total_tokens * width >= _INT32_LIMIT, + f"total_tokens * coff * head_dim must be < 2**31 for int32 flat offsets, got {total_tokens} * {width}", + ) + self._value_error_if( + total_comp * head_dim >= _INT32_LIMIT, + f"total_comp * head_dim must be < 2**31 for int32 flat offsets, got {total_comp} * {head_dim}", + ) + # APE is indexed as (k % ratio) * width + col with the same int32 arithmetic + # (only reachable at extreme head_dims, but cheap to pin explicitly). + self._value_error_if( + self.ratio * width >= _INT32_LIMIT, + f"ratio * coff * head_dim must be < 2**31 for int32 APE offsets, got {self.ratio} * {width}", + ) + # gridDim.y bound of the forward launch schedule (64/128-thread column groups): + # head_dims beyond this cannot be launched (the pre-vectorization schedule had + # the same 128 * 65535 envelope, just unchecked). + self._value_error_if( + head_dim > _MAX_HEAD_DIM, + f"head_dim must be <= {_MAX_HEAD_DIM} (forward launch gridDim.y bound), got {head_dim}", + ) + # The ratio=128 kernels are gated to the head_dims actually validated on + # hardware (numerics gate + ptxas 0-spill + benchmark, see + # compressor_sm100_r128.py); the kernels are generic and the gate can be + # widened per head_dim once validated. + if self.ratio == 128: + self._value_error_if( + head_dim not in (128, 512), + f"CSA compressor at ratio=128 is validated for head_dim in {{128, 512}} only, got head_dim={head_dim}", + ) + # Rows (including static-capacity padding rows) gather a window of `ratio` + # tokens; the eager gather has the same requirement. + self._value_error_if( + total_comp > 0 and total_tokens < self.ratio, + f"total_comp={total_comp} > 0 requires at least ratio={self.ratio} tokens, got {total_tokens}", + ) + for desc, name in ( + (self.kv_desc, "kv"), + (self.score_desc, "score"), + (self.ape_desc, "ape"), + (self.cu_desc, "cu_seqlens"), + (self.cuc_desc, "cu_seqlens_comp"), + (self.out_desc, "out/grad_out"), + ): + self._value_error_if(not desc.is_contiguous(), f"{name} must be contiguous") + + # Device resolution: all runtime (CUDA) descriptors on one device; meta + # descriptors are metadata-only stand-ins and pin nothing. + all_descs = (self.kv_desc, self.score_desc, self.ape_desc, self.cu_desc, self.cuc_desc, self.out_desc) + devices = {desc.device for desc in all_descs} + self._value_error_if( + any(dev.type not in ("cuda", "meta") for dev in devices), + f"all tensors must be CUDA tensors, got devices {sorted(str(dev) for dev in devices)}", + ) + cuda_devices = {dev for dev in devices if dev.type == "cuda"} + self._value_error_if( + len(cuda_devices) > 1, + f"all tensors must share one CUDA device, got {sorted(str(dev) for dev in cuda_devices)}", + ) + self._runtime_error_if(not torch.cuda.is_available(), "CSA compressor requires CUDA") + if cuda_devices: + target = next(iter(cuda_devices)) + if target.index is None: + target = torch.device("cuda", torch.cuda.current_device()) + else: + target = torch.device("cuda", torch.cuda.current_device()) + + capability = torch.cuda.get_device_capability(target) + self._runtime_error_if( + capability != SUPPORTED_COMPUTE_CAPABILITY, + f"CSA compressor requires compute capability {SUPPORTED_COMPUTE_CAPABILITY} (the only validated architecture so far), found SM{capability[0]}.{capability[1]} on {target}", + ) + + self.total_tokens = total_tokens + self.total_comp = total_comp + self.head_dim = head_dim + self.n_seg = self.cu_desc.shape[0] - 1 + self.target_device = target + self._is_supported = True + return True + + def _validate_runtime_tensor(self, tensor, name, shape, dtype, device, align): + """Cheap per-call validation of one runtime tensor.""" + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.dtype != dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if not tensor.is_cuda or tensor.device != device: + raise ValueError(f"{name} must be a CUDA tensor on {device}, got {tensor.device}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {tuple(tensor.stride())}") + # Contiguity does not imply base-pointer alignment (storage-offset views); the + # kernels' pointer wrappers assume it. + if tensor.data_ptr() % align: + raise ValueError(f"{name} base pointer must be {align}-byte aligned, got 0x{tensor.data_ptr():x}") + + @staticmethod + def _record_streams(tensors, current_stream: Optional[cuda.CUstream], device: torch.device) -> None: + """Keep tensor storages alive for work enqueued on an explicit external stream. + + The launch path takes raw pointers, so PyTorch's caching allocator does not know + the kernel on ``current_stream`` still reads/writes these tensors: without + ``record_stream`` it may recycle a storage freed by the caller while the kernel + is pending. Only needed for explicit streams — with ``current_stream=None`` the + launch lands on torch's current stream and ordinary stream semantics apply. + """ + if current_stream is None: + return + consumer = torch.cuda.get_stream_from_external(int(current_stream), device) + for t in tensors: + t.record_stream(consumer) + + +class CSACompressorForward(_CSACompressorBase): + """Fused CSA compressor forward: one kernel over the whole THD pack. + + Rows in ``[cu_seqlens_comp[-1], total_comp)`` are static-capacity padding (for + CUDA-graph static shapes) and are computed with first-in-segment semantics from + token 0, exactly like the eager gather. + """ + + def compile(self) -> None: + """JIT-compile the forward kernel for this ``(ratio, head_dim, coff, device)`` (idempotent).""" + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + # Route by ratio: the dedicated ratio=128 kernels share the launch machinery + # and cache pattern with the generic kernels (JIT cache keyed per + # (ratio, head_dim, coff, device) plus the compile-time schedule); the + # numerics contracts intentionally DIFFER per family — see the module + # docstring and docs/fe-oss-apis/csa.md. + if self.ratio == 128: + precompile_fwd_r128(self.ratio, self.head_dim, self.coff, self.target_device) + run = run_fwd_r128 + else: + precompile_fwd(self.ratio, self.head_dim, self.coff, self.target_device) + run = run_fwd + + ratio, head_dim, coff = self.ratio, self.head_dim, self.coff + + def tensor_api(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, stream_handle): + """Per-config closure: invoke the routed forward with the bound ``(ratio, head_dim, coff)``.""" + run(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, out.shape[0], ratio, head_dim, coff, stream_handle=stream_handle) + + self._compiled_kernel = tensor_api + self._logger.debug("Kernel compiled successfully") + + def execute( + self, + kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + score: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + ape: torch.Tensor, # (ratio, coff * head_dim) FP32, contiguous + cu_seqlens: torch.Tensor, # (B + 1,) INT32, contiguous + cu_seqlens_comp: torch.Tensor, # (B + 1,) INT32, contiguous + out: torch.Tensor, # (total_comp, head_dim) BF16, contiguous + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Run the compiled forward kernel; ``out`` is fully overwritten.""" + self._logger.debug("Entering execute") + if self._compiled_kernel is None: + raise ValueError("CSACompressorForward kernel not compiled") + device = self.target_device + width = self.coff * self.head_dim + self._validate_runtime_tensor(kv, "kv", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(score, "score", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(ape, "ape", (self.ratio, width), torch.float32, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens, "cu_seqlens", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens_comp, "cu_seqlens_comp", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(out, "out", (self.total_comp, self.head_dim), torch.bfloat16, device, PTR_ALIGN_BYTES) + if out.numel() == 0: + return + self._compiled_kernel(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, _resolve_stream_handle(current_stream)) + self._record_streams((kv, score, ape, cu_seqlens, cu_seqlens_comp, out), current_stream, device) + + +class CSACompressorBackward(_CSACompressorBase): + """Fused CSA compressor backward: recompute window probs, write grads in one kernel. + + ``grad_kv``/``grad_score`` may be UNINITIALIZED: the kernel writes every position — + consumed positions get their gradients (disjoint, atomic-free stores), and every + never-consumed position (segment-tail tokens; for ``coff == 2`` the first-half + columns of each segment's last block; tokens of segments shorter than ``ratio``; + tokens beyond ``cu_seqlens[-1]`` when the buffers carry static token-capacity + padding) gets an exact zero from its unique owning CTA — matching autograd without + separate zero-fill kernels. Exception: when ``total_comp == 0`` the kernel is not + launched and the buffers are left untouched, so a caller that needs autograd-exact + zeros in that case must zero them itself (the high-level wrapper does). + ``grad_ape`` must be zero-initialized by the caller before EVERY ``execute`` call + and before every CUDA-graph replay that reuses the buffer: the kernel only + ACCUMULATES into it (fp32 atomics, not bitwise run-to-run deterministic — + ``grad_kv``/``grad_score`` are) and never clears it, so a reused buffer otherwise + carries the previous invocation's sums. (The high-level wrapper allocates a fresh + zeroed buffer per call; because that zero-fill is captured together with the + kernel, wrapper graph replays re-zero automatically.) + ``execute`` raises under ``torch.use_deterministic_algorithms(True)``. Incoming + gradients on static-capacity padding rows (``[cu_seqlens_comp[-1], total_comp)``) + are ignored. + """ + + def compile(self) -> None: + """JIT-compile the backward kernel for this ``(ratio, head_dim, coff, device)`` (idempotent).""" + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + # Route by ratio, as in the forward. + if self.ratio == 128: + precompile_bwd_r128(self.ratio, self.head_dim, self.coff, self.target_device) + run = run_bwd_r128 + else: + precompile_bwd(self.ratio, self.head_dim, self.coff, self.target_device) + run = run_bwd + + ratio, head_dim, coff = self.ratio, self.head_dim, self.coff + + def tensor_api(kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape, stream_handle): + """Per-config closure: invoke the routed backward with the bound ``(ratio, head_dim, coff)``.""" + run( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + grad_out, + grad_kv, + grad_score, + grad_ape, + grad_out.shape[0], + ratio, + head_dim, + coff, + stream_handle=stream_handle, + ) + + self._compiled_kernel = tensor_api + self._logger.debug("Kernel compiled successfully") + + def execute( + self, + kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + score: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + ape: torch.Tensor, # (ratio, coff * head_dim) FP32, contiguous + cu_seqlens: torch.Tensor, # (B + 1,) INT32, contiguous + cu_seqlens_comp: torch.Tensor, # (B + 1,) INT32, contiguous + grad_out: torch.Tensor, # (total_comp, head_dim) BF16, contiguous + grad_kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16 (may be uninitialized; fully written when total_comp > 0) + grad_score: torch.Tensor, # (total_tokens, coff * head_dim) BF16 (may be uninitialized; fully written when total_comp > 0) + grad_ape: torch.Tensor, # (ratio, coff * head_dim) FP32, zero-initialized before EVERY call/replay (kernel accumulates) + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Run the compiled backward kernel into the gradient buffers (see class docs).""" + self._logger.debug("Entering execute") + if self._compiled_kernel is None: + raise ValueError("CSACompressorBackward kernel not compiled") + _reject_deterministic_backward() + device = self.target_device + width = self.coff * self.head_dim + self._validate_runtime_tensor(kv, "kv", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(score, "score", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(ape, "ape", (self.ratio, width), torch.float32, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens, "cu_seqlens", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens_comp, "cu_seqlens_comp", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(grad_out, "grad_out", (self.total_comp, self.head_dim), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(grad_kv, "grad_kv", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(grad_score, "grad_score", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(grad_ape, "grad_ape", (self.ratio, width), torch.float32, device, PTR_ALIGN_BYTES) + if grad_out.numel() == 0: + return + self._compiled_kernel(kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape, _resolve_stream_handle(current_stream)) + self._record_streams((kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape), current_stream, device) + + +# module-level bounded LRU cache of compiled API instances (thread-safe): +# (kind, ratio, head_dim, coff, total_tokens, total_comp, n_seg, device) -> api +# The compiled kernel underneath is shared per (ratio, head_dim, coff, device) through +# compressor_sm100's compile cache, so shape changes only rebuild the cheap wrapper +# object, never the JIT. +_api_cache = _LruDict() +# serializes verdict construction + JIT so concurrent same-config callers cannot +# compile the same kernel twice +_api_build_lock = threading.Lock() + + +def _get_api(kind, kv, score, ape, cu_seqlens, cu_seqlens_comp, out_shape, ratio, coff): + """Build (or fetch) a compiled forward/backward API instance for these tensors.""" + key = (kind, int(ratio), int(coff), out_shape[1], tuple(kv.shape), out_shape[0], cu_seqlens.shape[0], kv.device.index) + api = _api_cache.get(key) + if api is not None: + return api + with _api_build_lock: + api = _api_cache.get(key) + if api is not None: + return api + sample_out = torch.empty(out_shape, dtype=torch.bfloat16, device="meta") + cls = CSACompressorForward if kind == "fwd" else CSACompressorBackward + api = cls( + sample_kv=kv, + sample_score=score, + sample_ape=ape, + sample_cu_seqlens=cu_seqlens, + sample_cu_seqlens_comp=cu_seqlens_comp, + sample_out=sample_out, + ratio=ratio, + coff=coff, + ) + api.check_support() + api.compile() + _api_cache.put(key, api) + return api + + +def _infer_head_dim(kv: torch.Tensor, head_dim: Optional[int], coff: int) -> int: + """Infer ``head_dim`` from the packed kv width when not given explicitly.""" + if kv.ndim != 2: + raise ValueError(f"kv must be 2-D (total_tokens, coff * head_dim), got {tuple(kv.shape)}") + if head_dim is not None: + return int(head_dim) + width = kv.shape[1] + if coff < 1 or width % coff != 0: + raise ValueError(f"cannot infer head_dim from kv width {width} and coff {coff}") + return width // coff + + +def csa_compressor_forward_wrapper( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + ratio: int = 4, + head_dim: Optional[int] = None, + coff: int = 2, + total_comp: Optional[int] = None, + stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """High-level forward wrapper. Allocates and returns the pooled output. + + Args: + kv: ``(total_tokens, coff * head_dim)`` BF16 gate values (THD packed). + score: ``(total_tokens, coff * head_dim)`` BF16 gate scores. + ape: ``(ratio, coff * head_dim)`` FP32 additive position embedding. + cu_seqlens: ``(B + 1,)`` int32 cumulative token counts per segment. + cu_seqlens_comp: ``(B + 1,)`` int32 cumulative compressed-block counts, + ``cu_seqlens_comp[b + 1] - cu_seqlens_comp[b] == seqlen_b // ratio``. + ratio: compression ratio (tokens per output block); validated envelope: + {4, 128} (the wrappers route to the matching kernel family by ratio). + head_dim: output feature dimension; inferred from ``kv`` width when omitted. + coff: 1 for the own-block window form (window = ``ratio`` tokens, no overlap) or + 2 for the overlapping-window form (window = ``2 * ratio``); validated + envelope: {1, 2} at both ratios (ratio=128 additionally requires head_dim + in {128, 512}). + total_comp: output row count. Defaults to ``cu_seqlens_comp[-1]`` (synchronizes); + pass it explicitly (e.g. a static CUDA-graph capacity, which must be + ``>= cu_seqlens_comp[-1]``) to stay capture-safe. + stream: CUDA stream for allocation and kernel launch (None -> current stream). + + Returns: + ``{'out': (total_comp, head_dim) BF16}`` pooled output (pre-RMSNorm). + """ + head_dim = _infer_head_dim(kv, head_dim, coff) + if total_comp is None: + if cu_seqlens_comp.numel() < 1: + raise ValueError("cu_seqlens_comp must have B + 1 >= 2 entries") + total_comp = int(cu_seqlens_comp[-1].item()) + api = _get_api("fwd", kv, score, ape, cu_seqlens, cu_seqlens_comp, (int(total_comp), head_dim), ratio, coff) + with torch.cuda.device(kv.device), _torch_stream_context(stream, kv.device): + out = torch.empty(int(total_comp), head_dim, dtype=torch.bfloat16, device=kv.device) + with torch.cuda.nvtx.range("csa_compressor_fwd_kernel"): + api.execute(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, current_stream=stream) + return TupleDict(out=out) + + +def csa_compressor_backward_wrapper( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + grad_out: torch.Tensor, + ratio: int = 4, + head_dim: Optional[int] = None, + coff: int = 2, + stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """High-level backward wrapper. Allocates the grad buffers and fills them. + + ``grad_out`` is ``(total_comp, head_dim)`` BF16 (the incoming gradient of the + forward wrapper's ``out``); gradients on static-capacity padding rows are ignored. + ``grad_kv``/``grad_score`` are allocated UNINITIALIZED — the kernel writes every + position, storing exact zeros to never-consumed positions itself, so no zero-fill + kernels run; ``grad_ape`` is allocated zeroed (fp32 atomic accumulation). When + ``total_comp == 0`` the kernel is not launched and all three grads are allocated as + zeros instead, preserving autograd's exact-zero semantics. + Raises ``RuntimeError`` under ``torch.use_deterministic_algorithms(True)`` because + ``grad_ape`` is accumulated with fp32 atomics (``grad_kv``/``grad_score`` are + deterministic and bitwise reproducible). + + Returns: + ``{'grad_kv': (total_tokens, coff * head_dim) BF16, + 'grad_score': (total_tokens, coff * head_dim) BF16, + 'grad_ape': (ratio, coff * head_dim) FP32}`` + """ + head_dim = _infer_head_dim(kv, head_dim, coff) + if grad_out.ndim != 2: + raise ValueError(f"grad_out must be 2-D (total_comp, head_dim), got {tuple(grad_out.shape)}") + api = _get_api("bwd", kv, score, ape, cu_seqlens, cu_seqlens_comp, tuple(grad_out.shape), ratio, coff) + with torch.cuda.device(kv.device), _torch_stream_context(stream, kv.device): + if grad_out.shape[0] == 0: + # No kernel launch below -> the buffers must carry autograd's exact zeros. + grad_kv = torch.zeros_like(kv) + grad_score = torch.zeros_like(score) + else: + grad_kv = torch.empty_like(kv) + grad_score = torch.empty_like(score) + grad_ape = torch.zeros_like(ape, dtype=torch.float32) + with torch.cuda.nvtx.range("csa_compressor_bwd_kernel"): + api.execute(kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape, current_stream=stream) + return TupleDict(grad_kv=grad_kv, grad_score=grad_score, grad_ape=grad_ape) diff --git a/python/cudnn/csa/compressor/compressor_sm100.py b/python/cudnn/csa/compressor/compressor_sm100.py new file mode 100644 index 000000000..a821aec63 --- /dev/null +++ b/python/cudnn/csa/compressor/compressor_sm100.py @@ -0,0 +1,998 @@ +"""Fused CuTe-DSL forward+backward kernels for the CSA/HCA ``Compressor`` gated pooling. + +Ported from Megatron-LM (https://github.com/NVIDIA/Megatron-LM/pull/5984, see also +https://github.com/NVIDIA/Megatron-LM/issues/5968 for measurements and numerics); the +kernel math is unchanged. This module holds the kernels and the launch machinery; the +APIBase wrappers live in ``api.py``. + +The kernels fuse the gated-softmax pooling region of the CSA/HCA ``Compressor`` (THD +packed layout): the chain + + gather-index build -> gather -> ``+ APE`` -> overlap-window transform (``coff == 2``) + -> fp32 softmax over the window -> gated weighted sum -> bf16 cast + +is ONE forward kernel and ONE backward kernel (JIT-compiled per ``(ratio, head_dim, +coff)`` configuration). + +Semantics (ground truth = the eager region in Megatron-LM +``Compressor._forward_thd``): + For each segment ``s`` (``cu_seqlens[s]..cu_seqlens[s+1]``) and each output block ``b`` + of ``ratio`` tokens, the ``2 * ratio`` window (``coff == 2``) is + + - ``k in [0, ratio)``: previous block's token ``tok0 - ratio + k``, first-half + projection column ``j``, APE row ``k`` -> invalid for the segment's first block + (score ``-inf``, kv ``0``, exactly like the eager overlap-window transform). + - ``k in [ratio, 2 * ratio)``: own token ``tok0 + k - ratio``, second-half projection + column ``d + j``, APE row ``k - ratio``. + + ``out[b, j] = sum_k kv_k * softmax_k(score_k + ape_k)`` (fp32, single final bf16 + rounding). ``coff == 1`` (no overlap): the window is the block's own ``ratio`` tokens, + column ``j``, always valid. Per-segment tail tokens (``seqlen % ratio``) are dropped, + as in the eager code. + +Numerics: + All arithmetic is fp32 with a single final bf16 rounding. The fp32 accumulation + structure mirrors the eager ops (serial max, serial sum of exp, ``p = e / denom``, + serial sum ``kv * p``), with ``mul.rn.f32``/``fma.rn.f32`` pinned in PTX so results do + not depend on compiler FMA contraction. Against an fp32-intermediate eager reference, + ``dKV``/``dScore`` are bit-identical and the forward matches to within one bf16 + rounding step (see the tests and the Megatron-LM issue for data). + +Backward: + Atomic-free for ``dKV``/``dScore``: every consumed input element belongs to exactly one + pooling window (for ``coff == 2``, first-half columns are consumed by the NEXT block's + window and second-half columns by the OWN block's window), so gradient stores are + disjoint. Elements never consumed (segment-tail tokens; for ``coff == 2`` the + first-half projection columns of each segment's last block; all tokens of segments + shorter than ``ratio``; tokens beyond ``cu_seqlens[-1]`` when the gradient buffers + carry static token-capacity padding) are written as exact zeros by the kernel + itself — each such slot has a unique natural owner (see the kernel docstring) — so + ``dKV``/``dScore`` buffers need no zero-initialization and no separate fill kernels, + matching autograd output exactly. ``dAPE`` is accumulated in registers over ``rows_per_cta`` blocks and + then reduced with one fp32 atomic per ``(k, dim)`` per CTA into a buffer the caller + must still zero-initialize; ``dAPE`` is therefore not bitwise run-to-run deterministic + (forward, ``dKV`` and ``dScore`` are). When ``total_comp == 0`` no kernel is launched + and the buffers are left untouched (the wrapper falls back to allocating zeros). + +Static-capacity padding (``total_comp > cu_seqlens_comp[-1]``): + Forward computes the padding rows exactly like the eager code: they gather the window + from token 0 with first-in-segment semantics (requires ``total_tokens >= ratio``, like + the eager gather). Backward ignores incoming gradients on padding rows (they are tail + padding for CUDA-graph static shapes and are not consumed downstream). + +CUDA graphs: + The launch path is capture-compatible once the kernels for a given + ``(ratio, head_dim, coff)`` configuration have been JIT-compiled; compile (or run one + eager step) per configuration before capture. A call that would JIT under capture + raises a ``RuntimeError`` instead of corrupting the capture. + +``CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH=0`` disables only the cached-launch optimization +(see ``_FastLauncher``). +""" + +from __future__ import annotations + +import ctypes +import os +import threading + +import torch +import cuda.bindings.driver as cuda_driver + +import cutlass +import cutlass.base_dsl.typing as _cutlass_typing +import cutlass.cute as cute +import cutlass.cute.arch as cute_arch +import cutlass.cute.math as cute_math +from cutlass._mlir.dialects import llvm as _llvm +from cutlass.cute.runtime import make_ptr + +_ENV_FAST_LAUNCH = "CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH" + +# The only compute capability the kernels have been validated on so far. The kernels use +# no architecture-specific features (plain loads/stores, fp32 atomics, pinned mul/fma +# PTX), but wider coverage stays opt-in until validated per architecture. The +# ``cute.compile`` default arch resolution maps (10, 0) to ``sm_100a``, so the CC gate +# also pins the generated SASS target. +SUPPORTED_COMPUTE_CAPABILITY = (10, 0) + + +# ============================================================================= +# Cached fast-path launcher +# ============================================================================= +# Each steady-state CuTe-DSL call spends tens of microseconds of pure host Python +# (rebuilding pointer/scalar/stream argument objects, adapter lookups, fresh ctypes +# allocations, re-packing the void** array) to end at a ~3-4 us C launch call. For +# microsecond-scale kernels that overhead dominates the wall clock, so the launch state +# is snapshotted ONCE per (kernel, config, device, thread) and replayed with in-place +# mutation of the argument storages. +# ============================================================================= + +# torch's raw current-stream query (~0.5 us) vs `torch.cuda.current_stream` object +# construction (~2-3 us). Same handle the slow path ends up passing. Private API: guard +# the bind so module import survives torch builds that do not expose it. +_raw_stream = getattr(torch._C, "_cuda_getCurrentRawStream", None) +if _raw_stream is None: # pragma: no cover - older/stripped torch builds + + def _raw_stream(device_index=None): + """Fallback raw-stream query via the public torch API.""" + return torch.cuda.current_stream(device_index).cuda_stream + + +def _fast_launch_enabled() -> bool: + """Return True unless the cached-launch optimization is disabled via environment.""" + return os.environ.get(_ENV_FAST_LAUNCH, "1") == "1" + + +def _view_for_arg(arg, addr): + """Build a ctypes view over the storage backing one execution-args slot.""" + if isinstance(arg, _cutlass_typing.Numeric): + if isinstance(arg, _cutlass_typing.Boolean): + return ctypes.c_bool.from_address(addr) + if isinstance(arg, _cutlass_typing.Integer): + width = type(arg).width + signed = getattr(type(arg), "signed", True) + ctype = getattr(ctypes, f"c_{'int' if signed else 'uint'}{width}") + return ctype.from_address(addr) + if isinstance(arg, _cutlass_typing.Float32): + return ctypes.c_float.from_address(addr) + if isinstance(arg, _cutlass_typing.Float64): + return ctypes.c_double.from_address(addr) + raise TypeError(f"unsupported numeric scalar {type(arg)!r}") + # A cute runtime Pointer (make_ptr) stores its address in a c_void_p `_desc`; + # CUstream's storage is its own pointer-sized handle. + if hasattr(arg, "_desc") and isinstance(arg._desc, ctypes.c_void_p): + return ctypes.c_void_p.from_address(addr) + if type(arg).__name__ == "CUstream": + return ctypes.c_void_p.from_address(addr) + raise TypeError(f"unsupported argument type {type(arg)!r}") + + +class _FastLauncher: + """Replayable launch state for one compiled CuTe-DSL function. + + ``slots[i]`` is a ctypes view over the storage feeding runtime argument ``i`` (same + order as the tuple passed to ``fn(*args)``); write ``.value`` then call ``launch()``. + + Guards: + - Only flat argument tuples of cute runtime pointers (``make_ptr``), cutlass + scalars, and ``CUstream`` are eligible; anything else raises during construction + and the wrapper stays on its slow path. + - Construction introspects private-but-stable DSL internals + (``_default_executor``, ``_get_invoke_packed_args``); any structural mismatch on + a future ``nvidia-cutlass-dsl`` upgrade raises during construction, and the + wrapper permanently falls back to the regular (slow) launch path rather than + attempting a launch from a partially built snapshot. + """ + + __slots__ = ("slots", "_capi", "_packed", "_res", "_has_res", "_keep") + + def __init__(self, fn, args): + """Snapshot the launch state of one ``fn(*args)`` call for later replay.""" + # Must run after the wrapper's first real `fn(*args)` call so the default + # executor (device context, loaded modules) exists. + exe_args, adapted = fn.execution_args.generate_execution_args(args, {}) + executor = fn._default_executor + if executor is None: + raise RuntimeError("build _FastLauncher after the first fn() call") + if len(exe_args) != len(args): + # struct/dlpack args expand to multiple slots -> unsupported. + raise TypeError(f"non-flat exe_args ({len(exe_args)} slots for {len(args)} args)") + # Private copy of the packed void** array: the executor's own is a shared + # thread-local scratch buffer that any interleaved slow-path call would repoint + # to its (dead) per-call storages. + tls_packed = executor._get_invoke_packed_args(exe_args) + total = len(exe_args) + executor._num_extra_args + packed = (ctypes.c_void_p * total)() + for i in range(total): + packed[i] = tls_packed[i] + views = [] + for arg, exe_arg in zip(args, exe_args): + addr = exe_arg.value if isinstance(exe_arg, ctypes.c_void_p) else int(exe_arg) + views.append(_view_for_arg(arg, addr)) + self.slots = views + self._capi = executor.capi_func + self._has_res = executor._has_cuda_result + if self._has_res: + # Private result storage: the executor's own `cuda_result` is shared by + # every launcher built from this compiled function, so concurrent launches + # from different threads would overwrite one another's CUDA status. The + # result address is the first extra slot after the base arguments (see + # jit_executor._get_invoke_packed_args). + self._res = type(executor.cuda_result)() + packed[len(exe_args)] = ctypes.addressof(self._res) + else: + self._res = None + self._packed = packed + # Keep every object owning a storage referenced by `packed` alive. + self._keep = (args, exe_args, adapted, fn, executor) + + def launch(self): + """Replay the snapshotted launch with the current slot values.""" + self._capi(self._packed) + if self._has_res: + result = self._res.value + if result != 0: + raise RuntimeError(f"CUDA error {result} in CuTe-DSL fast launch (set {_ENV_FAST_LAUNCH}=0 to fall back to the slow launch path)") + + +class _FastCache: + """Thread-local ``{key: _FastLauncher}`` with build-once semantics. + + ``get`` returns a launcher or None (not built / build failed / disabled). ``put`` + attempts to build; a failed build is remembered so the wrapper pays the (cheap) + attempt exactly once per thread and stays on its slow path afterwards. The cache is + thread-local because callers may run forward on the main thread while backward runs + on an autograd thread, and slot mutation is not thread-safe. + """ + + def __init__(self): + """Create the empty per-thread launcher storage.""" + self._tls = threading.local() + + def _map(self): + """Return this thread's key -> launcher map.""" + cache_map = getattr(self._tls, "m", None) + if cache_map is None: + cache_map = {} + self._tls.m = cache_map + return cache_map + + def get(self, key): + """Return the cached launcher for ``key`` or None.""" + launcher = self._map().get(key) + return launcher if launcher is not None and launcher is not False else None + + def put(self, key, fn, args): + """Try to build and cache a launcher for ``key``; never fails the call.""" + if not _fast_launch_enabled(): + return + cache_map = self._map() + if key in cache_map: + return + try: + cache_map[key] = _FastLauncher(fn, args) + except Exception: # pylint: disable=broad-except + # Structural mismatch (DSL upgrade, exotic arg): remember and keep the + # wrapper on its slow path. Never fail the call. + cache_map[key] = False + + +_FAST = _FastCache() + + +# ============================================================================= +# CuTe-DSL kernel definitions +# ============================================================================= + +_NEG_INF = float("-inf") + + +@cutlass.dsl_user_op +def _fmul_rn(a, b, *, loc=None, ip=None): + """fp32 multiply pinned to ``mul.rn.f32`` (opaque to FMA contraction). + + The eager ``(kv * weights).sum(dim=1)`` and the softmax-backward inner sum both + accumulate ROUNDED products serially; letting the compiler contract mul+add into + FMA breaks bit-exactness against the fp32 eager reference. + """ + return cutlass.Float32( + _llvm.inline_asm( + cutlass.Float32.mlir_type, + [ + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + ], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cutlass.dsl_user_op +def _ffma_rn(a, b, c, *, loc=None, ip=None): + """fp32 fma pinned to ``fma.rn.f32``. + + The eager softmax-backward epilogue is ``ds = fma(p, -S, round(p * dp))``; pinning + removes any dependence on compiler contraction choices. + """ + return cutlass.Float32( + _llvm.inline_asm( + cutlass.Float32.mlir_type, + [ + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + cutlass.Float32(c).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.kernel +def _compressor_fwd_kernel( + mKV: cute.Tensor, # flat [T * W] bf16, W = coff * d + mScore: cute.Tensor, # flat [T * W] bf16 + mAPE: cute.Tensor, # flat [ratio * W] fp32 + mCu: cute.Tensor, # [n_seq + 1] int32 (token cu_seqlens) + mCuComp: cute.Tensor, # [n_seq + 1] int32 (block cu_seqlens) + mOut: cute.Tensor, # flat [nb_total * d] bf16 + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """Forward: one thread per (output block, ``vec`` adjacent head dims). + + ``vec == 2`` widens every bf16 access to one 32-bit load/store (``vec == 1`` is the + scalar layout for odd ``head_dim``). The per-thread window slices are contiguous and + ``vec``-aligned by construction (``W``, ``d`` and the thread's first column are all + multiples of ``vec``), which ``cute.assume`` makes provable so ``autovec_copy`` + lowers each slice to a single ``vec * 16``-bit universal copy. Wider vectors were + measured and rejected: 64/128/256-bit variants cut instructions but blow up + registers (80/147/255 per thread) and occupancy, losing to ``vec == 2`` on every + production shape. + + The per-lane fp32 math is IDENTICAL to the scalar kernel (same op order per output + element, one lane per head dim), so the output stays bitwise stable across the + ``vec`` configurations. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + ncol: cutlass.Constexpr = d // vec # thread-column count per output row + col = bidy * threads + tidx + W: cutlass.Constexpr = coff * d + win: cutlass.Constexpr = 2 * ratio if coff == 2 else ratio + + if col < ncol: + cvec = col * vec # first head-dim column of this thread's lane group + + # Hoist APE loads: constant per (k, lane) across all rows. + ape_k = [] + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + colbase = cvec + else: + colbase = (d + cvec) if cutlass.const_expr(coff == 2) else cvec + fr_a = cute.make_rmem_tensor((vec,), cutlass.Float32) + aoff = cute.assume((k % ratio) * W + colbase, divby=vec) + gA = cute.make_tensor(mAPE.iterator + aoff, cute.make_layout(vec)) + cute.autovec_copy(gA, fr_a) + for j in cutlass.range_constexpr(vec): + ape_k.append(cutlass.Float32(fr_a[j])) + + # True compressed row count; rows in [nb_valid, nb_total) are static-capacity + # padding and gather the window from token 0 with first-in-segment semantics, + # like the eager code. + nb_valid = mCuComp[n_seq] + + for rr in cutlass.range_constexpr(rows_per_cta): + bb = bidx * rows_per_cta + rr + if bb < nb_total: + # Per-segment boundary scan (n_seq is small). + seq_idx = cutlass.Int32(0) + bis = cutlass.Int32(0) + if bb < nb_valid: + bis = cutlass.Int32(bb) + for s in cutlass.range(n_seq): + cs = mCuComp[s] + ce = mCuComp[s + 1] + if bb >= cs: + if bb < ce: + seq_idx = s + bis = bb - cs + tok0 = mCu[seq_idx] + bis * ratio + + sv = [] + kvv = [] + for k in cutlass.range_constexpr(win): + fr_s = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_k = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + if cutlass.const_expr(coff == 2 and k < ratio): + if bis > 0: + off = cute.assume((tok0 - ratio + k) * W + cvec, divby=vec) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gK = cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gK, fr_k) + # Same value construction as the scalar kernel: the invalid + # window contributes the CONSTANT -inf score (no APE add — APE + # values are not required to be finite) and a zero kv lane. + for j in cutlass.range_constexpr(vec): + v = cutlass.Float32(_NEG_INF) + u = cutlass.Float32(0.0) + if bis > 0: + v = cutlass.Float32(fr_s[j]) + ape_k[k * vec + j] + u = cutlass.Float32(fr_k[j]) + sv.append(v) + kvv.append(u) + else: + if cutlass.const_expr(coff == 2): + off = cute.assume((tok0 + k - ratio) * W + d + cvec, divby=vec) + else: + off = cute.assume((tok0 + k) * W + cvec, divby=vec) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gK = cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gK, fr_k) + for j in cutlass.range_constexpr(vec): + sv.append(cutlass.Float32(fr_s[j]) + ape_k[k * vec + j]) + kvv.append(cutlass.Float32(fr_k[j])) + + fr_o = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + for j in cutlass.range_constexpr(vec): + mx = sv[j] + for k in cutlass.range_constexpr(1, win): + if sv[k * vec + j] > mx: + mx = sv[k * vec + j] + den = cutlass.Float32(0.0) + ex = [] + for k in cutlass.range_constexpr(win): + e = cute_math.exp(sv[k * vec + j] - mx) + den = den + e + ex.append(e) + acc = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(win): + acc = acc + _fmul_rn(kvv[k * vec + j], ex[k] / den) + fr_o[j] = cutlass.BFloat16(acc) + ooff = cute.assume(bb * d + cvec, divby=vec) + gO = cute.make_tensor(mOut.iterator + ooff, cute.make_layout(vec)) + cute.autovec_copy(fr_o, gO) + + +@cute.kernel +def _compressor_bwd_kernel( + mKV: cute.Tensor, # flat [T * W] bf16 + mScore: cute.Tensor, # flat [T * W] bf16 + mAPE: cute.Tensor, # flat [ratio * W] fp32 + mCu: cute.Tensor, # [n_seq + 1] int32 + mCuComp: cute.Tensor, # [n_seq + 1] int32 + mGO: cute.Tensor, # flat [nb_total * d] bf16 + mGKV: cute.Tensor, # flat [T * W] bf16 (fully written; may be uninitialized) + mGS: cute.Tensor, # flat [T * W] bf16 (fully written; may be uninitialized) + mGAPE: cute.Tensor, # flat [ratio * W] fp32 (zero-initialized) + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + total_tokens: cutlass.Int32, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """Backward: recompute window probs, disjoint ``dKV``/``dScore`` stores, ``dAPE`` atomics. + + ``dKV``/``dScore`` are FULLY written by the kernel: consumed positions get their + gradients, and every never-consumed position gets an exact zero from its unique + natural owner (see below), so the caller can pass uninitialized buffers instead of + paying two tensor-wide zero-fills. The zero-write ownership keeps all stores + disjoint (no atomics, bitwise run-to-run deterministic): + + - for ``coff == 2``, the first-half columns of each segment's LAST block's own + tokens (no next block consumes them) — written by that last block; + - per-segment tail tokens (``seqlen % ratio``, both halves) — written by the + segment's last block; + - all tokens of segments with zero output blocks (``seqlen < ratio``) — written + by the CTA column ``bidx == 0``; + - tokens beyond ``cu_seqlens[-1]`` (static token-capacity padding of the + gradient buffers) — grid-strided across the CTA columns. + + ``dAPE`` is still accumulated into a caller-zero-initialized buffer with one fp32 + atomic per ``(k, dim)`` per CTA. Rows in ``[cu_seqlens_comp[-1], nb_total)`` are + static-capacity padding; their incoming gradients are ignored. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + dim = bidy * threads + tidx + W: cutlass.Constexpr = coff * d + win: cutlass.Constexpr = 2 * ratio if coff == 2 else ratio + ZERO_BF16 = cutlass.BFloat16(0.0) + + if dim < d: + ape_k = [] + dape = [] + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + col = dim + else: + col = (d + dim) if cutlass.const_expr(coff == 2) else dim + ape_k.append(mAPE[(k % ratio) * W + col]) + dape.append(cutlass.Float32(0.0)) + + nb_valid = mCuComp[n_seq] + + # CTA column (0, bidy) zeroes both halves of every token in segments that have + # zero output blocks (seqlen < ratio): those tokens are never consumed by any + # pooling window, so no block-owning CTA would otherwise write them. + if bidx == 0: + for s in cutlass.range(n_seq): + if mCuComp[s + 1] == mCuComp[s]: + t0 = mCu[s] + t1 = mCu[s + 1] + for tt in cutlass.range(t1 - t0): + mGKV[(t0 + tt) * W + dim] = ZERO_BF16 + mGS[(t0 + tt) * W + dim] = ZERO_BF16 + if cutlass.const_expr(coff == 2): + mGKV[(t0 + tt) * W + d + dim] = ZERO_BF16 + mGS[(t0 + tt) * W + d + dim] = ZERO_BF16 + + # Tokens in [cu_seqlens[-1], total_tokens) are static token-capacity padding of + # the gradient buffers (CUDA-graph static shapes): no segment owns them, so the + # CTA columns zero them in a grid-strided sweep. count == 0 in the common + # exact-size case. The quotient/remainder split keeps every intermediate within + # int32 for any count < 2**31. + gdimx, _, _ = cute.arch.grid_dim() + pad0 = mCu[n_seq] + pad_count = total_tokens - pad0 + if bidx < pad_count: + my_count = pad_count // gdimx + if bidx < pad_count % gdimx: + my_count = my_count + 1 + for i in cutlass.range(my_count): + t = pad0 + bidx + i * gdimx + mGKV[t * W + dim] = ZERO_BF16 + mGS[t * W + dim] = ZERO_BF16 + if cutlass.const_expr(coff == 2): + mGKV[t * W + d + dim] = ZERO_BF16 + mGS[t * W + d + dim] = ZERO_BF16 + + for rr in cutlass.range_constexpr(rows_per_cta): + bb = bidx * rows_per_cta + rr + if bb < nb_valid: + seq_idx = cutlass.Int32(0) + bis = cutlass.Int32(bb) + for s in cutlass.range(n_seq): + cs = mCuComp[s] + ce = mCuComp[s + 1] + if bb >= cs: + if bb < ce: + seq_idx = s + bis = bb - cs + tok0 = mCu[seq_idx] + bis * ratio + + # Recompute window probs (same order as forward). + sv = [] + kvv = [] + offs = [] + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + off = (tok0 - ratio + k) * W + dim + v = cutlass.Float32(_NEG_INF) + u = cutlass.Float32(0.0) + if bis > 0: + v = cutlass.Float32(mScore[off]) + ape_k[k] + u = cutlass.Float32(mKV[off]) + else: + if cutlass.const_expr(coff == 2): + off = (tok0 + k - ratio) * W + d + dim + else: + off = (tok0 + k) * W + dim + v = cutlass.Float32(mScore[off]) + ape_k[k] + u = cutlass.Float32(mKV[off]) + sv.append(v) + kvv.append(u) + offs.append(off) + + mx = sv[0] + for k in cutlass.range_constexpr(1, win): + if sv[k] > mx: + mx = sv[k] + den = cutlass.Float32(0.0) + ex = [] + for k in cutlass.range_constexpr(win): + e = cute_math.exp(sv[k] - mx) + den = den + e + ex.append(e) + + go = cutlass.Float32(mGO[bb * d + dim]) + + # Same expression tree as torch's softmax_backward_data: + # dp_k = go * kv_k ; S = serial sum of ROUNDED dp_k * p_k ; + # ds_k = fma(p_k, -S, round(dp_k * p_k)) ; dkv_k = go * p_k. + p = [] + dp = [] + S = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(win): + pk = ex[k] / den + dpk = go * kvv[k] + S = S + _fmul_rn(dpk, pk) + p.append(pk) + dp.append(dpk) + + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + if bis > 0: + ds = _ffma_rn(p[k], -S, _fmul_rn(dp[k], p[k])) + mGKV[offs[k]] = cutlass.BFloat16(go * p[k]) + mGS[offs[k]] = cutlass.BFloat16(ds) + dape[k] = dape[k] + ds + else: + ds = _ffma_rn(p[k], -S, _fmul_rn(dp[k], p[k])) + mGKV[offs[k]] = cutlass.BFloat16(go * p[k]) + mGS[offs[k]] = cutlass.BFloat16(ds) + dape[k] = dape[k] + ds + + # The segment's last block additionally zeroes the never-consumed slots + # it is the unique natural owner of: (a) for coff == 2 the first-half + # columns of its own tokens (there is no next block to consume them), + # (b) the segment's tail tokens (seqlen % ratio, both halves). + is_last = bb + 1 == mCuComp[seq_idx + 1] + if is_last: + if cutlass.const_expr(coff == 2): + for k in cutlass.range_constexpr(ratio): + mGKV[(tok0 + k) * W + dim] = ZERO_BF16 + mGS[(tok0 + k) * W + dim] = ZERO_BF16 + tail0 = tok0 + ratio + tail1 = mCu[seq_idx + 1] + for tt in cutlass.range(tail1 - tail0): + mGKV[(tail0 + tt) * W + dim] = ZERO_BF16 + mGS[(tail0 + tt) * W + dim] = ZERO_BF16 + if cutlass.const_expr(coff == 2): + mGKV[(tail0 + tt) * W + d + dim] = ZERO_BF16 + mGS[(tail0 + tt) * W + d + dim] = ZERO_BF16 + + # One fp32 atomic per (k, dim) per CTA (amortized over rows_per_cta rows). + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + col = dim + else: + col = (d + dim) if cutlass.const_expr(coff == 2) else dim + cute_arch.atomic_add(mGAPE.iterator + ((k % ratio) * W + col), dape[k]) + + +_EXT = (1 << 31) - 1 # flat extent placeholder (int32 offsets, no bounds checks) + + +@cute.jit +def _compressor_fwd_launch( + kv_ptr: cute.Pointer, + score_ptr: cute.Pointer, + ape_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + cuc_ptr: cute.Pointer, + out_ptr: cute.Pointer, + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + stream: cuda_driver.CUstream, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """JIT entry point that wraps raw pointers into tensors and launches forward.""" + lay = cute.make_layout(_EXT) + mKV = cute.make_tensor(kv_ptr, lay) + mScore = cute.make_tensor(score_ptr, lay) + mAPE = cute.make_tensor(ape_ptr, lay) + mCu = cute.make_tensor(cu_ptr, lay) + mCuComp = cute.make_tensor(cuc_ptr, lay) + mOut = cute.make_tensor(out_ptr, lay) + ncol = d // vec + gx = (nb_total + rows_per_cta - 1) // rows_per_cta + gy = (ncol + threads - 1) // threads + _compressor_fwd_kernel( + mKV, + mScore, + mAPE, + mCu, + mCuComp, + mOut, + nb_total, + n_seq, + ratio, + d, + coff, + vec, + rows_per_cta, + threads, + ).launch(grid=(gx, gy, 1), block=(threads, 1, 1), stream=stream) + + +@cute.jit +def _compressor_bwd_launch( + kv_ptr: cute.Pointer, + score_ptr: cute.Pointer, + ape_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + cuc_ptr: cute.Pointer, + go_ptr: cute.Pointer, + gkv_ptr: cute.Pointer, + gs_ptr: cute.Pointer, + gape_ptr: cute.Pointer, + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + total_tokens: cutlass.Int32, + stream: cuda_driver.CUstream, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """JIT entry point that wraps raw pointers into tensors and launches backward.""" + lay = cute.make_layout(_EXT) + mKV = cute.make_tensor(kv_ptr, lay) + mScore = cute.make_tensor(score_ptr, lay) + mAPE = cute.make_tensor(ape_ptr, lay) + mCu = cute.make_tensor(cu_ptr, lay) + mCuComp = cute.make_tensor(cuc_ptr, lay) + mGO = cute.make_tensor(go_ptr, lay) + mGKV = cute.make_tensor(gkv_ptr, lay) + mGS = cute.make_tensor(gs_ptr, lay) + mGAPE = cute.make_tensor(gape_ptr, lay) + gx = (nb_total + rows_per_cta - 1) // rows_per_cta + gy = (d + threads - 1) // threads + _compressor_bwd_kernel( + mKV, + mScore, + mAPE, + mCu, + mCuComp, + mGO, + mGKV, + mGS, + mGAPE, + nb_total, + n_seq, + total_tokens, + ratio, + d, + coff, + rows_per_cta, + threads, + ).launch(grid=(gx, gy, 1), block=(threads, 1, 1), stream=stream) + + +_COMPILED = {} +# Serializes JIT compilation so concurrent same-config callers cannot compile the same +# kernel twice (the compiled-function cache itself is a plain dict guarded by the GIL). +_COMPILE_LOCK = threading.Lock() +_BWD_ROWS, _BWD_THREADS = 8, 128 + + +def _fwd_schedule(d): + """Forward launch schedule ``(vec, rows_per_cta, threads)`` for ``head_dim == d``. + + ``vec = 2`` (32-bit paired bf16 accesses) whenever ``d`` is even, else the scalar + ``vec = 1`` layout. One output row per CTA with 64-thread column groups: measured + optimum across the production shapes (1x/3x 8192-token packs, head_dim 128/512) — + smaller CTAs raise the sub-wave grid width that limits the small shapes, and wider + per-thread vectors trade instructions for registers/occupancy at a loss (see the + kernel docstring). For enormous head_dims whose column count would overflow the + 65535 ``gridDim.y`` limit at 64 threads, fall back to 128-thread CTAs (the previous + schedule's capability envelope). + """ + vec = 2 if d % 2 == 0 else 1 + ncol = d // vec + threads = 64 if ncol >= 64 else ncol + if (ncol + threads - 1) // threads > 65535: + threads = 128 + return vec, 1, threads + + +# make_ptr assumed alignments below. Contiguity does NOT imply base-pointer alignment +# (storage-offset views), so the API layer checks every runtime tensor's data_ptr() +# against these before launching. +PTR_ALIGN_BYTES = 16 # bf16 / fp32 operands +CU_ALIGN_BYTES = 4 # int32 cu_seqlens operands + + +def _bf16_ptr(t): + """Wrap a bf16 tensor's data pointer for the DSL.""" + return make_ptr(cutlass.BFloat16, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def _f32_ptr(t): + """Wrap an fp32 tensor's data pointer for the DSL.""" + return make_ptr(cutlass.Float32, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def _i32_ptr(t): + """Wrap an int32 tensor's data pointer for the DSL.""" + return make_ptr(cutlass.Int32, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=4) + + +def _compile_fwd(key, args, ratio, d, coff): + """JIT-compile the forward launch entry for ``key`` (capture-guarded).""" + with _COMPILE_LOCK: + fn = _COMPILED.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"CSA compressor: first call for config {key} happened under CUDA " + "graph capture (JIT compilation is not capture-safe); compile() or " + "run one eager step for this configuration before capturing." + ) + fn = cute.compile(_compressor_fwd_launch, *args, ratio, d, coff, *_fwd_schedule(d)) + _COMPILED[key] = fn + return fn + + +def _compile_bwd(key, args, ratio, d, coff): + """JIT-compile the backward launch entry for ``key`` (capture-guarded).""" + with _COMPILE_LOCK: + fn = _COMPILED.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"CSA compressor: first call for config {key} happened under CUDA " + "graph capture (JIT compilation is not capture-safe); compile() or " + "run one eager step for this configuration before capturing." + ) + fn = cute.compile(_compressor_bwd_launch, *args, ratio, d, coff, _BWD_ROWS, _BWD_THREADS) + _COMPILED[key] = fn + return fn + + +def precompile_fwd(ratio, d, coff, device): + """Ensure the forward kernel for ``(ratio, d, coff, device)`` is JIT-compiled. + + Compilation only traces types (pointers/scalars/stream are runtime arguments), so + tiny scratch buffers stand in for the real tensors; nothing is launched. + """ + key = ("fwd", ratio, d, coff, device.index) + if key in _COMPILED: + return + with torch.cuda.device(device): + scratch_bf16 = torch.zeros(16, device=device, dtype=torch.bfloat16) + scratch_f32 = torch.zeros(16, device=device, dtype=torch.float32) + scratch_i32 = torch.zeros(16, device=device, dtype=torch.int32) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + args = ( + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + _i32_ptr(scratch_i32), + _i32_ptr(scratch_i32), + _bf16_ptr(scratch_bf16), + cutlass.Int32(0), + cutlass.Int32(1), + stream, + ) + _compile_fwd(key, args, ratio, d, coff) + + +def precompile_bwd(ratio, d, coff, device): + """Ensure the backward kernel for ``(ratio, d, coff, device)`` is JIT-compiled.""" + key = ("bwd", ratio, d, coff, device.index) + if key in _COMPILED: + return + with torch.cuda.device(device): + scratch_bf16 = torch.zeros(16, device=device, dtype=torch.bfloat16) + scratch_f32 = torch.zeros(16, device=device, dtype=torch.float32) + scratch_i32 = torch.zeros(16, device=device, dtype=torch.int32) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + args = ( + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + _i32_ptr(scratch_i32), + _i32_ptr(scratch_i32), + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + stream, + ) + _compile_bwd(key, args, ratio, d, coff) + + +def run_fwd(kv, score, ape, cu_i, cuc_i, out, nb_total, ratio, d, coff, stream_handle=None): + """Launch the forward kernel (cached fast path -> compiled slow path -> JIT). + + ``stream_handle`` is an integer CUDA stream handle; None uses torch's current + stream on ``kv``'s device. The launch is anchored in ``kv``'s device context: the + compiled module and the default-stream query are per-device, and launching from a + foreign current device silently misbehaves. + """ + dev = kv.device.index + key = ("fwd", ratio, d, coff, dev) + if stream_handle is None: + stream_handle = _raw_stream(dev) + with torch.cuda.device(dev): + launcher = _FAST.get(key) + if launcher is not None: + # Cached launch: mutate the snapshotted argument storages in place; this is + # the same launch the slow path below performs. + slots = launcher.slots + slots[0].value = kv.data_ptr() + slots[1].value = score.data_ptr() + slots[2].value = ape.data_ptr() + slots[3].value = cu_i.data_ptr() + slots[4].value = cuc_i.data_ptr() + slots[5].value = out.data_ptr() + slots[6].value = nb_total + slots[7].value = cu_i.numel() - 1 + slots[8].value = stream_handle + launcher.launch() + return + stream = cuda_driver.CUstream(stream_handle) + args = ( + _bf16_ptr(kv), + _bf16_ptr(score), + _f32_ptr(ape), + _i32_ptr(cu_i), + _i32_ptr(cuc_i), + _bf16_ptr(out), + cutlass.Int32(nb_total), + cutlass.Int32(cu_i.numel() - 1), + stream, + ) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile_fwd(key, args, ratio, d, coff) + fn(*args) + _FAST.put(key, fn, args) + + +def run_bwd(kv, score, ape, cu_i, cuc_i, go, gkv, gs, gape, nb_total, ratio, d, coff, stream_handle=None): + """Launch the backward kernel (cached fast path -> compiled slow path -> JIT). + + Device-context anchoring as in :func:`run_fwd`. The gradient-buffer token capacity + (for the kernel's padding-token zero sweep) is derived from ``kv``'s element count. + """ + dev = kv.device.index + key = ("bwd", ratio, d, coff, dev) + total_tokens = kv.numel() // (coff * d) + if stream_handle is None: + stream_handle = _raw_stream(dev) + with torch.cuda.device(dev): + launcher = _FAST.get(key) + if launcher is not None: + slots = launcher.slots + slots[0].value = kv.data_ptr() + slots[1].value = score.data_ptr() + slots[2].value = ape.data_ptr() + slots[3].value = cu_i.data_ptr() + slots[4].value = cuc_i.data_ptr() + slots[5].value = go.data_ptr() + slots[6].value = gkv.data_ptr() + slots[7].value = gs.data_ptr() + slots[8].value = gape.data_ptr() + slots[9].value = nb_total + slots[10].value = cu_i.numel() - 1 + slots[11].value = total_tokens + slots[12].value = stream_handle + launcher.launch() + return + stream = cuda_driver.CUstream(stream_handle) + args = ( + _bf16_ptr(kv), + _bf16_ptr(score), + _f32_ptr(ape), + _i32_ptr(cu_i), + _i32_ptr(cuc_i), + _bf16_ptr(go), + _bf16_ptr(gkv), + _bf16_ptr(gs), + _f32_ptr(gape), + cutlass.Int32(nb_total), + cutlass.Int32(cu_i.numel() - 1), + cutlass.Int32(total_tokens), + stream, + ) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile_bwd(key, args, ratio, d, coff) + fn(*args) + _FAST.put(key, fn, args) diff --git a/python/cudnn/csa/compressor/compressor_sm100_r128.py b/python/cudnn/csa/compressor/compressor_sm100_r128.py new file mode 100644 index 000000000..2db31300a --- /dev/null +++ b/python/cudnn/csa/compressor/compressor_sm100_r128.py @@ -0,0 +1,1237 @@ +"""Dedicated ratio=128 CSA/HCA ``Compressor`` forward + backward kernels. + +The generic ``compressor_sm100.py`` kernels keep the ENTIRE pooling window in per-thread +registers (a thread owns one ``(output block, head-dim vec-group)`` and materializes all +``win = coff * ratio`` window positions), which is optimal at the production +``ratio = 4`` (48 registers, 0 spill) but hits the 255-register cap from ``ratio = 32`` +and spills kilobytes of local memory per thread at ``ratio = 128`` for both ``coff`` +(measured: 2.5 KB st + 2.5 KB ld at ``coff = 1``, ~33 KB at ``coff = 2``). + +The forward kernel removes the window residency entirely with a **chunked softmax** +whose per-chunk accumulation form is selected per ``nb_total`` schedule bucket +(online-rescale by default, two-phase where it measured faster; see +``_fwd_schedule_r128``): + + - The softmax reduction axis (window position ``k``) is independent per head-dim + column, so each column only needs ``(m, den, acc)`` (max, exp-sum, weighted sum) — + 3 fp32 registers per dim regardless of the window length. The default buckets + accumulate the triple **online** (classic running max with a predicated rescale; + each operand is loaded once). The **two-phase** buckets first compute the exact + chunk max (order-independent fp32 max, no exp), then accumulate ``den``/``acc`` + against that fixed max in ascending ``k`` — the eager two-pass softmax form per + chunk, re-reading score/ape once from the L1/L2-resident chunk lines (kv is read + once in both forms). + - Lanes stay along head-dim columns (``vec`` adjacent dims per lane), so every + window-position step is one fully coalesced 128-byte line per warp per operand — + the layout is ``[token, coff * d]`` row-major, and splitting the window across + lanes instead would make each lane read the same column of 32 different tokens + (32 uncoalesced sectors per step). + - The window is split across ``tchunks`` (``threadIdx.y``) chunk-rows of the CTA; + each chunk-thread streams ``win / tchunks`` positions serially, and the + ``tchunks`` partial triples per ``(block, dim)`` are merged once per output row + through a small smem buffer (``3 * tchunks * threads_x * vec`` fp32) in a fixed + serial order. ``tchunks = 1`` degenerates to a pure streaming kernel: no smem, + no barrier. + +Work decomposition: one CTA per output block row (``gridDim.x = nb_total``), +``gridDim.y`` spans head-dim column groups of ``threads_x * vec`` dims, CTA shape +``(threads_x, tchunks)``. Everything else (THD segment scan, static-capacity padding +rows with first-in-segment semantics, the ``coff == 2`` overlap window with the invalid +previous-block half for each segment's first block, fp32-only window math with a single +final bf16 rounding, pinned ``mul.rn``/``fma.rn``) matches ``compressor_sm100.py``. + +Numerics: identical fp32 dataflow (bf16 kv/score loads widened to fp32, fp32 APE add, +fp32 exp / sums, one final bf16 rounding), but the forward reduction ORDER differs from +the ratio=4 kernel: per-chunk accumulation (online-rescale or exact-chunk-max two-phase, +schedule-selected) + fixed chunk merge instead of a two-pass serial pass over the whole +window. The result is a few fp32 ulps from the whole-window two-pass value, run-to-run +bitwise deterministic (fixed chunk boundaries and merge order, no atomics), and within +the r128 tolerance contract against the fp32-intermediate eager reference — same +values within final-bf16 rounding at the gate tolerances (absolute thresholds +calibrated on the gate's documented input distribution), the reference's NaN/Inf +propagation where its fp32 intermediates overflow, and fp64-oracle parity where the +fp32 intermediates stay finite (gate-checked; reproduce via +``benchmark/csa/gate_csa_compressor_r128.py``). ``exp`` is +``cute.math.exp`` exactly as in the ratio=4 kernel, EXCEPT in schedule buckets that +adopted the ``fastexp`` field (see ``_exp_fast``): those use the tolerance-contract +ex2.approx path, gated per bucket. + +The backward kernel (``_compressor_bwd_r128_kernel``) stages each row's window into +shared memory chunk-parallel, accumulates per-chunk partial ``den`` / ``S`` sums +inside the e-pass (the approved ratio=128 deterministic tolerance contract), merges +them per column in a FIXED chunk order, and stores gradients with a hoisted ``1/den`` +multiply — no serial sweeps, no per-element division. dKV/dScore are deterministic +(fixed orders, no atomics) and match the fp32-intermediate eager autograd within the +r128 tolerance contract (same values within the gate tolerances — absolute thresholds +calibrated on the gate's documented input distribution — the eager reference's NaN/Inf +propagation where its fp32 intermediates overflow, and fp64-oracle parity where the +fp32 intermediates stay finite), NOT bitwise. ``dAPE`` keeps the +ratio=4 contract (one fp32 atomic per ``(k, dim)`` per CTA into a caller-zeroed +buffer, amortized over ``rows_per_cta`` rows; not run-to-run deterministic), and the +kernel-side zero-writes to never-consumed ``dKV``/``dScore`` slots keep the ratio=4 +ownership rules verbatim, parallelized across the CTA (the zero classes are up to 127 +tokens at ratio=128). + +This file intentionally does not touch the ratio=4 path; launch machinery mirrors +``compressor_sm100.py`` and reuses its cached fast-launcher infrastructure. +""" + +from __future__ import annotations + +import threading + +import torch +import cuda.bindings.driver as cuda_driver + +import cutlass +import cutlass.cute as cute +import cutlass.cute.arch as cute_arch +import cutlass.cute.math as cute_math + +from .compressor_sm100 import ( + _EXT, + _FastCache, + _NEG_INF, + _bf16_ptr, + _f32_ptr, + _ffma_rn, + _fmul_rn, + _i32_ptr, + _raw_stream, +) + +# One (block row, dim vec-group) chunk-softmax partial: chunk max / exp-sum / weighted +# sum. Merged across the CTA's tchunks chunk-rows through smem. + +# Fast exp (the ``fastexp`` schedule field; tolerance-contract buckets only): +# ``exp(x) = 2^(x * log2e)`` through ``ex2.approx.ftz.f32`` (MUFU.EX2). The log2e +# multiply is split hi + lo and recombined with one fma, which removes the +# REPRESENTATION error of rounding log2(e) to a single fp32 constant; the rounding of +# the hi product itself and MUFU.EX2's approximation error remain. This is an +# empirically gated approximation, NOT a bounded-error exp: the API places no range +# restriction on ``score + APE``, so no a-priori accuracy bound is claimed — every +# bucket that enables the field must pass the r128 contract gate (tolerance vs fp32 +# eager + fp64-oracle parity, ``benchmark/csa/gate_csa_compressor_r128.py``) on top of +# a measured win. Three issued instructions (FMUL + FFMA + MUFU.EX2) replace the +# ~8-instruction full-range expf sequence — measured ~25-30% of the window loop's +# issued instructions at the instruction-bound buckets (B200 SASS audit). +# exp(-inf) == +0 and NaN -> NaN are preserved by ex2.approx; ftz flushes only +# sub-2^-126 results, which softmax consumes as exact zeros (den >= exp(0) = 1 always +# survives). +_L2E_HI = 1.4426950216293335 # fp32(log2(e)) +_L2E_LO = 1.925963033500011e-08 # fp32(log2(e) - _L2E_HI); residual ~4e-16 + + +def _exp_fast(x): + y = _ffma_rn(x, cutlass.Float32(_L2E_LO), _fmul_rn(x, cutlass.Float32(_L2E_HI))) + return cute_math.exp2(y, approx=True, ftz=True) + + +@cute.kernel +def _compressor_fwd_r128_kernel( + mKV: cute.Tensor, # flat [T * W] bf16, W = coff * d + mScore: cute.Tensor, # flat [T * W] bf16 + mAPE: cute.Tensor, # flat [ratio * W] fp32 + mCu: cute.Tensor, # [n_seq + 1] int32 (token cu_seqlens) + mCuComp: cute.Tensor, # [n_seq + 1] int32 (block cu_seqlens) + mOut: cute.Tensor, # flat [nb_total * d] bf16 + n_seq: cutlass.Int32, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + tchunks: cutlass.Constexpr, + threads_x: cutlass.Constexpr, + twophase: cutlass.Constexpr, + fastexp: cutlass.Constexpr, +): + """Forward: one CTA per output row; window chunked over threadIdx.y, chunk softmax. + + Thread ``(tidx, tidy)`` owns head dims ``[col * vec, col * vec + vec)`` (with + ``col = bidy * threads_x + tidx``) and window positions + ``[tidy * C, (tidy + 1) * C)`` (``C = win / tchunks``). Loads are vec-wide and + coalesced along dims. For ``coff == 2`` a chunk never straddles the half-window + boundary (``ratio % C == 0`` is enforced by the schedule), so the previous-block + half (invalid for each segment's first block, ``bis == 0``) is skipped as a whole + chunk: its partial stays empty (``den == 0``) and the merge ignores it. + """ + tidx, tidy, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + ncol: cutlass.Constexpr = d // vec + win: cutlass.Constexpr = 2 * ratio if coff == 2 else ratio + C: cutlass.Constexpr = win // tchunks + W: cutlass.Constexpr = coff * d + col = bidy * threads_x + tidx + bb = bidx # one output row per CTA + + smem = cutlass.utils.SmemAllocator() + # Partial-merge buffer, [tchunks][threads_x][vec] per quantity; unused (0 B) when + # tchunks == 1 (the allocation below is skipped at trace time). + if cutlass.const_expr(tchunks > 1): + npart: cutlass.Constexpr = tchunks * threads_x * vec + sM = smem.allocate_tensor(cutlass.Float32, cute.make_layout(npart), 16) + sD = smem.allocate_tensor(cutlass.Float32, cute.make_layout(npart), 16) + sA = smem.allocate_tensor(cutlass.Float32, cute.make_layout(npart), 16) + + if col < ncol: + cvec = col * vec + + # THD segment scan (identical to compressor_sm100): rows beyond the true + # compressed count are static-capacity padding and gather the window from + # token 0 with first-in-segment semantics, like the eager code. + nb_valid = mCuComp[n_seq] + seq_idx = cutlass.Int32(0) + bis = cutlass.Int32(0) + if bb < nb_valid: + bis = cutlass.Int32(bb) + for sg in cutlass.range(n_seq): + cs = mCuComp[sg] + ce = mCuComp[sg + 1] + if bb >= cs: + if bb < ce: + seq_idx = sg + bis = bb - cs + tok0 = mCu[seq_idx] + bis * ratio + + # This thread's window chunk [k0, k0 + C). Both coff forms share one loop + # body: token row tok0 - ratio + k (== tok0 + k - ratio for the own half), + # only the projection column / APE row / validity differ per half. + k0 = tidy * C + run_chunk = cutlass.Boolean(True) + tok_row0 = tok0 + k0 + colbase = cvec + ape_row0 = cutlass.Int32(k0) + if cutlass.const_expr(coff == 2): + tok_row0 = tok0 - ratio + k0 + if k0 < ratio: + # Previous block's half-window: first-half projection columns, APE + # row k; invalid (contributes nothing) for the segment's first block + # and for static-capacity padding rows (bis == 0 in both cases). + if bis == 0: + run_chunk = cutlass.Boolean(False) + else: + colbase = d + cvec + ape_row0 = cutlass.Int32(k0 - ratio) + + # Chunk softmax over [k0, k0 + C), in one of two compile-time forms: + # + # - twophase (large d=128 contexts): phase A computes the exact chunk max + # (order-independent fp32 max, no exp), phase B accumulates den/acc against + # that fixed max in ascending k. Versus the online-rescale form this deletes + # the predicated rescale path (a second exp sequence + two multiplies that + # issued EVERY iteration for ~5-in-128 taken updates — ~40% of the loop's + # instructions, B200 SASS audit) at the cost of + # re-reading score/ape once per chunk. Measured win only where the kernel + # is issue-bound rather than latency/traffic-bound (131k-token d=128: + # c1 23.1 -> 18.5 us, c2 64.3 -> 49.6 us); the extra reload LOSES at small + # grids and at d=512 (DRAM ~40% there), so it is a schedule field, not a + # global change. Per chunk the arithmetic is exactly the eager two-pass + # softmax form (exp(s - max) with the true max). + # + # - online (all other buckets): classic running (m, den, acc) with a + # predicated rescale on running-max updates; loads each operand once. + # + # Both forms produce a (m, den, acc) partial with identical merge semantics; + # both are run-to-run bitwise deterministic. + fr_m = cute.make_rmem_tensor((vec,), cutlass.Float32) + fr_d = cute.make_rmem_tensor((vec,), cutlass.Float32) + fr_a = cute.make_rmem_tensor((vec,), cutlass.Float32) + for j in cutlass.range_constexpr(vec): + fr_m[j] = cutlass.Float32(_NEG_INF) + fr_d[j] = cutlass.Float32(0.0) + fr_a[j] = cutlass.Float32(0.0) + + if run_chunk: + if cutlass.const_expr(twophase): + for kk in cutlass.range(C, unroll=4): + off = cute.assume((tok_row0 + kk) * W + colbase, divby=vec) + aoff = cute.assume((ape_row0 + kk) * W + colbase, divby=vec) + fr_s = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_p = cute.make_rmem_tensor((vec,), cutlass.Float32) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gA = cute.make_tensor(mAPE.iterator + aoff, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gA, fr_p) + for j in cutlass.range_constexpr(vec): + s = cutlass.Float32(fr_s[j]) + cutlass.Float32(fr_p[j]) + if s > cutlass.Float32(fr_m[j]): + fr_m[j] = s + for kk in cutlass.range(C, unroll=4): + off = cute.assume((tok_row0 + kk) * W + colbase, divby=vec) + aoff = cute.assume((ape_row0 + kk) * W + colbase, divby=vec) + fr_s = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_k = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_p = cute.make_rmem_tensor((vec,), cutlass.Float32) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gK = cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)) + gA = cute.make_tensor(mAPE.iterator + aoff, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gK, fr_k) + cute.autovec_copy(gA, fr_p) + for j in cutlass.range_constexpr(vec): + s = cutlass.Float32(fr_s[j]) + cutlass.Float32(fr_p[j]) + u = cutlass.Float32(fr_k[j]) + if cutlass.const_expr(fastexp): + e = _exp_fast(s - cutlass.Float32(fr_m[j])) + else: + e = cute_math.exp(s - cutlass.Float32(fr_m[j])) + fr_d[j] = cutlass.Float32(fr_d[j]) + e + fr_a[j] = _ffma_rn(u, e, fr_a[j]) + else: + for kk in cutlass.range(C, unroll=4): + off = cute.assume((tok_row0 + kk) * W + colbase, divby=vec) + aoff = cute.assume((ape_row0 + kk) * W + colbase, divby=vec) + fr_s = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_k = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_p = cute.make_rmem_tensor((vec,), cutlass.Float32) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gK = cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)) + gA = cute.make_tensor(mAPE.iterator + aoff, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gK, fr_k) + cute.autovec_copy(gA, fr_p) + for j in cutlass.range_constexpr(vec): + s = cutlass.Float32(fr_s[j]) + cutlass.Float32(fr_p[j]) + u = cutlass.Float32(fr_k[j]) + m_old = cutlass.Float32(fr_m[j]) + if s > m_old: + # New running max: rescale den/acc by exp(m_old - s); the + # first valid position rescales by exp(-inf) == 0 exactly. + if cutlass.const_expr(fastexp): + scale = _exp_fast(m_old - s) + else: + scale = cute_math.exp(m_old - s) + fr_d[j] = _fmul_rn(fr_d[j], scale) + fr_a[j] = _fmul_rn(fr_a[j], scale) + fr_m[j] = s + if cutlass.const_expr(fastexp): + e = _exp_fast(s - cutlass.Float32(fr_m[j])) + else: + e = cute_math.exp(s - cutlass.Float32(fr_m[j])) + fr_d[j] = cutlass.Float32(fr_d[j]) + e + fr_a[j] = _ffma_rn(u, e, fr_a[j]) + + if cutlass.const_expr(tchunks == 1): + fr_o = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + for j in cutlass.range_constexpr(vec): + fr_o[j] = cutlass.BFloat16(cutlass.Float32(fr_a[j]) / cutlass.Float32(fr_d[j])) + ooff = cute.assume(bb * d + cvec, divby=vec) + gO = cute.make_tensor(mOut.iterator + ooff, cute.make_layout(vec)) + cute.autovec_copy(fr_o, gO) + else: + base = (tidy * threads_x + tidx) * vec + for j in cutlass.range_constexpr(vec): + sM[base + j] = cutlass.Float32(fr_m[j]) + sD[base + j] = cutlass.Float32(fr_d[j]) + sA[base + j] = cutlass.Float32(fr_a[j]) + + if cutlass.const_expr(tchunks > 1): + cute.arch.barrier() + if col < ncol: + if tidy == 0: + # Fixed-order serial merge of the tchunks partials per dim. Empty + # partials (den == 0: the skipped previous-block half) are ignored; + # every row has at least one valid position (the own half), so the + # merged den is >= 1 and the final division is safe. + cvec = col * vec + fr_o = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + for j in cutlass.range_constexpr(vec): + m = cutlass.Float32(sM[tidx * vec + j]) + den = cutlass.Float32(sD[tidx * vec + j]) + acc = cutlass.Float32(sA[tidx * vec + j]) + for t in cutlass.range_constexpr(1, tchunks): + slot = (t * threads_x + tidx) * vec + j + d2 = cutlass.Float32(sD[slot]) + if d2 > 0: + m2 = cutlass.Float32(sM[slot]) + a2 = cutlass.Float32(sA[slot]) + mn = m + if m2 > mn: + mn = m2 + # exp(-inf - mn) == 0 handles a still-empty running + # partial; both-empty never reaches here (d2 > 0). + s1 = cute_math.exp(m - mn) + s2 = cute_math.exp(m2 - mn) + den = _ffma_rn(den, s1, _fmul_rn(d2, s2)) + acc = _ffma_rn(acc, s1, _fmul_rn(a2, s2)) + m = mn + fr_o[j] = cutlass.BFloat16(acc / den) + ooff = cute.assume(bb * d + cvec, divby=vec) + gO = cute.make_tensor(mOut.iterator + ooff, cute.make_layout(vec)) + cute.autovec_copy(fr_o, gO) + + +@cute.jit +def _compressor_fwd_r128_launch( + kv_ptr: cute.Pointer, + score_ptr: cute.Pointer, + ape_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + cuc_ptr: cute.Pointer, + out_ptr: cute.Pointer, + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + stream: cuda_driver.CUstream, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + tchunks: cutlass.Constexpr, + threads_x: cutlass.Constexpr, + twophase: cutlass.Constexpr, + fastexp: cutlass.Constexpr, +): + """JIT entry point that wraps raw pointers into tensors and launches forward.""" + lay = cute.make_layout(_EXT) + mKV = cute.make_tensor(kv_ptr, lay) + mScore = cute.make_tensor(score_ptr, lay) + mAPE = cute.make_tensor(ape_ptr, lay) + mCu = cute.make_tensor(cu_ptr, lay) + mCuComp = cute.make_tensor(cuc_ptr, lay) + mOut = cute.make_tensor(out_ptr, lay) + ncol = d // vec + gy = (ncol + threads_x - 1) // threads_x + _compressor_fwd_r128_kernel( + mKV, + mScore, + mAPE, + mCu, + mCuComp, + mOut, + n_seq, + ratio, + d, + coff, + vec, + tchunks, + threads_x, + twophase, + fastexp, + ).launch(grid=(nb_total, gy, 1), block=(threads_x, tchunks, 1), stream=stream) + + +_COMPILED = {} +_COMPILE_LOCK = threading.Lock() +_FAST = _FastCache() + +# nb_total-bucketed schedule tables (B200-measured; see _fwd_schedule_r128 docstring). +# Key (coff, d) -> (vec, tchunks, twophase, fastexp); threads_x is always derived +# (one warp per chunk-row). Configs absent from a table use the default schedule in +# that bucket. +_SMALL_NB_MAX = 128 +_LARGE_NB_MIN = 1024 +# 3rd field: two-phase chunk softmax (exact chunk max, then fixed-max accumulation) +# instead of the online-rescale loop — wins only where the kernel is issue-bound +# AND the fast exp alone does not already clear the issue bottleneck (after the +# fastexp lever only the c1d128 large bucket keeps it; at c2d128-131k the +# online+fastexp form measured 33.7 us vs two-phase 49.6 / two-phase+fastexp 57.8). +# 4th field: fast exp (ex2.approx path, see _exp_fast) — the remaining instruction +# diet; deterministic but tolerance-contract (not the ratio=4 expf bit pattern), so +# it is enabled per bucket by measured win (nsys pure-kernel, B200; the inline +# numbers) + the contract gate. +_SMALL_SCHEDULES = { + (1, 128): (2, 8, False, False), + (2, 128): (2, 16, False, True), # 8.21 -> 6.42 us (1x8192) + (2, 512): (2, 8, False, True), # 14.42 -> 10.91 us (1x8192) +} +_LARGE_SCHEDULES = { + (1, 128): (4, 4, True, True), # 18.58 -> 15.07 us (1x131072) + (1, 512): (4, 4, False, True), # 85.41 (v2 exact) -> 77.55 us (1x131072) + (2, 512): (4, 4, False, True), # 152.35 -> 117.71 us (1x131072) +} +# Configs whose DEFAULT bucket also runs the fast exp (measured win across the whole +# band: c2d128 3x8192 12.19 -> 8.99, 1x32768 14.43 -> 10.91, 1x65536 25.56 -> 16.39 +# us; c2d128 >= 1024 rows intentionally falls through to the same schedule, +# measured 49.57 (two-phase) -> 33.67 us at 1x131072). +_DEFAULT_FASTEXP = {(2, 128)} + + +def _fwd_schedule_r128(ratio, d, coff, nb_total=None): + """Launch schedule ``(vec, tchunks, threads_x, twophase, fastexp)`` for the + ratio=128 forward. + + ``vec = 2`` (32-bit paired bf16 accesses) for even ``d`` — except the widest rows + (``coff * d >= 1024``), where ``vec = 4`` halves the issued loads and measured + ~1.4-1.9x faster at the large shapes (d=512/coff=2) — scalar ``vec = 1`` for odd + ``d``. ``threads_x = 32`` keeps one warp per chunk-row (a warp's ``32 * vec`` + adjacent dims are one or two full 128-byte lines) and doubles the CTA count at + ``d = 128`` versus 64-wide CTAs. ``tchunks`` window chunks give ``tchunks``-way + window parallelism per row at a ~KB smem merge cost. ``coff == 2`` requires + ``tchunks >= 2`` so a chunk never straddles the half-window boundary. + + The schedule is additionally **bucketed by ``nb_total`` (output rows)** — the JIT + cache key contains the schedule, so each bucket compiles once per config: + + - **small packs** (``nb_total <= _SMALL_NB_MAX``): a 64-row pack launches only + ``nb * gy`` CTAs (128 for d=128) on a 148-SM B200 — most SMs hold a single + 4-warp CTA and global-load latency is unhidden (measured achieved occupancy + 6%, eligible warps 0.08/cycle). Doubling/quadrupling ``tchunks`` puts 2-4x the + warps on the same rows: measured on B200 (nsys pure-kernel) + c1d128 1x8192 7.75 -> ~5 us, c2d128 1x8192 10.5 -> ~7.3 us, c2d512 1x8192 + 18.9 -> ~13.3 us (with ``vec = 2``; at 64 rows occupancy beats load width). + c1d512 keeps the default (its ``gy = 8`` grid already fills the machine). + - **large contexts** (``nb_total >= _LARGE_NB_MIN``): c1d128 switches to + ``vec = 4`` / ``gy = 1`` — 1024 CTAs at 131k tokens = 0.58 waves (the default's + 2048 CTAs = 1.153 waves leave a 272-CTA tail wave that costs ~10%) — plus the + two-phase softmax and fast exp; c2d512 keeps its default geometry and adds the + fast exp. c2d128 intentionally has NO large entry: its default schedule (with + the fast exp, ``_DEFAULT_FASTEXP``) measured faster than every two-phase + variant at 131k. + - The boundaries are measured at 64 rows (win) and 256 rows (default wins); the + 64..128-row band applies the small schedule on the occupancy argument (256 + CTAs still leave SMs at <=2 of 12 resident CTAs); 129..1023 rows use the + default; >= 1024 rows use the large table where it exists. + + The ``twophase`` and ``fastexp`` fields select the per-bucket arithmetic form + (see the kernel docstring and ``_exp_fast``); ``fastexp`` buckets are the + tolerance-contract instruction diet, adopted strictly per measured win + gate. + Every schedule field is compile-time (part of the JIT cache key). All entries + and boundaries are B200-measured (nsys pure-kernel A/B per bucket) and gated on + the r128 numerics contract (``benchmark/csa/gate_csa_compressor_r128.py``). + """ + W = coff * d + if W >= 1024 and d % 4 == 0: + vec = 4 + else: + vec = 2 if d % 2 == 0 else 1 + tchunks = 8 if coff == 2 and vec == 2 else 4 + twophase = False + fastexp = (coff, d) in _DEFAULT_FASTEXP + if nb_total is not None: + if nb_total <= _SMALL_NB_MAX and (coff, d) in _SMALL_SCHEDULES: + vec, tchunks, twophase, fastexp = _SMALL_SCHEDULES[(coff, d)] + elif nb_total >= _LARGE_NB_MIN and (coff, d) in _LARGE_SCHEDULES: + vec, tchunks, twophase, fastexp = _LARGE_SCHEDULES[(coff, d)] + if d % vec != 0: + raise ValueError(f"vec={vec} must divide head_dim ({d})") + ncol = d // vec + threads_x = 32 if ncol >= 32 else ncol + win = 2 * ratio if coff == 2 else ratio + if win % tchunks != 0: + raise ValueError(f"tchunks={tchunks} must divide the window ({win})") + if coff == 2 and ratio % (win // tchunks) != 0: + raise ValueError(f"coff=2 chunks must not straddle the half-window boundary (ratio={ratio}, chunk={win // tchunks})") + if threads_x * tchunks > 1024: + raise ValueError(f"CTA too large: {threads_x} x {tchunks}") + return vec, tchunks, threads_x, twophase, fastexp + + +def _compile_fwd_r128(key, args, ratio, d, coff, schedule): + """JIT-compile the forward launch entry for ``key`` (capture-guarded).""" + with _COMPILE_LOCK: + fn = _COMPILED.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"CSA compressor (r128): first call for config {key} happened under " + "CUDA graph capture (JIT compilation is not capture-safe); compile() " + "or run one eager step for this configuration before capturing." + ) + fn = cute.compile(_compressor_fwd_r128_launch, *args, ratio, d, coff, *schedule) + _COMPILED[key] = fn + return fn + + +def precompile_fwd_r128(ratio, d, coff, device, nb_total=None): + """Ensure the ratio=128 forward kernel(s) for this configuration are JIT-compiled. + + With ``nb_total`` given, compiles exactly the schedule bucket that shape will use; + without it, compiles every bucket this config can select at runtime (small / + default / large), so a subsequent CUDA-graph capture cannot hit a cold bucket. + """ + if nb_total is not None: + candidates = [nb_total] + else: + candidates = [1, _SMALL_NB_MAX + 1, _LARGE_NB_MIN] + schedules = {_fwd_schedule_r128(ratio, d, coff, nb) for nb in candidates} + with torch.cuda.device(device): + scratch_bf16 = torch.zeros(16, device=device, dtype=torch.bfloat16) + scratch_f32 = torch.zeros(16, device=device, dtype=torch.float32) + scratch_i32 = torch.zeros(16, device=device, dtype=torch.int32) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + args = ( + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + _i32_ptr(scratch_i32), + _i32_ptr(scratch_i32), + _bf16_ptr(scratch_bf16), + cutlass.Int32(0), + cutlass.Int32(1), + stream, + ) + for schedule in schedules: + key = ("r128fwd", ratio, d, coff, schedule, device.index) + if key in _COMPILED: + continue + _compile_fwd_r128(key, args, ratio, d, coff, schedule) + + +def run_fwd_r128(kv, score, ape, cu_i, cuc_i, out, nb_total, ratio, d, coff, stream_handle=None): + """Launch the ratio=128 forward kernel (cached fast path -> slow path -> JIT). + + Same contract as ``compressor_sm100.run_fwd``: flat contiguous bf16 kv/score/out, + fp32 ape, int32 cu_seqlens/cu_seqlens_comp, launch anchored in ``kv``'s device. + ``nb_total == 0`` launches nothing. The launch schedule is bucketed by ``nb_total`` + (see ``_fwd_schedule_r128``); each bucket JIT-compiles once per config. + """ + if nb_total == 0: + return + dev = kv.device.index + schedule = _fwd_schedule_r128(ratio, d, coff, nb_total) + key = ("r128fwd", ratio, d, coff, schedule, dev) + if stream_handle is None: + stream_handle = _raw_stream(dev) + with torch.cuda.device(dev): + launcher = _FAST.get(key) + if launcher is not None: + slots = launcher.slots + slots[0].value = kv.data_ptr() + slots[1].value = score.data_ptr() + slots[2].value = ape.data_ptr() + slots[3].value = cu_i.data_ptr() + slots[4].value = cuc_i.data_ptr() + slots[5].value = out.data_ptr() + slots[6].value = nb_total + slots[7].value = cu_i.numel() - 1 + slots[8].value = stream_handle + launcher.launch() + return + stream = cuda_driver.CUstream(stream_handle) + args = ( + _bf16_ptr(kv), + _bf16_ptr(score), + _f32_ptr(ape), + _i32_ptr(cu_i), + _i32_ptr(cuc_i), + _bf16_ptr(out), + cutlass.Int32(nb_total), + cutlass.Int32(cu_i.numel() - 1), + stream, + ) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile_fwd_r128(key, args, ratio, d, coff, schedule) + fn(*args) + _FAST.put(key, fn, args) + + +# ============================================================================= +# Backward +# ============================================================================= + + +@cute.kernel +def _compressor_bwd_r128_kernel( + mKV: cute.Tensor, # flat [T * W] bf16, W = coff * d + mScore: cute.Tensor, # flat [T * W] bf16 + mAPE: cute.Tensor, # flat [ratio * W] fp32 + mCu: cute.Tensor, # [n_seq + 1] int32 + mCuComp: cute.Tensor, # [n_seq + 1] int32 + mGO: cute.Tensor, # flat [nb_total * d] bf16 + mGKV: cute.Tensor, # flat [T * W] bf16 (fully written; may be uninitialized) + mGS: cute.Tensor, # flat [T * W] bf16 (fully written; may be uninitialized) + mGAPE: cute.Tensor, # flat [ratio * W] fp32 (zero-initialized) + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + total_tokens: cutlass.Int32, + rows_per_cta: cutlass.Int32, # runtime: one compiled kernel serves every row count + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + tchunks: cutlass.Constexpr, + threads_x: cutlass.Constexpr, + fastexp: cutlass.Constexpr, +): + """Backward: staged-smem chunk-parallel phases with fused reductions. + + Phases per output row — stage -> e-pass+partials -> fixed-order merge -> store, + 4 barriers/row (CTA per ``rows_per_cta`` row group, CTA shape + ``(threads_x, tchunks)``, ``gridDim.y`` spanning head-dim column groups). This + reduction structure is the approved ratio=128 deterministic tolerance contract. + + 1. **Stage (chunk-parallel):** as in the forward, chunk-row ``tidy`` owns window + positions ``[tidy * C, (tidy + 1) * C)`` (``C = win / tchunks``); it loads its + score/APE/kv slices coalesced (lanes on head-dim columns, ``vec`` dims per + lane), stages ``s_k = f32(score) + ape`` into an fp32 smem tile and ``kv_k`` + into a bf16 smem tile, and tracks its chunk max. The invalid previous-block + half of a segment's first block (``coff == 2``, ``bis == 0``) is staged as the + constant ``(-inf, 0)`` pair, so the reductions need no validity special-casing. + 2. **e-pass + partials (chunk-parallel):** every chunk-row merges the ``tchunks`` + staged maxes for its columns (order-independent), overwrites its slots with + ``e_k = exp(s_k - mx)``, and accumulates the chunk-partial sums + ``den_c = sum e_k`` and ``S'_c = sum fma(dp_k, e_k)`` (``dp_k = go * kv_k``) + into one smem slot per (chunk, column). Invalid slots contribute exactly 0. + 3. **Merge (one lane per column):** ``den`` and ``S'`` are reduced over the + ``tchunks`` partials in a FIXED chunk order (deterministic), then the lane + publishes ``1/den`` and ``S = S'/den`` — two IEEE divisions per COLUMN instead + of one per element. + 4. **Store (chunk-parallel):** ``p_k = e_k * (1/den)`` (multiply, not divide), + ``dkv_k = bf16(go * p_k)``, ``ds_k = bf16(fma.rn(p_k, -S, mul.rn(dp_k, p_k)))``. + ``dAPE`` is accumulated into ``C * vec`` per-thread registers across the CTA's + ``rows_per_cta`` rows and reduced with one fp32 atomic per ``(k, dim)`` per CTA + at the end. + + The den/S reduction ORDER (per-chunk then fixed merge) and the reciprocal + rounding differ from the eager serial scan — dKV/dScore are deterministic but + tolerance-vs-eager, gate-checked against the forward-style tolerances and an + fp64 oracle. Measured 1.10-1.47x on B200 across the shipped envelope over a + bitwise-pinned reduction structure (serial ascending-``k`` ``den``/``S`` sweeps + + a per-element ``div.rn`` p-pass, the ratio=4 op order). + + Kernel-side zero-writes keep the ratio=4 ownership verbatim, parallelized across + the CTA (each class is up to 127 tokens at ratio=128, or 128 rows for the + ``coff == 2`` first-half class): + + - per-segment tail tokens (``seqlen % ratio``, both halves) and, for + ``coff == 2``, the first-half columns of the segment's LAST block's own tokens + — written by that last block's CTA, tokens strided over ``tidy``; + - all tokens of segments with zero output blocks (``seqlen < ratio``) — written + by the ``bidx == 0`` CTA column, tokens strided over ``tidy``; + - tokens beyond ``cu_seqlens[-1]`` (static token-capacity padding) — strided + over ``(bidx, tidy)`` grid rows. + + Rows in ``[cu_seqlens_comp[-1], nb_total)`` are static-capacity padding; their + incoming gradients are ignored, as in the ratio=4 backward. + """ + tidx, tidy, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + gdimx, _, _ = cute.arch.grid_dim() + ncol: cutlass.Constexpr = d // vec # vec-group count per output row + win: cutlass.Constexpr = 2 * ratio if coff == 2 else ratio + C: cutlass.Constexpr = win // tchunks + W: cutlass.Constexpr = coff * d + cols_pc: cutlass.Constexpr = threads_x * vec # head-dim columns per CTA + col = bidy * threads_x + tidx + cvec = col * vec + ZERO_BF16 = cutlass.BFloat16(0.0) + + smem = cutlass.utils.SmemAllocator() + # [win][cols_pc] fp32 tile: holds s_k after stage, e_k after the e-pass (``p_k`` + # is formed at store time as ``e_k * (1/den)``). [win][cols_pc] bf16 tile: staged + # kv. Small buffers for the chunk maxes and the published per-column den / S. + sE = smem.allocate_tensor(cutlass.Float32, cute.make_layout(win * cols_pc), 16) + sKV = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(win * cols_pc), 16) + sMax = smem.allocate_tensor(cutlass.Float32, cute.make_layout(tchunks * cols_pc), 16) + sDen = smem.allocate_tensor(cutlass.Float32, cute.make_layout(cols_pc), 16) + sS = smem.allocate_tensor(cutlass.Float32, cute.make_layout(cols_pc), 16) + # Per-chunk partial sums of den and S' = sum_k fma(dp_k, e_k) (the tolerance + # contract's fused reductions), merged in fixed chunk order by one lane per + # column — no serial sweeps, and the per-element p-pass division becomes + # p = e * (1/den) at store time. + sDenP = smem.allocate_tensor(cutlass.Float32, cute.make_layout(tchunks * cols_pc), 16) + sSP = smem.allocate_tensor(cutlass.Float32, cute.make_layout(tchunks * cols_pc), 16) + + nb_valid = mCuComp[n_seq] + + fr_z = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + for j in cutlass.range_constexpr(vec): + fr_z[j] = ZERO_BF16 + + # --- zero sweeps for never-consumed slots with no owning block row --- + if col < ncol: + # Segments with zero output blocks (seqlen < ratio, up to 127 tokens each): + # owned by the bidx == 0 CTA column, tokens strided over the chunk-rows. + if bidx == 0: + for sg in cutlass.range(n_seq): + if mCuComp[sg + 1] == mCuComp[sg]: + t0 = mCu[sg] + cnt = mCu[sg + 1] - t0 + for i in cutlass.range((cnt - tidy + tchunks - 1) // tchunks): + t = t0 + tidy + i * tchunks + off = cute.assume(t * W + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + off, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + off, cute.make_layout(vec))) + if cutlass.const_expr(coff == 2): + off2 = cute.assume(t * W + d + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + off2, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + off2, cute.make_layout(vec))) + + # Static token-capacity padding of the gradient buffers: strided over + # (bidx, tidy) grid rows. The quotient/remainder split keeps every + # intermediate within int32 for any count < 2**31, as in the ratio=4 kernel. + pad0 = mCu[n_seq] + pad_count = total_tokens - pad0 + gr = bidx * tchunks + tidy + nrows = gdimx * tchunks + if gr < pad_count: + my_count = pad_count // nrows + if gr < pad_count % nrows: + my_count = my_count + 1 + for i in cutlass.range(my_count): + t = pad0 + gr + i * nrows + off = cute.assume(t * W + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + off, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + off, cute.make_layout(vec))) + if cutlass.const_expr(coff == 2): + off2 = cute.assume(t * W + d + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + off2, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + off2, cute.make_layout(vec))) + + # Per-thread chunk mapping, row-independent parts. A chunk never straddles the + # coff == 2 half-window boundary (ratio % C == 0 enforced by the schedule), so + # projection column base and APE row base are per-thread constants. + k0 = tidy * C + colbase = cutlass.Int32(cvec) + ape_row0 = cutlass.Int32(k0) + if cutlass.const_expr(coff == 2): + if k0 >= ratio: + colbase = cutlass.Int32(d + cvec) + ape_row0 = cutlass.Int32(k0 - ratio) + + # dAPE accumulator: C * vec fp32 registers (<= 64 under the schedule defaults; + # ptxas-verified spill-free), accumulated across the CTA's rows_per_cta rows, + # one atomic per slot at the end. ``cutlass.range(..., unroll_full=True)``, not + # ``range_constexpr``: same unrolled register zero-init with identical ptxas + # register/spill counts and a slightly smaller prologue (the constexpr form + # emitted one mov per slot and trips the DSL's slow-compile warning at the + # 64-iteration schedules). + fr_dape = cute.make_rmem_tensor((C * vec,), cutlass.Float32) + for q in cutlass.range(C * vec, unroll_full=True): + fr_dape[q] = cutlass.Float32(0.0) + + for rr in cutlass.range(rows_per_cta): + bb = bidx * rows_per_cta + rr + if bb < nb_valid: + # Per-segment boundary scan (n_seq is small), as in the ratio=4 kernel. + seq_idx = cutlass.Int32(0) + bis = cutlass.Int32(bb) + for sg in cutlass.range(n_seq): + cs = mCuComp[sg] + ce = mCuComp[sg + 1] + if bb >= cs: + if bb < ce: + seq_idx = sg + bis = bb - cs + tok0 = mCu[seq_idx] + bis * ratio + is_last = bb + 1 == mCuComp[seq_idx + 1] + + # Both coff forms share one chunk loop body: token row tok0 - ratio + k + # (== tok0 + k - ratio for the own half); only column base / APE row / + # validity differ per half (all per-thread constants above). + tok_row0 = tok0 + k0 + run_chunk = cutlass.Boolean(True) + if cutlass.const_expr(coff == 2): + tok_row0 = tok0 - ratio + k0 + if k0 < ratio: + if bis == 0: + run_chunk = cutlass.Boolean(False) + + # ---- phase 1: chunk-parallel stage (s, kv -> smem tiles, chunk max) ---- + if col < ncol: + sbase = k0 * cols_pc + tidx * vec + fr_m = cute.make_rmem_tensor((vec,), cutlass.Float32) + for j in cutlass.range_constexpr(vec): + fr_m[j] = cutlass.Float32(_NEG_INF) + if run_chunk: + for kk in cutlass.range(C, unroll=4): + off = cute.assume((tok_row0 + kk) * W + colbase, divby=vec) + aoff = cute.assume((ape_row0 + kk) * W + colbase, divby=vec) + fr_s = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_k = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_a = cute.make_rmem_tensor((vec,), cutlass.Float32) + cute.autovec_copy(cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)), fr_s) + cute.autovec_copy(cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)), fr_k) + cute.autovec_copy(cute.make_tensor(mAPE.iterator + aoff, cute.make_layout(vec)), fr_a) + for j in cutlass.range_constexpr(vec): + s = cutlass.Float32(fr_s[j]) + cutlass.Float32(fr_a[j]) + sE[sbase + kk * cols_pc + j] = s + sKV[sbase + kk * cols_pc + j] = fr_k[j] + if s > fr_m[j]: + fr_m[j] = s + else: + # Invalid previous-block half: the constant (-inf, 0) pair, exactly + # the values the ratio=4 kernel feeds its serial window. + for kk in cutlass.range(C, unroll=4): + for j in cutlass.range_constexpr(vec): + sE[sbase + kk * cols_pc + j] = cutlass.Float32(_NEG_INF) + sKV[sbase + kk * cols_pc + j] = ZERO_BF16 + for j in cutlass.range_constexpr(vec): + sMax[tidy * cols_pc + tidx * vec + j] = cutlass.Float32(fr_m[j]) + cute.arch.barrier() + + # ---- phase 2: chunk-parallel e-pass (mx merge + exp + den/S' partials) ---- + # This pass ALSO accumulates the chunk's partial den / S' sums + # (registers, then one smem slot per (chunk, column)). + if col < ncol: + sbase = k0 * cols_pc + tidx * vec + fr_go2 = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + gooff2 = cute.assume(bb * d + cvec, divby=vec) + cute.autovec_copy(cute.make_tensor(mGO.iterator + gooff2, cute.make_layout(vec)), fr_go2) + for j in cutlass.range_constexpr(vec): + # Chunk-max merge: max is order-independent, so the chunked max + # equals the ratio=4 kernel's serial scan bitwise. Redundant per + # chunk-row (cheap smem broadcasts), which keeps the barrier count + # down. mx is finite: the own half always has valid positions. + mx = cutlass.Float32(sMax[tidx * vec + j]) + for t in cutlass.range_constexpr(1, tchunks): + v = cutlass.Float32(sMax[t * cols_pc + tidx * vec + j]) + if v > mx: + mx = v + # Invalid slots become exp(-inf - mx) == 0 exactly, the value the + # ratio=4 kernel's serial window feeds its den sum. + go2 = cutlass.Float32(fr_go2[j]) + den_p = cutlass.Float32(0.0) + sp_p = cutlass.Float32(0.0) + for kk in cutlass.range(C, unroll=8): + slot = sbase + kk * cols_pc + j + if cutlass.const_expr(fastexp): + e = _exp_fast(cutlass.Float32(sE[slot]) - mx) + else: + e = cute_math.exp(cutlass.Float32(sE[slot]) - mx) + sE[slot] = e + den_p = den_p + e + dp = go2 * cutlass.Float32(sKV[slot]) + sp_p = _ffma_rn(dp, e, sp_p) + sDenP[tidy * cols_pc + tidx * vec + j] = den_p + sSP[tidy * cols_pc + tidx * vec + j] = sp_p + cute.arch.barrier() + + # ---- phase 3: fused merge (tolerance contract): den and S' merged in + # fixed chunk order by one lane per column; publish 1/den (the store + # phase multiplies) and S = S'/den. Deterministic: fixed order, no + # atomics; the reduction ORDER differs from the eager serial scan, and + # the per-element division becomes a reciprocal multiply — both inside + # the approved tolerance contract (gate-checked). ---- + if tidy < vec: + c = tidy * threads_x + tidx + if bidy * cols_pc + c < d: + den = cutlass.Float32(0.0) + sp = cutlass.Float32(0.0) + for t in cutlass.range_constexpr(tchunks): + den = den + cutlass.Float32(sDenP[t * cols_pc + c]) + sp = sp + cutlass.Float32(sSP[t * cols_pc + c]) + # den >= 1 (the max element contributes exp(0)); both divisions + # are exact IEEE div.rn, once per column instead of per element. + sDen[c] = cutlass.Float32(1.0) / den + sS[c] = sp / den + cute.arch.barrier() + + # ---- phase 4: chunk-parallel gradient stores + dAPE accumulation ---- + if col < ncol: + sbase = k0 * cols_pc + tidx * vec + fr_go = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + gooff = cute.assume(bb * d + cvec, divby=vec) + cute.autovec_copy(cute.make_tensor(mGO.iterator + gooff, cute.make_layout(vec)), fr_go) + if run_chunk: + for kk in cutlass.range_constexpr(C): + off = cute.assume((tok_row0 + kk) * W + colbase, divby=vec) + fr_gkv = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_gs = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + for j in cutlass.range_constexpr(vec): + # sE holds e_k; sDen holds 1/den — the p-pass division + # became one multiply per element (tolerance contract). + p = _fmul_rn(cutlass.Float32(sE[sbase + kk * cols_pc + j]), cutlass.Float32(sDen[tidx * vec + j])) + go = cutlass.Float32(fr_go[j]) + dp = go * cutlass.Float32(sKV[sbase + kk * cols_pc + j]) + ds = _ffma_rn(p, -cutlass.Float32(sS[tidx * vec + j]), _fmul_rn(dp, p)) + fr_gkv[j] = cutlass.BFloat16(go * p) + fr_gs[j] = cutlass.BFloat16(ds) + fr_dape[kk * vec + j] = cutlass.Float32(fr_dape[kk * vec + j]) + ds + cute.autovec_copy(fr_gkv, cute.make_tensor(mGKV.iterator + off, cute.make_layout(vec))) + cute.autovec_copy(fr_gs, cute.make_tensor(mGS.iterator + off, cute.make_layout(vec))) + + # The segment's last block zeroes the never-consumed slots it uniquely + # owns: (a) for coff == 2 the first-half columns of its own tokens + # (no next block consumes them), (b) the segment's tail tokens + # (seqlen % ratio, both halves). Tokens strided over the chunk-rows. + if is_last: + if cutlass.const_expr(coff == 2): + for i in cutlass.range((ratio - tidy + tchunks - 1) // tchunks): + t = tok0 + tidy + i * tchunks + offz = cute.assume(t * W + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + offz, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + offz, cute.make_layout(vec))) + tail0 = tok0 + ratio + cnt = mCu[seq_idx + 1] - tail0 + for i in cutlass.range((cnt - tidy + tchunks - 1) // tchunks): + t = tail0 + tidy + i * tchunks + offz = cute.assume(t * W + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + offz, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + offz, cute.make_layout(vec))) + if cutlass.const_expr(coff == 2): + offz2 = cute.assume(t * W + d + cvec, divby=vec) + cute.autovec_copy(fr_z, cute.make_tensor(mGKV.iterator + offz2, cute.make_layout(vec))) + cute.autovec_copy(fr_z, cute.make_tensor(mGS.iterator + offz2, cute.make_layout(vec))) + # Tiles are reused by the next row's stage phase. + cute.arch.barrier() + + # One fp32 atomic per owned (k, dim) per CTA (amortized over rows_per_cta rows). + # Chunks that never ran (invalid halves, padding-row CTAs) accumulated 0.0. + if col < ncol: + for kk in cutlass.range_constexpr(C): + for j in cutlass.range_constexpr(vec): + cute_arch.atomic_add(mGAPE.iterator + ((ape_row0 + kk) * W + colbase + j), cutlass.Float32(fr_dape[kk * vec + j])) + + +@cute.jit +def _compressor_bwd_r128_launch( + kv_ptr: cute.Pointer, + score_ptr: cute.Pointer, + ape_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + cuc_ptr: cute.Pointer, + go_ptr: cute.Pointer, + gkv_ptr: cute.Pointer, + gs_ptr: cute.Pointer, + gape_ptr: cute.Pointer, + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + total_tokens: cutlass.Int32, + rows_per_cta: cutlass.Int32, + stream: cuda_driver.CUstream, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + tchunks: cutlass.Constexpr, + threads_x: cutlass.Constexpr, + fastexp: cutlass.Constexpr, +): + """JIT entry point that wraps raw pointers into tensors and launches backward.""" + lay = cute.make_layout(_EXT) + mKV = cute.make_tensor(kv_ptr, lay) + mScore = cute.make_tensor(score_ptr, lay) + mAPE = cute.make_tensor(ape_ptr, lay) + mCu = cute.make_tensor(cu_ptr, lay) + mCuComp = cute.make_tensor(cuc_ptr, lay) + mGO = cute.make_tensor(go_ptr, lay) + mGKV = cute.make_tensor(gkv_ptr, lay) + mGS = cute.make_tensor(gs_ptr, lay) + mGAPE = cute.make_tensor(gape_ptr, lay) + ncol = d // vec + gx = (nb_total + rows_per_cta - 1) // rows_per_cta + gy = (ncol + threads_x - 1) // threads_x + _compressor_bwd_r128_kernel( + mKV, + mScore, + mAPE, + mCu, + mCuComp, + mGO, + mGKV, + mGS, + mGAPE, + nb_total, + n_seq, + total_tokens, + rows_per_cta, + ratio, + d, + coff, + vec, + tchunks, + threads_x, + fastexp, + ).launch(grid=(gx, gy, 1), block=(threads_x, tchunks, 1), stream=stream) + + +# Small-pack backward bucket (mirrors the forward's nb_total bucketing; the JIT key +# contains the schedule so each bucket compiles once per config). At d=128 a small +# pack launches only nb * 2 CTAs at vec=2; vec=1 halves the columns per CTA and +# doubles gridDim.y, which fills the underoccupied machine (clock-safe interleaved +# A/B, B200: c1d128 1x8192 13.5 -> 11.5 us, 3x8192 21.8 -> 19.7, +# c2d128 1x8192 17.3 -> 13.8) and LOSES from 256 rows on (1x32768: 0.85x/0.65x) — +# boundary measured at 192 (win) / 256 (loss). The vec=1 buckets keep the exact exp: +# fastexp measured 0.905x (c1) / 1.003x (c2) on top of them. +_BWD_SMALL_NB_MAX = 192 +# (coff, d) -> (vec, tchunks, fastexp); threads_x is always derived (one warp per +# chunk-row). +_BWD_SMALL_SCHEDULES = { + (1, 128): (1, 8, False), + (2, 128): (1, 8, False), +} + +_SM_COUNT_CACHE = {} + + +def _sm_count(dev): + """SM count for device index ``dev`` (cached; used by the rows_per_cta pick).""" + n = _SM_COUNT_CACHE.get(dev) + if n is None: + n = torch.cuda.get_device_properties(dev).multi_processor_count + _SM_COUNT_CACHE[dev] = n + return n + + +def _bwd_schedule_r128(ratio, d, coff, nb_total=None): + """Launch schedule ``(vec, tchunks, threads_x, fastexp)`` for the ratio=128 + backward. + + ``vec = 2`` (32-bit paired bf16 accesses) for even ``d``, scalar ``vec = 1`` for + odd ``d`` — no ``vec = 4`` variant: it doubles the smem tiles and registers + (ptxas 186-195 regs) and measured 0.42-0.63x from the residency collapse. + **Small packs** (``nb_total <= _BWD_SMALL_NB_MAX``, d=128) switch to ``vec = 1`` + (see ``_BWD_SMALL_SCHEDULES``). ``threads_x = 32`` keeps one warp per chunk-row. + ``tchunks`` trades chunk-parallel width against resident CTAs per SM (more CTAs + overlap other CTAs' phases): measured optimum ``tchunks = 8`` across the + envelope except ``coff = 1`` at ``d >= 512``, where the 8x wider column grid + already fills the machine and ``tchunks = 4`` (fewer threads, one more CTA/SM) + wins ~25% at long context. ``coff == 2`` additionally requires chunks not to + straddle the half-window boundary. ``fastexp`` (see ``_exp_fast``) measured a + uniform +3-8% everywhere EXCEPT the vec=1 small buckets — the default is True + outside them. + """ + W = coff * d + vec = 2 if d % 2 == 0 else 1 + tchunks = 4 if coff == 1 and d >= 512 else 8 + fastexp = True + if nb_total is not None: + if nb_total <= _BWD_SMALL_NB_MAX and (coff, d) in _BWD_SMALL_SCHEDULES: + vec, tchunks, fastexp = _BWD_SMALL_SCHEDULES[(coff, d)] + if d % vec != 0: + raise ValueError(f"vec={vec} must divide head_dim ({d})") + ncol = d // vec + threads_x = 32 if ncol >= 32 else ncol + win = 2 * ratio if coff == 2 else ratio + if win % tchunks != 0: + raise ValueError(f"tchunks={tchunks} must divide the window ({win})") + if coff == 2 and ratio % (win // tchunks) != 0: + raise ValueError(f"coff=2 chunks must not straddle the half-window boundary (ratio={ratio}, chunk={win // tchunks})") + if threads_x * tchunks > 1024: + raise ValueError(f"CTA too large: {threads_x} x {tchunks}") + # smem: fp32 s/e tile + bf16 kv tile + max merge + den/S' partials + den/S + # publish, 16 B aligned. + smem_bytes = win * threads_x * vec * 6 + 3 * tchunks * threads_x * vec * 4 + 2 * threads_x * vec * 4 + 64 + if smem_bytes > 227 * 1024: + raise ValueError(f"backward smem tile too large ({smem_bytes} B) for W={W}") + return vec, tchunks, threads_x, fastexp + + +def _bwd_rows_per_cta(nb_total, ratio, d, coff, dev): + """Runtime ``rows_per_cta`` for the backward launch. + + The measured optimum is the smallest R whose grid fits ONE resident wave + (``ctas_per_sm * sm_count``): the per-CTA row loop is a serial pipeline, so a + grid a few percent over one wave costs a whole extra wave (measured cliffs: + 1.16 waves = 1.5x one wave at c2/d512/65536-token). Larger R additionally + amortizes the dAPE atomics and the per-row segment scan, which is why R grows + with the pack instead of capping at one row. ``ctas_per_sm`` is the schedule's + measured residency (register/smem-bound): 2 for the coff=2 schedule (T=8, + ~101 KB smem), 3 for coff=1 T=8, 4 for the coff=1 T=4 (d >= 512) schedule. + Capped at 16 (beyond one wave per SM the pipeline depth stops paying). + """ + vec, tchunks, threads_x = _bwd_schedule_r128(ratio, d, coff, nb_total)[:3] + gy = (d // vec + threads_x - 1) // threads_x + ctas_per_sm = 2 if coff == 2 else (4 if tchunks <= 4 else 3) + slots = ctas_per_sm * _sm_count(dev) + return max(1, min(16, -((nb_total * gy) // -slots))) + + +def _compile_bwd_r128(key, args, ratio, d, coff, schedule): + """JIT-compile the backward launch entry for ``key`` (capture-guarded).""" + with _COMPILE_LOCK: + fn = _COMPILED.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"CSA compressor (r128): first call for config {key} happened under " + "CUDA graph capture (JIT compilation is not capture-safe); compile() " + "or run one eager step for this configuration before capturing." + ) + fn = cute.compile(_compressor_bwd_r128_launch, *args, ratio, d, coff, *schedule) + _COMPILED[key] = fn + return fn + + +def precompile_bwd_r128(ratio, d, coff, device, nb_total=None): + """Ensure the ratio=128 backward kernel(s) for this configuration are JIT-compiled. + + With ``nb_total`` given, compiles exactly the schedule bucket that shape will use; + without it, compiles every bucket this config can select at runtime (small / + default), so a subsequent CUDA-graph capture cannot hit a cold bucket. + """ + if nb_total is not None: + candidates = [nb_total] + else: + candidates = [1, _BWD_SMALL_NB_MAX + 1] + schedules = {_bwd_schedule_r128(ratio, d, coff, nb) for nb in candidates} + with torch.cuda.device(device): + scratch_bf16 = torch.zeros(16, device=device, dtype=torch.bfloat16) + scratch_f32 = torch.zeros(16, device=device, dtype=torch.float32) + scratch_i32 = torch.zeros(16, device=device, dtype=torch.int32) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + args = ( + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + _i32_ptr(scratch_i32), + _i32_ptr(scratch_i32), + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + cutlass.Int32(1), + stream, + ) + for schedule in schedules: + key = ("r128bwd", ratio, d, coff, schedule, device.index) + if key in _COMPILED: + continue + _compile_bwd_r128(key, args, ratio, d, coff, schedule) + + +def run_bwd_r128(kv, score, ape, cu_i, cuc_i, go, gkv, gs, gape, nb_total, ratio, d, coff, stream_handle=None): + """Launch the ratio=128 backward kernel (cached fast path -> slow path -> JIT). + + Same contract as ``compressor_sm100.run_bwd``: recompute-from-inputs (flat + contiguous bf16 kv/score/grad_out, fp32 ape, int32 cu_seqlens/cu_seqlens_comp), + ``gkv``/``gs`` fully written (may be uninitialized), ``gape`` zero-initialized + (fp32 atomics), launch anchored in ``kv``'s device. ``nb_total == 0`` launches + nothing. + """ + if nb_total == 0: + return + dev = kv.device.index + schedule = _bwd_schedule_r128(ratio, d, coff, nb_total) + key = ("r128bwd", ratio, d, coff, schedule, dev) + total_tokens = kv.numel() // (coff * d) + rows = _bwd_rows_per_cta(nb_total, ratio, d, coff, dev) + if stream_handle is None: + stream_handle = _raw_stream(dev) + with torch.cuda.device(dev): + launcher = _FAST.get(key) + if launcher is not None: + slots = launcher.slots + slots[0].value = kv.data_ptr() + slots[1].value = score.data_ptr() + slots[2].value = ape.data_ptr() + slots[3].value = cu_i.data_ptr() + slots[4].value = cuc_i.data_ptr() + slots[5].value = go.data_ptr() + slots[6].value = gkv.data_ptr() + slots[7].value = gs.data_ptr() + slots[8].value = gape.data_ptr() + slots[9].value = nb_total + slots[10].value = cu_i.numel() - 1 + slots[11].value = total_tokens + slots[12].value = rows + slots[13].value = stream_handle + launcher.launch() + return + stream = cuda_driver.CUstream(stream_handle) + args = ( + _bf16_ptr(kv), + _bf16_ptr(score), + _f32_ptr(ape), + _i32_ptr(cu_i), + _i32_ptr(cuc_i), + _bf16_ptr(go), + _bf16_ptr(gkv), + _bf16_ptr(gs), + _f32_ptr(gape), + cutlass.Int32(nb_total), + cutlass.Int32(cu_i.numel() - 1), + cutlass.Int32(total_tokens), + cutlass.Int32(rows), + stream, + ) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile_bwd_r128(key, args, ratio, d, coff, schedule) + fn(*args) + _FAST.put(key, fn, args) diff --git a/test/python/fe_api/csa/test_CSA_compressor.py b/test/python/fe_api/csa/test_CSA_compressor.py new file mode 100644 index 000000000..718fad596 --- /dev/null +++ b/test/python/fe_api/csa/test_CSA_compressor.py @@ -0,0 +1,1318 @@ +"""Tests for the fused CSA/HCA Compressor gated-pooling kernels (``cudnn.csa``). + +Ported with the kernels from Megatron-LM (https://github.com/NVIDIA/Megatron-LM/pull/5984, +measurements and numerics in https://github.com/NVIDIA/Megatron-LM/issues/5968). Covers: + + - numerics of the fused region vs an fp32-intermediate eager reference, per ratio + family: at ``ratio == 4`` ``dKV``/``dScore`` are bit-identical and the forward is + within one bf16 rounding step on a tiny fraction of elements; at ``ratio == 128`` + the contract is faithfulness to the fp32 eager reference: all three match it + within tolerance thresholds (differing elements <= max(1, 0.1%), max_abs <= + 1.6e-2) calibrated on this suite's input distribution — bf16's grid is relative, + so absolute deviations scale with input magnitude — and carry fp64-oracle parity + on inputs whose fp32 intermediates stay finite (inputs that overflow the + reference's fp32 intermediates reproduce its NaN pattern instead: the committed + gate's overflow-intermediate case; see docs/fe-oss-apis/csa.md and the gate's + scaled-input case). Both families are also compared against the verbatim + upstream eager numerics (tolerance), over ragged THD packs including segments + shorter than ``ratio``; + - static-capacity padding rows (``total_comp > cu_seqlens_comp[-1]``); + - kernel-side zero-writes to never-consumed ``dKV``/``dScore`` slots: NaN-canary + (uninitialized) gradient buffers stay bitwise-equal to zero-initialized runs, + every never-consumed slot class is asserted EXACTLY zero, the consumed slots + match the eager reference per the ratio's contract, and the ``total_comp == 0`` + host fallback still hands back exact zeros; + - run-to-run determinism of forward / ``dKV`` / ``dScore`` (``dAPE`` uses fp32 atomics + and is exempt by design; the backward refuses to run under + ``torch.use_deterministic_algorithms(True)``, and only accumulates into + ``grad_ape`` — re-zeroing is the caller's job on the class API); + - the ratio=128 dispatch envelope: schedule selection at every nb_total bucket + boundary, and (L1) one execution of every shipped (config, schedule) kernel — + fast-exp, two-phase and vec=1 buckets included — against the contract; + - CUDA graph capture: warmup -> capture fwd+bwd -> replay (including replay with new + data and a smaller device-side true row count, checked on all four outputs and + bitwise against a direct call), and the loud error when the first call for a + configuration would JIT under capture; + - ``check_support`` boundaries (validated envelope: CC 10.0; ratio 4 with coff + {1, 2}, ratio 128 with coff {1, 2} x head_dim {128, 512}; BF16 kv/score, FP32 ape, + int32 cu_seqlens and int32 flat-offset bounds). + +The eager reference below mirrors the exact region of Megatron-LM +``Compressor._forward_thd`` (non-pre-grouped THD path) that the fused kernels replace: +gather-index build -> gather -> ``+ APE`` -> overlap-window transform (``coff == 2`` +only, ``Compressor._overlap_transform_thd``; ``coff == 1`` keeps the block's own +``ratio``-token window) -> fp32 softmax -> gated weighted sum -> bf16 cast. ``mode`` +selects the numerics: "upstream" reproduces the eager code exactly +(softmax weights rounded to bf16, bf16 multiply); "fp32" keeps all intermediates fp32 +with a single final bf16 rounding (the fused kernels' numerics); "fp64" is an oracle. +""" + +import pytest +import torch + + +def _import_compressor(): + """Import ``cudnn.csa.compressor``, skipping only on a missing cutedsl stack.""" + # Skip only when the optional cutedsl dependency stack is missing; a broken + # cudnn.csa package itself must fail the tests, not skip them. + pytest.importorskip("cutlass", reason="Environment not supported: cudnn[cutedsl] not installed") + pytest.importorskip("cuda.bindings", reason="Environment not supported: cuda-python not installed") + from cudnn.csa import compressor + + return compressor + + +def _require_sm100(): + """Skip the test unless a CC 10.0 (Blackwell) CUDA GPU is available.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA GPU required") + if torch.cuda.get_device_capability() != (10, 0): + pytest.skip("compute capability 10.0 GPU required") + + +# --------------------------------------------------------------------------- +# Eager reference (self-contained mirror of the Megatron-LM eager region) +# --------------------------------------------------------------------------- + + +def _batch_of_row(cu_seqlens, total): + """Segment index owning each packed row (mirror of Megatron-LM ``batch_of_row``).""" + n_seg = cu_seqlens.shape[0] - 1 + row_idx = torch.arange(total, device=cu_seqlens.device, dtype=torch.int64) + return torch.bucketize(row_idx, cu_seqlens[1:], right=True).clamp(max=max(n_seg - 1, 0)) + + +def _overlap_transform_thd(tensor, is_first_in_seg, head_dim, fill_value=0): + """Mirror of Megatron-LM ``Compressor._overlap_transform_thd``. + + Input shape: [total_comp, ratio, b, coff * head_dim] + Output shape: [total_comp, 2 * ratio, b, head_dim] + """ + n, ratio, b_dim, _ = tensor.size() + d = head_dim + new_tensor = tensor.new_full((n, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + # Previous group's first-half data -- shift by 1 along dim-0. + prev_data = torch.roll(tensor[:, :, :, :d], shifts=1, dims=0) + # Zero-fill (or fill_value-fill) segment boundaries. + prev_data[is_first_in_seg] = fill_value + new_tensor[:, :ratio] = prev_data + return new_tensor + + +def _eager_pool(kv, score, ape, cu_seqlens, cu_seqlens_comp, total_comp, ratio, d, coff, mode): + """Eager pooling region (see module docstring); ``coff == 1`` skips the overlap transform.""" + device = kv.device + row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_comp.dtype) + batch_ids = _batch_of_row(cu_seqlens_comp, total_comp) + valid_comp = row_idx < cu_seqlens_comp[-1] + local_pos = row_idx - cu_seqlens_comp[batch_ids] + local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) + base = cu_seqlens[batch_ids].unsqueeze(1) + local_pos.unsqueeze(1) * ratio + base = torch.where(valid_comp.unsqueeze(1), base, torch.zeros_like(base)) + offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) + gather_idx = base + offsets # (total_comp, ratio) + + if mode == "fp32": + kv = kv.float() + score = score.float() + elif mode == "fp64": + kv = kv.double() + score = score.double() + ape = ape.double() + + kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) + score_grouped = score[gather_idx] + score_grouped = score_grouped + ape.view(1, ratio, 1, -1) + + if coff == 2: + is_first = local_pos == 0 + kv_grouped = _overlap_transform_thd(kv_grouped, is_first, d, fill_value=0) + score_grouped = _overlap_transform_thd(score_grouped, is_first, d, fill_value=float("-inf")) + + if mode == "upstream": + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32).to(kv_grouped.dtype) + out = (kv_grouped * weights).sum(dim=1) + elif mode == "fp32": + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32) + out = (kv_grouped * weights).sum(dim=1).to(torch.bfloat16) + else: # fp64 oracle + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float64) + out = (kv_grouped * weights).sum(dim=1) + return out # (total_comp, 1, d) + + +# --------------------------------------------------------------------------- +# Input construction and runners +# --------------------------------------------------------------------------- + + +def _make_inputs(lens, d, ratio, coff, seed=1234, device="cuda"): + """Build a seeded random THD pack (kv, score, ape, cu, cuc, total_comp, grad_out) for ``lens``.""" + total = sum(lens) + w = coff * d + gen = torch.Generator(device="cpu").manual_seed(seed) + kv = torch.randn(total, 1, w, generator=gen, dtype=torch.float32).to(torch.bfloat16) + score = (torch.randn(total, 1, w, generator=gen, dtype=torch.float32).mul_(1.5)).to(torch.bfloat16) + ape = torch.randn(ratio, w, generator=gen, dtype=torch.float32).mul_(0.25) + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device=device) + seg_comp = torch.tensor([seg_len // ratio for seg_len in lens]) + cuc = torch.tensor([0] + list(seg_comp.cumsum(0)), dtype=torch.int32, device=device) + total_comp = int(cuc[-1].item()) + go = torch.randn(total_comp, 1, d, generator=gen, dtype=torch.float32).to(torch.bfloat16) + return kv.to(device), score.to(device), ape.to(device), cu, cuc, total_comp, go.to(device) + + +def _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode): + """Forward + backward through the eager reference; returns (out, dKV, dScore, dAPE).""" + dtype = torch.float64 if mode == "fp64" else None + kv_l = (kv.to(dtype) if dtype else kv.clone()).requires_grad_(True) + score_l = (score.to(dtype) if dtype else score.clone()).requires_grad_(True) + ape_l = (ape.to(dtype) if dtype else ape.clone()).requires_grad_(True) + out = _eager_pool(kv_l, score_l, ape_l, cu, cuc, total_comp, ratio, d, coff, mode) + out.backward(go.to(out.dtype)) + torch.cuda.synchronize() + return out.detach(), kv_l.grad.detach(), score_l.grad.detach(), ape_l.grad.detach() + + +def _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go): + """Forward + backward through the fused wrappers; returns (out, dKV, dScore, dAPE).""" + compressor = _import_compressor() + total = kv.shape[0] + out = compressor.csa_compressor_forward_wrapper( + kv.view(total, -1), + score.view(total, -1), + ape, + cu, + cuc, + ratio=ratio, + head_dim=d, + coff=coff, + total_comp=total_comp, + )["out"] + grads = compressor.csa_compressor_backward_wrapper( + kv.view(total, -1), + score.view(total, -1), + ape, + cu, + cuc, + go.view(total_comp, d), + ratio=ratio, + head_dim=d, + coff=coff, + ) + torch.cuda.synchronize() + return ( + out.view(total_comp, 1, d), + grads["grad_kv"].view_as(kv), + grads["grad_score"].view_as(score), + grads["grad_ape"], + ) + + +def _assert_grads_vs_fp32(gkv, gs, ref_kv, ref_s, ratio): + """dKV/dScore vs the fp32-intermediate eager reference, per the ratio's contract. + + ratio=4: bit-identical (the production bitwise contract, unchanged). + ratio=128: deterministic tolerance contract — the fused backward reorders the + den/S reductions (fixed chunk merge) and hoists 1/den, and some forward buckets + use the ex2.approx fast exp, so dKV/dScore match eager within the forward-style + tolerances instead of bitwise (thresholds calibrated on this suite's input + distribution; they stay bitwise run-to-run, and the fp64-oracle parity assertion + below keeps the accuracy honest on the tested finite-intermediate inputs). + """ + if ratio == 4: + assert torch.equal(gkv, ref_kv), "dKV must be bit-identical to the fp32 reference at ratio=4" + assert torch.equal(gs, ref_s), "dScore must be bit-identical to the fp32 reference at ratio=4" + return + for name, fused_t, ref_t in (("dKV", gkv, ref_kv), ("dScore", gs, ref_s)): + diff = (fused_t.float() - ref_t.float()).abs() + n_diff = (fused_t != ref_t).sum().item() + assert n_diff <= max(1, int(0.001 * fused_t.numel())), (name, n_diff) + assert diff.max().item() <= 1.6e-2, (name, diff.max().item()) + + +_SHAPES = [ + # (lens, head_dim, ratio, coff) + pytest.param([2048], 128, 4, 2, id="b1-d128-r4"), + pytest.param([1023, 2048, 509], 128, 4, 2, id="ragged3-d128-r4"), + pytest.param([2048], 512, 4, 2, id="b1-d512-r4"), + pytest.param([3, 515, 1024, 129], 128, 4, 2, id="short-seg-d128-r4"), + # odd head_dim exercises the scalar (vec == 1) forward layout; even head_dims all + # take the vectorized (vec == 2) one. + pytest.param([260], 65, 4, 2, id="b1-d65-odd-r4"), + # a zero-length segment inside the pack (degenerate cu_seqlens entry). + pytest.param([64, 0, 253, 3], 128, 4, 2, id="empty-seg-d128-r4"), + # coff == 1: the non-overlapping window form (win = ratio, own-block tokens only). + pytest.param([2048], 128, 4, 1, id="b1-d128-r4-coff1"), + pytest.param([1023, 2048, 509], 128, 4, 1, id="ragged3-d128-r4-coff1"), + pytest.param([2048], 512, 4, 1, id="b1-d512-r4-coff1"), + pytest.param([3, 515, 1024, 129], 128, 4, 1, id="short-seg-d128-r4-coff1"), + pytest.param([260], 65, 4, 1, id="b1-d65-odd-r4-coff1"), + pytest.param([64, 0, 253, 3], 128, 4, 1, id="empty-seg-d128-r4-coff1"), + # ratio=128 (dedicated r128 kernels; coff {1, 2} x head_dim {128, 512}). The edge + # pack covers zero-block segments (127, 3), a literal empty segment, an + # exactly-one-block segment (128), a 1-token tail (129) and other tails. + pytest.param([8192], 128, 128, 1, id="b1x8192-d128-r128c1"), + pytest.param([1023, 2048, 509], 128, 128, 1, id="ragged3-d128-r128c1"), + pytest.param([127, 8192, 0, 129, 128, 3, 515, 1024], 128, 128, 1, id="edgepack-d128-r128c1"), + pytest.param([8192], 128, 128, 2, id="b1x8192-d128-r128c2"), + pytest.param([127, 8192, 0, 129, 128, 3, 515, 1024], 128, 128, 2, id="edgepack-d128-r128c2"), + pytest.param([2048, 509], 512, 128, 1, id="ragged2-d512-r128c1"), + pytest.param([8192], 512, 128, 2, id="b1x8192-d512-r128c2"), +] + + +# --------------------------------------------------------------------------- +# Numerics +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.parametrize("lens,d,ratio,coff", _SHAPES) +def test_numerics_vs_references(lens, d, ratio, coff): + """Fused fwd+bwd vs fp32-eager (bitwise dKV/dScore at ratio=4, tolerance at + ratio=128), upstream eager, and fp64 oracle.""" + _require_sm100() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs(lens, d, ratio, coff) + + r_fused = _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go) + r_fp32 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp32") + r_up = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="upstream") + r_fp64 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp64") + + # vs fp32-intermediate eager reference (the fused kernels' numerics contract): + # dKV / dScore bit-identical at ratio=4 / within the forward-style tolerances at + # ratio=128 (see _assert_grads_vs_fp32); forward within one bf16 rounding step on + # a tiny fraction of elements; dAPE within fp32 atomics reorder noise. + _assert_grads_vs_fp32(r_fused[1], r_fused[2], r_fp32[1], r_fp32[2], ratio) + fwd_diff = (r_fused[0].float() - r_fp32[0].float()).abs() + n_diff = (r_fused[0] != r_fp32[0]).sum().item() + assert n_diff <= max(1, int(0.001 * r_fused[0].numel())), n_diff + assert fwd_diff.max().item() <= 1.6e-2 + assert (r_fused[3] - r_fp32[3]).abs().max().item() <= 1e-3 + + # vs the verbatim upstream eager numerics: not bit-identical (the eager path rounds + # softmax weights to bf16 and multiplies in bf16), but close. + for fused_t, up_t in zip(r_fused, r_up): + assert torch.allclose(fused_t.float(), up_t.float(), rtol=0, atol=0.1) + + # vs the fp64 oracle. ratio=128: the contract's fp64-parity gate — per tensor the + # fused output must be at least as close to the oracle as the FP32-INTERMEDIATE + # eager reference, the comparator the contract names (comparing against the + # bf16-weight upstream path instead would be materially looser). ratio=4 keeps its + # historical check against the upstream numerics it replaced (its contract pins + # dKV/dScore bitwise-to-fp32-eager above and has no fp64-parity clause). + eager_ref = r_fp32 if ratio == 128 else r_up + for i in range(4): + err_fused = (r_fused[i].double() - r_fp64[i].double()).abs().max().item() + err_eager = (eager_ref[i].double() - r_fp64[i].double()).abs().max().item() + assert err_fused <= err_eager * (1 + 1e-6) + 1e-4, (i, err_fused, err_eager) + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_replay_determinism(coff): + """Forward, dKV and dScore replay bitwise identically run to run (dAPE is exempt).""" + _require_sm100() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 2048, 509], 128, 4, coff) + runs = [_run_fused(kv, score, ape, cu, cuc, total_comp, 4, 128, coff, go) for _ in range(3)] + for other in runs[1:]: + assert torch.equal(runs[0][0], other[0]) + assert torch.equal(runs[0][1], other[1]) + assert torch.equal(runs[0][2], other[2]) + # dAPE is accumulated with fp32 atomics; equality is not guaranteed, closeness is. + assert torch.allclose(runs[0][3], other[3], rtol=0, atol=1e-3) + + +@pytest.mark.L0 +def test_replay_determinism_r128(): + """ratio=128 mirror of test_replay_determinism (fixed chunk boundaries, no dKV/dScore atomics).""" + _require_sm100() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 2048, 509], 128, 128, 1) + runs = [_run_fused(kv, score, ape, cu, cuc, total_comp, 128, 128, 1, go) for _ in range(3)] + for other in runs[1:]: + assert torch.equal(runs[0][0], other[0]) + assert torch.equal(runs[0][1], other[1]) + assert torch.equal(runs[0][2], other[2]) + assert torch.allclose(runs[0][3], other[3], rtol=0, atol=1e-3) + + +# --------------------------------------------------------------------------- +# ratio=128 dispatch envelope: every shipped schedule bucket +# --------------------------------------------------------------------------- + + +def _r128_module(): + _import_compressor() + from cudnn.csa.compressor import compressor_sm100_r128 + + return compressor_sm100_r128 + + +# Expected shipped schedules per (coff, d, nb_total) bucket — hardcoded on purpose: an +# edit that silently changes any shipped launch geometry or bucket boundary must fail +# here. Forward tuples are (vec, tchunks, threads_x, twophase, fastexp); backward +# tuples are (vec, tchunks, threads_x, fastexp). +_FWD_BUCKETS = [ + (1, 128, 64, (2, 8, 32, False, False)), # small + (1, 128, 256, (2, 4, 32, False, False)), # default + (1, 128, 1024, (4, 4, 32, True, True)), # large: two-phase + fast exp + (2, 128, 64, (2, 16, 32, False, True)), # small (fast exp) + (2, 128, 256, (2, 8, 32, False, True)), # default (fast exp; no large entry) + (1, 512, 256, (2, 4, 32, False, False)), # default (no small entry) + (1, 512, 1024, (4, 4, 32, False, True)), # large + (2, 512, 64, (2, 8, 32, False, True)), # small + (2, 512, 256, (4, 4, 32, False, False)), # default + (2, 512, 1024, (4, 4, 32, False, True)), # large +] +_BWD_BUCKETS = [ + (1, 128, 64, (1, 8, 32, False)), # small pack: vec=1, exact exp + (1, 128, 256, (2, 8, 32, True)), # default (fast exp) + (2, 128, 64, (1, 8, 32, False)), # small + (2, 128, 256, (2, 8, 32, True)), # default + (1, 512, 256, (2, 4, 32, True)), # default (tchunks=4 at coff=1, d>=512) + (2, 512, 64, (2, 8, 32, True)), # no bwd small entry at d=512 -> default + (2, 512, 256, (2, 8, 32, True)), # default +] + + +@pytest.mark.L0 +def test_r128_dispatch_boundaries(): + """Schedule selection flips exactly at the documented nb_total bucket boundaries.""" + M = _r128_module() + for coff, d in ((1, 128), (2, 128), (1, 512), (2, 512)): + small = M._fwd_schedule_r128(128, d, coff, M._SMALL_NB_MAX) + first_default = M._fwd_schedule_r128(128, d, coff, M._SMALL_NB_MAX + 1) + last_default = M._fwd_schedule_r128(128, d, coff, M._LARGE_NB_MIN - 1) + large = M._fwd_schedule_r128(128, d, coff, M._LARGE_NB_MIN) + assert first_default == last_default, (coff, d) # one default bucket in between + if (coff, d) in M._SMALL_SCHEDULES: + assert small != first_default, (coff, d) + else: + assert small == first_default, (coff, d) + if (coff, d) in M._LARGE_SCHEDULES: + assert large != last_default, (coff, d) + else: + assert large == last_default, (coff, d) + bwd_small = M._bwd_schedule_r128(128, d, coff, M._BWD_SMALL_NB_MAX) + bwd_default = M._bwd_schedule_r128(128, d, coff, M._BWD_SMALL_NB_MAX + 1) + if (coff, d) in M._BWD_SMALL_SCHEDULES: + assert bwd_small != bwd_default, (coff, d) + else: + assert bwd_small == bwd_default, (coff, d) + + +@pytest.mark.L0 +def test_r128_dispatch_expected_schedules(): + """Every shipped (config, nb_total bucket) selects exactly the audited schedule.""" + M = _r128_module() + for coff, d, nb, expected in _FWD_BUCKETS: + assert M._fwd_schedule_r128(128, d, coff, nb) == expected, ("fwd", coff, d, nb) + for coff, d, nb, expected in _BWD_BUCKETS: + assert M._bwd_schedule_r128(128, d, coff, nb) == expected, ("bwd", coff, d, nb) + + +_ENVELOPE_CASES = [ + # (coff, d, nb_rows): one case per unique shipped (config, schedule) kernel, fwd + # and bwd together — the 10 forward + 6 backward shipped kernels across 10 cases + # (single-segment packs of nb_rows * ratio tokens land in the intended buckets). + pytest.param(1, 128, 64, id="c1d128-small"), + pytest.param(1, 128, 256, id="c1d128-default"), + pytest.param(1, 128, 1024, id="c1d128-large"), + pytest.param(2, 128, 64, id="c2d128-small"), + pytest.param(2, 128, 256, id="c2d128-default"), + pytest.param(1, 512, 256, id="c1d512-default"), + pytest.param(1, 512, 1024, id="c1d512-large"), + pytest.param(2, 512, 64, id="c2d512-small"), + pytest.param(2, 512, 256, id="c2d512-default"), + pytest.param(2, 512, 1024, id="c2d512-large"), +] + + +@pytest.mark.L1 +@pytest.mark.parametrize("coff,d,nb_rows", _ENVELOPE_CASES) +def test_r128_envelope_execution(coff, d, nb_rows): + """Execute every shipped schedule bucket once against the full contract: tolerance + vs the fp32 eager reference, fp64-oracle parity, and 2-run bitwise determinism — + covering the fast-exp, two-phase and vec=1 variants the (smaller) L0 shapes never + reach.""" + _require_sm100() + M = _r128_module() + ratio = 128 + # The shape must land in the intended bucket for BOTH directions (the wrappers + # select schedules through these exact functions). + exp_fwd = {(c, dd, n): e for c, dd, n, e in _FWD_BUCKETS}[(coff, d, nb_rows)] + assert M._fwd_schedule_r128(ratio, d, coff, nb_rows) == exp_fwd + exp_bwd = {(c, dd, n): e for c, dd, n, e in _BWD_BUCKETS}.get((coff, d, nb_rows)) + if exp_bwd is not None: + assert M._bwd_schedule_r128(ratio, d, coff, nb_rows) == exp_bwd + + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([nb_rows * ratio], d, ratio, coff) + assert total_comp == nb_rows + r1 = _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go) + r2 = _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go) + # Bitwise run-to-run determinism holds in EVERY bucket (incl. fast exp/two-phase). + for i in range(3): + assert torch.equal(r1[i], r2[i]), i + assert torch.allclose(r1[3], r2[3], rtol=0, atol=1e-3) + + r_fp32 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp32") + _assert_grads_vs_fp32(r1[1], r1[2], r_fp32[1], r_fp32[2], ratio) + assert (r1[0] != r_fp32[0]).sum().item() <= max(1, int(0.001 * r1[0].numel())) + assert (r1[0].float() - r_fp32[0].float()).abs().max().item() <= 1.6e-2 + assert (r1[3] - r_fp32[3]).abs().max().item() <= 1e-3 + + # fp64-oracle parity in EVERY shipped bucket (the contract's accuracy clause on + # finite-intermediate inputs), so the large-bucket schedules are held to the full + # contract here too. + r_fp64 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp64") + for i in range(4): + err_fused = (r1[i].double() - r_fp64[i].double()).abs().max().item() + err_eager = (r_fp32[i].double() - r_fp64[i].double()).abs().max().item() + assert err_fused <= err_eager * (1 + 1e-6) + 1e-4, (i, err_fused, err_eager) + + +_PADDING_SHAPES = [ + # (lens, head_dim, ratio, coff, pad); the first case has a LEADING segment shorter + # than ratio (0 compressed blocks), so padding rows gather tokens [0, ratio) that + # span a segment boundary -- exactly like the eager gather. + ([3, 515, 1024, 129], 128, 4, 2, 8), + ([1023, 2048, 509], 128, 4, 2, 8), + ([3, 515, 1024, 129], 128, 4, 1, 8), + ([1023, 2048, 509], 128, 4, 1, 8), +] + + +@pytest.mark.L0 +@pytest.mark.parametrize("lens,d,ratio,coff,pad", _PADDING_SHAPES) +def test_static_capacity_padding(lens, d, ratio, coff, pad): + """Static-capacity padding rows: eager-matching forward, ignored padding gradients.""" + _require_sm100() + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + capacity = total_true + pad + gen = torch.Generator(device="cpu").manual_seed(7) + go = torch.randn(capacity, 1, d, generator=gen, dtype=torch.float32) + go = go.to(torch.bfloat16).cuda() + go_zero_pad = go.clone() + go_zero_pad[total_true:] = 0 + + r_fused = _run_fused(kv, score, ape, cu, cuc, capacity, ratio, d, coff, go) + r_fp32 = _run_eager(kv, score, ape, cu, cuc, capacity, ratio, d, coff, go_zero_pad, mode="fp32") + + # Forward: padding rows replicate row 0's window exactly like the eager code, so the + # full padded output (valid + padding rows) obeys the same criteria as the unpadded + # comparison. + assert (r_fused[0] != r_fp32[0]).sum().item() <= max(1, int(0.001 * r_fused[0].numel())) + assert (r_fused[0].float() - r_fp32[0].float()).abs().max().item() <= 1.6e-2 + + # Backward: incoming gradients on padding rows are ignored by design -- the fused + # gradients (computed with NONZERO padding-row grads) match the eager reference run + # with zeroed padding-row grads bit-for-bit on dKV/dScore. + assert torch.equal(r_fused[1], r_fp32[1]) + assert torch.equal(r_fused[2], r_fp32[2]) + assert (r_fused[3] - r_fp32[3]).abs().max().item() <= 1e-3 + + # And explicitly: nonzero vs zero padding-row grads produce identical fused grads. + r_fused_zero = _run_fused(kv, score, ape, cu, cuc, capacity, ratio, d, coff, go_zero_pad) + assert torch.equal(r_fused[1], r_fused_zero[1]) + assert torch.equal(r_fused[2], r_fused_zero[2]) + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_empty_output(coff): + """total_comp == 0 launches nothing and returns well-formed empty/zero tensors.""" + _require_sm100() + compressor = _import_compressor() + d, ratio = 128, 4 + w = coff * d + kv = torch.randn(2, w, device="cuda").to(torch.bfloat16) + score = torch.randn(2, w, device="cuda").to(torch.bfloat16) + ape = torch.randn(ratio, w, device="cuda") + cu = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, 0], dtype=torch.int32, device="cuda") + out = compressor.csa_compressor_forward_wrapper(kv, score, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=0)["out"] + assert out.shape == (0, d) and out.dtype == torch.bfloat16 + grads = compressor.csa_compressor_backward_wrapper(kv, score, ape, cu, cuc, out, ratio=ratio, head_dim=d, coff=coff) + assert grads["grad_kv"].abs().sum().item() == 0 + assert grads["grad_score"].abs().sum().item() == 0 + assert grads["grad_ape"].abs().sum().item() == 0 + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_backward_wrapper_zeros_when_no_blocks(coff): + """Multi-segment packs where NO segment reaches ratio tokens: total_comp == 0, so + the kernel cannot launch and the wrapper must hand back exact-zero grads (the + host-side zeros fallback behind the uninitialized-buffer optimization).""" + _require_sm100() + compressor = _import_compressor() + d, ratio = 128, 4 + w = coff * d + lens = [3, 2, 1] + total = sum(lens) + kv = torch.randn(total, w, device="cuda").to(torch.bfloat16) + score = torch.randn(total, w, device="cuda").to(torch.bfloat16) + ape = torch.randn(ratio, w, device="cuda") + cu = torch.tensor([0, 3, 5, 6], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, 0, 0, 0], dtype=torch.int32, device="cuda") + go = torch.empty(0, d, dtype=torch.bfloat16, device="cuda") + grads = compressor.csa_compressor_backward_wrapper(kv, score, ape, cu, cuc, go, ratio=ratio, head_dim=d, coff=coff) + assert grads["grad_kv"].shape == kv.shape and grads["grad_kv"].abs().sum().item() == 0 + assert grads["grad_score"].shape == score.shape and grads["grad_score"].abs().sum().item() == 0 + assert grads["grad_ape"].abs().sum().item() == 0 + + +def _never_consumed_mask(lens, total_tokens, d, ratio, coff): + """Boolean ``(total_tokens, coff * d)`` mask of dKV/dScore slots no output row reads. + + Mirrors the kernel's zero-write ownership classes exactly: per-segment tail tokens + (``seqlen % ratio``, all columns), whole segments shorter than ``ratio``, for + ``coff == 2`` the first-half columns of each segment's LAST block's own tokens (no + next in-segment block consumes them), and token-capacity padding beyond + ``cu_seqlens[-1]``. + """ + w = coff * d + mask = torch.zeros(total_tokens, w, dtype=torch.bool) + pos = 0 + for seg_len in lens: + nb = seg_len // ratio + if nb == 0: + mask[pos : pos + seg_len] = True + else: + mask[pos + nb * ratio : pos + seg_len] = True + if coff == 2: + mask[pos + (nb - 1) * ratio : pos + nb * ratio, :d] = True + pos += seg_len + mask[pos:] = True + return mask + + +_CANARY_SHAPES = [ + # (lens, head_dim, ratio, coff, pad, tok_pad) — every never-consumed dKV/dScore + # slot class must be hit: segment tails (seqlen % ratio), the last block's + # first-half columns (coff == 2 only), whole segments shorter than ratio (zero + # blocks), static-capacity padding rows (pad > 0 extra grad_out rows), and static + # token-capacity padding of the gradient buffers themselves (tok_pad > 0 tokens + # beyond cu_seqlens[-1]). + pytest.param([2048], 128, 4, 2, 0, 0, id="b1-d128"), + pytest.param([1023, 2048, 509], 128, 4, 2, 0, 0, id="ragged3-d128"), + pytest.param([3, 515, 1024, 129], 128, 4, 2, 0, 0, id="short-seg-d128"), + pytest.param([5, 6, 7], 128, 4, 2, 0, 0, id="all-tiny-d128"), + pytest.param([1023, 2048, 509], 512, 4, 2, 0, 0, id="ragged3-d512"), + pytest.param([3, 515, 1024, 129], 128, 4, 2, 8, 0, id="short-seg-d128-padded"), + pytest.param([1023, 2048, 509], 128, 4, 2, 8, 0, id="ragged3-d128-padded"), + pytest.param([1023, 2048, 509], 128, 4, 2, 0, 37, id="ragged3-d128-tokpad"), + pytest.param([3, 515, 1024, 129], 128, 4, 2, 8, 21, id="short-seg-d128-padded-tokpad"), + # ratio=4, coff == 1 (own-block window): same shapes as the coff=2 rows above; the + # zero classes drop the first-half-column class (no overlap halves at coff=1). + pytest.param([2048], 128, 4, 1, 0, 0, id="b1-d128-coff1"), + pytest.param([1023, 2048, 509], 128, 4, 1, 0, 0, id="ragged3-d128-coff1"), + pytest.param([3, 515, 1024, 129], 128, 4, 1, 0, 0, id="short-seg-d128-coff1"), + pytest.param([5, 6, 7], 128, 4, 1, 0, 0, id="all-tiny-d128-coff1"), + pytest.param([1023, 2048, 509], 512, 4, 1, 0, 0, id="ragged3-d512-coff1"), + pytest.param([3, 515, 1024, 129], 128, 4, 1, 8, 0, id="short-seg-d128-padded-coff1"), + pytest.param([1023, 2048, 509], 128, 4, 1, 8, 0, id="ragged3-d128-padded-coff1"), + pytest.param([1023, 2048, 509], 128, 4, 1, 0, 37, id="ragged3-d128-tokpad-coff1"), + pytest.param([3, 515, 1024, 129], 128, 4, 1, 8, 21, id="short-seg-d128-padded-tokpad-coff1"), + # ratio=128: the zero classes are up to 127 tokens each (tails, zero-block + # segments) plus the coff=2 last-block first-half (128 rows). + pytest.param([1023, 2048, 509], 128, 128, 1, 0, 0, id="ragged3-d128-r128c1"), + pytest.param([127, 8192, 0, 129, 128, 3, 515, 1024], 128, 128, 1, 8, 21, id="edgepack-d128-r128c1-padded-tokpad"), + pytest.param([127, 8192, 0, 129, 128, 3, 515, 1024], 128, 128, 2, 8, 21, id="edgepack-d128-r128c2-padded-tokpad"), + pytest.param([2048, 509], 512, 128, 2, 8, 21, id="ragged2-d512-r128c2-padded-tokpad"), +] + + +@pytest.mark.L0 +@pytest.mark.parametrize("lens,d,ratio,coff,pad,tok_pad", _CANARY_SHAPES) +def test_backward_fills_uninitialized_buffers(lens, d, ratio, coff, pad, tok_pad): + """NaN-canary: the backward kernel fully overwrites garbage dKV/dScore buffers. + + The kernel writes exact zeros to every never-consumed slot itself (there are no + separate zero-fill kernels anymore), so running it into NaN-poisoned buffers must + produce bitwise the same dKV/dScore as running it into zero-initialized buffers. + """ + _require_sm100() + compressor = _import_compressor() + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + total_comp = total_true + pad + gen = torch.Generator(device="cpu").manual_seed(11) + go = torch.randn(total_comp, d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda() + total = kv.shape[0] + tok_pad + kv2 = torch.cat([kv.view(kv.shape[0], -1), torch.randn(tok_pad, coff * d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda()]) + score2 = torch.cat([score.view(score.shape[0], -1), torch.randn(tok_pad, coff * d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda()]) + + bwd = compressor.CSACompressorBackward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert bwd.check_support() + bwd.compile() + + def run(poison): + """One backward into poisoned (NaN) or zeroed grad buffers; returns the grads.""" + grad_kv = torch.empty_like(kv2) + grad_score = torch.empty_like(score2) + if poison: + grad_kv.fill_(float("nan")) + grad_score.fill_(float("nan")) + else: + grad_kv.zero_() + grad_score.zero_() + grad_ape = torch.zeros_like(ape) + bwd.execute(kv2, score2, ape, cu, cuc, go, grad_kv, grad_score, grad_ape) + torch.cuda.synchronize() + return grad_kv, grad_score, grad_ape + + gkv_ref, gs_ref, gape_ref = run(poison=False) + gkv_nan, gs_nan, gape_nan = run(poison=True) + assert not torch.isnan(gkv_nan).any(), "unwritten dKV slots survived (NaN canary)" + assert not torch.isnan(gs_nan).any(), "unwritten dScore slots survived (NaN canary)" + assert torch.equal(gkv_nan, gkv_ref) + assert torch.equal(gs_nan, gs_ref) + assert torch.allclose(gape_nan, gape_ref, rtol=0, atol=1e-3) + + # And the zero-slot pattern matches autograd: never-consumed slots are exact zeros, + # exactly as the fp32 eager reference computes them (bitwise at ratio=4, tolerance + # at ratio=128 — the zero slots themselves are exact in both). (The fused backward + # ignores incoming gradients on static-capacity padding rows by design, so the + # eager reference runs with those rows zeroed.) + go_ref = go.clone() + go_ref[total_true:] = 0 + r_fp32 = _run_eager( + kv2.view(total, 1, -1), + score2.view(total, 1, -1), + ape, + cu, + cuc, + total_comp, + ratio, + d, + coff, + go_ref.view(total_comp, 1, d), + mode="fp32", + ) + _assert_grads_vs_fp32(gkv_nan.view_as(r_fp32[1]), gs_nan.view_as(r_fp32[2]), r_fp32[1], r_fp32[2], ratio) + + # The zero-owned classes themselves must be EXACT zeros, not merely within the + # ratio=128 tolerance (a deterministic small nonzero written into a never-consumed + # slot would otherwise pass). The mask is sanity-checked against the eager + # reference first: autograd produces exact zeros on precisely these slots. + mask = _never_consumed_mask(lens, total, d, ratio, coff).cuda() + assert (r_fp32[1].view(total, -1)[mask] == 0).all(), "mask does not match the eager zero slots (dKV)" + assert (r_fp32[2].view(total, -1)[mask] == 0).all(), "mask does not match the eager zero slots (dScore)" + assert (gkv_nan.view(total, -1)[mask] == 0).all(), "never-consumed dKV slots must be exact zeros" + assert (gs_nan.view(total, -1)[mask] == 0).all(), "never-consumed dScore slots must be exact zeros" + + +# --------------------------------------------------------------------------- +# Deterministic mode +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_backward_rejects_deterministic_mode(coff): + """The backward raises under torch.use_deterministic_algorithms (dAPE fp32 atomics).""" + _require_sm100() + compressor = _import_compressor() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([512, 256], 128, 4, coff) + total = kv.shape[0] + # Forward is deterministic and keeps working. + prev_det = torch.are_deterministic_algorithms_enabled() + prev_warn = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True, warn_only=False) + try: + out = compressor.csa_compressor_forward_wrapper( + kv.view(total, -1), score.view(total, -1), ape, cu, cuc, ratio=4, head_dim=128, coff=coff, total_comp=total_comp + )["out"] + assert out.shape == (total_comp, 128) + with pytest.raises(RuntimeError, match="not deterministic"): + compressor.csa_compressor_backward_wrapper( + kv.view(total, -1), score.view(total, -1), ape, cu, cuc, go.view(total_comp, 128), ratio=4, head_dim=128, coff=coff + ) + finally: + torch.use_deterministic_algorithms(prev_det, warn_only=prev_warn) + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_backward_warns_in_warn_only_deterministic_mode(coff): + """warn_only deterministic mode warns (torch semantics) and still runs the backward.""" + _require_sm100() + compressor = _import_compressor() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([512, 256], 128, 4, coff) + r_ref = _run_fused(kv, score, ape, cu, cuc, total_comp, 4, 128, coff, go) + prev_det = torch.are_deterministic_algorithms_enabled() + prev_warn = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True, warn_only=True) + try: + with pytest.warns(RuntimeWarning, match="not deterministic") as record: + grads = compressor.csa_compressor_backward_wrapper( + kv.view(kv.shape[0], -1), score.view(score.shape[0], -1), ape, cu, cuc, go.view(total_comp, 128), ratio=4, head_dim=128, coff=coff + ) + finally: + torch.use_deterministic_algorithms(prev_det, warn_only=prev_warn) + # Exactly ONE warning on the wrapper path: execute() is the single check point + # (the wrapper does not duplicate it). + det_warnings = [w for w in record if issubclass(w.category, RuntimeWarning) and "not deterministic" in str(w.message)] + assert len(det_warnings) == 1, f"expected exactly one deterministic-mode warning, got {len(det_warnings)}" + assert torch.equal(grads["grad_kv"].view_as(r_ref[1]), r_ref[1]) + assert torch.equal(grads["grad_score"].view_as(r_ref[2]), r_ref[2]) + + +# --------------------------------------------------------------------------- +# CUDA graph capture +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.filterwarnings("ignore::UserWarning") +@pytest.mark.parametrize("coff", [1, 2]) +def test_cuda_graph_capture(coff): + """Warmup -> capture fwd+bwd -> replay; JIT under capture raises a clear error.""" + _require_sm100() + compressor = _import_compressor() + ratio, d = 4, 128 + lens = [512, 256] + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + capacity = total_true + 8 # static capacity, as with CUDA-graph static shapes + total = kv.shape[0] + + kv_s = kv.view(total, -1).clone() + score_s = score.view(total, -1).clone() + ape_s = ape.clone() + go_s = torch.zeros(capacity, d, device="cuda", dtype=torch.bfloat16) + go_s[:total_true] = torch.randn(total_true, d, device="cuda").to(torch.bfloat16) + + def _fused_fwd_bwd(): + """One fused forward + backward over the static-capacity buffers.""" + out = compressor.csa_compressor_forward_wrapper(kv_s, score_s, ape_s, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=capacity)["out"] + grads = compressor.csa_compressor_backward_wrapper(kv_s, score_s, ape_s, cu, cuc, go_s, ratio=ratio, head_dim=d, coff=coff) + return out, grads + + # One warmup per configuration on a side stream (JIT-compiles both kernels). + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _fused_fwd_bwd() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_c, grads_c = _fused_fwd_bwd() + gkv_c, gscore_c, gape_c = grads_c["grad_kv"], grads_c["grad_score"], grads_c["grad_ape"] + + # Replay on the same data must reproduce the direct (non-captured) fused results + # bitwise on forward/dKV/dScore. + graph.replay() + torch.cuda.synchronize() + ref_out, ref_grads = _fused_fwd_bwd() + torch.cuda.synchronize() + assert torch.equal(out_c, ref_out) + assert torch.equal(gkv_c, ref_grads["grad_kv"]) + assert torch.equal(gscore_c, ref_grads["grad_score"]) + assert torch.allclose(gape_c, ref_grads["grad_ape"], rtol=0, atol=1e-3) + + # Replay with new data and a SMALLER device-side true row count (the fixed capacity + # stays static, cu/cuc contents change) -- the graph-replayed gradients must match + # the fp32 eager reference bitwise on dKV/dScore. + lens2 = [384, 128] + kv2, score2, ape2, cu2, cuc2, total2, _ = _make_inputs(lens2, d, ratio, coff, seed=99) + n2 = sum(lens2) + kv_s.zero_() + score_s.zero_() + kv_s[:n2] = kv2.view(n2, -1) + score_s[:n2] = score2.view(n2, -1) + ape_s.copy_(ape2) + cu.copy_(cu2) + cuc.copy_(cuc2) + go_s.zero_() + go_s[:total2] = torch.randn(total2, d, device="cuda").to(torch.bfloat16) + graph.replay() + torch.cuda.synchronize() + r_fp32 = _run_eager( + kv_s.view(-1, 1, coff * d), + score_s.view(-1, 1, coff * d), + ape_s, + cu, + cuc, + capacity, + ratio, + d, + coff, + go_s.view(capacity, 1, d), + mode="fp32", + ) + assert torch.equal(gkv_c.view_as(r_fp32[1]), r_fp32[1]) + assert torch.equal(gscore_c.view_as(r_fp32[2]), r_fp32[2]) + # The narrowed replay is checked on ALL FOUR outputs, not just dKV/dScore: + # forward vs eager (contract tolerance), dAPE vs eager (atomics tolerance), and + # the whole replay vs a direct (non-captured) fused call on the same mutated + # inputs — bitwise on forward/dKV/dScore. + assert (out_c.view_as(r_fp32[0]) != r_fp32[0]).sum().item() <= max(1, int(0.001 * r_fp32[0].numel())) + assert (out_c.view_as(r_fp32[0]).float() - r_fp32[0].float()).abs().max().item() <= 1.6e-2 + assert (gape_c - r_fp32[3]).abs().max().item() <= 1e-3 + direct_out, direct_grads = _fused_fwd_bwd() + torch.cuda.synchronize() + assert torch.equal(out_c, direct_out) + assert torch.equal(gkv_c, direct_grads["grad_kv"]) + assert torch.equal(gscore_c, direct_grads["grad_score"]) + + # A first call for a NEW configuration under capture must raise loudly instead of + # JIT-compiling (which is not capture-safe). head_dim 192 is used by no other test + # in this module (and kernels are compiled per (ratio, head_dim, coff)), so this + # configuration is guaranteed to be uncompiled regardless of test execution order. + d_new = 192 + kv3 = torch.randn(256, coff * d_new, device="cuda").to(torch.bfloat16) + score3 = torch.randn(256, coff * d_new, device="cuda").to(torch.bfloat16) + ape3 = torch.randn(ratio, coff * d_new, device="cuda") + cu3 = torch.tensor([0, 256], dtype=torch.int32, device="cuda") + cuc3 = torch.tensor([0, 64], dtype=torch.int32, device="cuda") + graph2 = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="CUDA graph capture"): + with torch.cuda.graph(graph2): + compressor.csa_compressor_forward_wrapper(kv3, score3, ape3, cu3, cuc3, ratio=ratio, head_dim=d_new, coff=coff, total_comp=64) + # The CUDA context must remain usable after the aborted capture. + torch.cuda.synchronize() + probe = torch.ones(8, device="cuda") + assert probe.sum().item() == 8 + + +@pytest.mark.L0 +@pytest.mark.filterwarnings("ignore::UserWarning") +def test_cuda_graph_capture_r128(): + """ratio=128 mirror of test_cuda_graph_capture: warmup -> capture fwd+bwd -> replay. + + The ratio=128 backward's rows_per_cta launch parameter is derived from the static + row capacity on the host, so capture/replay must reproduce the direct call bitwise + on forward/dKV/dScore exactly like the ratio=4 path. + """ + _require_sm100() + compressor = _import_compressor() + ratio, d, coff = 128, 128, 1 + lens = [640, 259] # 5 + 2 blocks, 3-token tail + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + capacity = total_true + 8 + total = kv.shape[0] + + kv_s = kv.view(total, -1).clone() + score_s = score.view(total, -1).clone() + ape_s = ape.clone() + go_s = torch.zeros(capacity, d, device="cuda", dtype=torch.bfloat16) + go_s[:total_true] = torch.randn(total_true, d, device="cuda").to(torch.bfloat16) + + def _fused_fwd_bwd(): + out = compressor.csa_compressor_forward_wrapper(kv_s, score_s, ape_s, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=capacity)["out"] + grads = compressor.csa_compressor_backward_wrapper(kv_s, score_s, ape_s, cu, cuc, go_s, ratio=ratio, head_dim=d, coff=coff) + return out, grads + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _fused_fwd_bwd() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_c, grads_c = _fused_fwd_bwd() + gkv_c, gscore_c, gape_c = grads_c["grad_kv"], grads_c["grad_score"], grads_c["grad_ape"] + + graph.replay() + torch.cuda.synchronize() + ref_out, ref_grads = _fused_fwd_bwd() + torch.cuda.synchronize() + assert torch.equal(out_c, ref_out) + assert torch.equal(gkv_c, ref_grads["grad_kv"]) + assert torch.equal(gscore_c, ref_grads["grad_score"]) + assert torch.allclose(gape_c, ref_grads["grad_ape"], rtol=0, atol=1e-3) + + # Replay with new data and a SMALLER device-side true row count (the fixed capacity + # stays static, cu/cuc contents change) -- the graph-replayed gradients must match + # the fp32 eager reference within the ratio=128 tolerance contract (capture/replay + # itself stays bitwise vs the direct call, asserted above). + lens2 = [384, 128] + kv2, score2, ape2, cu2, cuc2, total2, _ = _make_inputs(lens2, d, ratio, coff, seed=99) + n2 = sum(lens2) + kv_s.zero_() + score_s.zero_() + kv_s[:n2] = kv2.view(n2, -1) + score_s[:n2] = score2.view(n2, -1) + ape_s.copy_(ape2) + cu.copy_(cu2) + cuc.copy_(cuc2) + go_s.zero_() + go_s[:total2] = torch.randn(total2, d, device="cuda").to(torch.bfloat16) + graph.replay() + torch.cuda.synchronize() + r_fp32 = _run_eager( + kv_s.view(-1, 1, coff * d), + score_s.view(-1, 1, coff * d), + ape_s, + cu, + cuc, + capacity, + ratio, + d, + coff, + go_s.view(capacity, 1, d), + mode="fp32", + ) + _assert_grads_vs_fp32(gkv_c.view_as(r_fp32[1]), gscore_c.view_as(r_fp32[2]), r_fp32[1], r_fp32[2], ratio) + # The narrowed replay is checked on ALL FOUR outputs, not just dKV/dScore: + # forward vs eager (contract tolerance), dAPE vs eager (atomics tolerance), and + # the whole replay vs a direct (non-captured) fused call on the same mutated + # inputs — bitwise on forward/dKV/dScore. + assert (out_c.view_as(r_fp32[0]) != r_fp32[0]).sum().item() <= max(1, int(0.001 * r_fp32[0].numel())) + assert (out_c.view_as(r_fp32[0]).float() - r_fp32[0].float()).abs().max().item() <= 1.6e-2 + assert (gape_c - r_fp32[3]).abs().max().item() <= 1e-3 + direct_out, direct_grads = _fused_fwd_bwd() + torch.cuda.synchronize() + assert torch.equal(out_c, direct_out) + assert torch.equal(gkv_c, direct_grads["grad_kv"]) + assert torch.equal(gscore_c, direct_grads["grad_score"]) + + +# --------------------------------------------------------------------------- +# check_support boundaries +# --------------------------------------------------------------------------- + + +def _meta(shape, dtype, stride=None): + """Metadata-only sample tensor (meta device) for check_support tests.""" + if stride is None: + stride = [] + acc = 1 + for s in reversed(shape): + stride.append(acc) + acc *= s + stride = tuple(reversed(stride)) + return torch.empty_strided(shape, stride, dtype=dtype, device="meta") + + +def _meta_samples( + total=512, + d=128, + ratio=4, + coff=2, + n_seg=2, + total_comp=None, + kv_dtype=torch.bfloat16, + ape_dtype=torch.float32, + cu_dtype=torch.int32, + out_dtype=torch.bfloat16, + score_shape=None, + cuc_len=None, + kv_stride=None, +): + """Consistent meta-device sample-tensor kwargs, with per-field overrides for negatives.""" + w = coff * d + if total_comp is None: + total_comp = total // ratio + kv = _meta((total, w), kv_dtype, stride=kv_stride) + score = _meta(score_shape or (total, w), kv_dtype) + ape = _meta((ratio, w), ape_dtype) + cu = _meta((n_seg + 1,), cu_dtype) + cuc = _meta((cuc_len or (n_seg + 1),), cu_dtype) + out = _meta((total_comp, d), out_dtype) + return dict( + sample_kv=kv, + sample_score=score, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=out, + ) + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_check_support_accepts_envelope(coff): + """Metadata-only samples inside the validated envelope pass check_support.""" + _require_sm100() + compressor = _import_compressor() + for cls in (compressor.CSACompressorForward, compressor.CSACompressorBackward): + api = cls(**_meta_samples(coff=coff), ratio=4, coff=coff) + assert api.check_support() is True + assert api.head_dim == 128 and api.total_tokens == 512 and api.total_comp == 128 + + +@pytest.mark.L0 +@pytest.mark.parametrize("d", [128, 512]) +@pytest.mark.parametrize("coff", [1, 2]) +def test_check_support_accepts_envelope_r128(d, coff): + """ratio=128 envelope: coff {1, 2} x head_dim {128, 512} pass check_support.""" + _require_sm100() + compressor = _import_compressor() + for cls in (compressor.CSACompressorForward, compressor.CSACompressorBackward): + api = cls(**_meta_samples(total=1024, d=d, ratio=128, coff=coff), ratio=128, coff=coff) + assert api.check_support() is True + assert api.head_dim == d and api.total_tokens == 1024 and api.total_comp == 8 + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs,ctor,match", + [ + (dict(), dict(ratio=128, coff=3), "ratio=128 supports coff"), + (dict(), dict(ratio=8, coff=2), "ratio in \\{4, 128\\}"), + (dict(coff=3), dict(ratio=4, coff=3), "coff in"), + (dict(), dict(ratio=4, coff=0), "coff in"), + (dict(kv_dtype=torch.float16), dict(), "kv"), + (dict(ape_dtype=torch.bfloat16), dict(), "ape"), + (dict(cu_dtype=torch.int64), dict(), "cu_seqlens"), + (dict(out_dtype=torch.float32), dict(), "out"), + (dict(score_shape=(512, 128)), dict(), "score shape"), + (dict(cuc_len=4), dict(), "B \\+ 1"), + (dict(total=2**25, d=128), dict(), "int32 flat offsets"), + (dict(total=4, d=8388482, total_comp=1), dict(), "head_dim"), + (dict(total=2, total_comp=1), dict(), "requires at least ratio"), + (dict(kv_stride=(512, 2)), dict(), "contiguous"), + (dict(total=1024, d=96, ratio=128, coff=1), dict(ratio=128, coff=1), "ratio=128 is validated for head_dim"), + ], +) +def test_check_support_rejects(kwargs, ctor, match): + """check_support raises ValueError for configurations outside the envelope.""" + _require_sm100() + compressor = _import_compressor() + samples = _meta_samples(**kwargs) + api = compressor.CSACompressorForward(**samples, ratio=ctor.get("ratio", 4), coff=ctor.get("coff", 2)) + with pytest.raises(ValueError, match=match): + api.check_support() + + +@pytest.mark.L0 +def test_check_support_rejects_cpu_tensors(): + """CPU sample tensors are rejected with a clear error.""" + _require_sm100() + compressor = _import_compressor() + samples = _meta_samples() + samples["sample_kv"] = torch.empty(512, 256, dtype=torch.bfloat16, device="cpu") + api = compressor.CSACompressorForward(**samples, ratio=4, coff=2) + with pytest.raises(ValueError, match="CUDA"): + api.check_support() + + +# --------------------------------------------------------------------------- +# Class API vs wrapper equivalence +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.parametrize("coff", [1, 2]) +def test_class_api_matches_wrapper(coff): + """The explicit class API produces bitwise-identical results to the wrappers.""" + _require_sm100() + compressor = _import_compressor() + ratio, d = 4, 128 + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 509], d, ratio, coff) + total = kv.shape[0] + kv2, score2, go2 = kv.view(total, -1), score.view(total, -1), go.view(total_comp, d) + + r_wrapped = _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go) + + fwd = compressor.CSACompressorForward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert fwd.check_support() + fwd.compile() + out = torch.empty(total_comp, d, dtype=torch.bfloat16, device="cuda") + fwd.execute(kv2, score2, ape, cu, cuc, out) + + bwd = compressor.CSACompressorBackward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert bwd.check_support() + bwd.compile() + grad_kv = torch.zeros_like(kv2) + grad_score = torch.zeros_like(score2) + grad_ape = torch.zeros_like(ape) + bwd.execute(kv2, score2, ape, cu, cuc, go2, grad_kv, grad_score, grad_ape) + torch.cuda.synchronize() + + assert torch.equal(out.view_as(r_wrapped[0]), r_wrapped[0]) + assert torch.equal(grad_kv.view_as(r_wrapped[1]), r_wrapped[1]) + assert torch.equal(grad_score.view_as(r_wrapped[2]), r_wrapped[2]) + assert torch.allclose(grad_ape, r_wrapped[3], rtol=0, atol=1e-3) + + +@pytest.mark.L0 +@pytest.mark.filterwarnings("ignore::UserWarning") +def test_backward_grad_ape_zeroing_contract(): + """grad_ape ownership: the kernel only ACCUMULATES, so the class-API caller must + re-zero grad_ape before every execute and before every CUDA-graph replay that + reuses the buffer. The high-level wrapper allocates a fresh zeroed buffer per call + (and its zero-fill is captured with the kernel, so wrapper replays re-zero).""" + _require_sm100() + compressor = _import_compressor() + ratio, d, coff = 128, 128, 1 + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([640, 259], d, ratio, coff) + total = kv.shape[0] + kv2, score2, go2 = kv.view(total, -1), score.view(total, -1), go.view(total_comp, d) + + bwd = compressor.CSACompressorBackward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert bwd.check_support() + bwd.compile() + grad_kv = torch.empty_like(kv2) + grad_score = torch.empty_like(score2) + + # Single-run reference (also warms the launch path up for the capture below). + ref = torch.zeros_like(ape) + bwd.execute(kv2, score2, ape, cu, cuc, go2, grad_kv, grad_score, ref) + torch.cuda.synchronize() + + # Two executes WITHOUT re-zeroing accumulate (the documented sharp edge). + acc = torch.zeros_like(ape) + bwd.execute(kv2, score2, ape, cu, cuc, go2, grad_kv, grad_score, acc) + bwd.execute(kv2, score2, ape, cu, cuc, go2, grad_kv, grad_score, acc) + torch.cuda.synchronize() + assert torch.allclose(acc, 2 * ref, rtol=0, atol=2e-3) + + # Graph replays of a captured class-API execute accumulate the same way: the + # zero-fill happened BEFORE capture, so it is not part of the graph. + gape_graph = torch.zeros_like(ape) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + bwd.execute(kv2, score2, ape, cu, cuc, go2, grad_kv, grad_score, gape_graph) + graph.replay() + graph.replay() + torch.cuda.synchronize() + assert torch.allclose(gape_graph, 2 * ref, rtol=0, atol=2e-3) + + # The wrapper path does NOT accumulate: fresh zeroed buffer per call. + g1 = compressor.csa_compressor_backward_wrapper(kv2, score2, ape, cu, cuc, go2, ratio=ratio, head_dim=d, coff=coff) + g2 = compressor.csa_compressor_backward_wrapper(kv2, score2, ape, cu, cuc, go2, ratio=ratio, head_dim=d, coff=coff) + torch.cuda.synchronize() + assert torch.allclose(g1["grad_ape"], ref, rtol=0, atol=1e-3) + assert torch.allclose(g2["grad_ape"], ref, rtol=0, atol=1e-3) + + +# --------------------------------------------------------------------------- +# Runtime hazards: alignment, explicit streams, multi-device +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_misaligned_base_pointer_rejected(): + """Contiguous storage-offset views with unaligned base pointers raise ValueError.""" + _require_sm100() + compressor = _import_compressor() + d, ratio, coff = 128, 4, 2 + w = coff * d + total = 512 + buf = torch.randn(total * w + 8, device="cuda").to(torch.bfloat16) + kv = buf[2 : 2 + total * w].view(total, w) # contiguous, 4-byte-aligned base + assert kv.is_contiguous() and kv.data_ptr() % 16 != 0 + score = torch.randn(total, w, device="cuda").to(torch.bfloat16) + ape = torch.randn(ratio, w, device="cuda") + cu = torch.tensor([0, total], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, total // ratio], dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="aligned"): + compressor.csa_compressor_forward_wrapper(kv, score, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total // ratio) + + +@pytest.mark.L0 +def test_explicit_stream_input_lifetime(): + """Inputs released right after an explicit-stream call are not recycled early. + + The launch path takes raw pointers, so the wrapper must ``record_stream`` its + operands on the external stream; otherwise the caching allocator can hand the freed + input storage to allocations on another stream while the kernel is still pending. + """ + _require_sm100() + compressor = _import_compressor() + if not hasattr(torch.cuda, "_sleep"): + pytest.skip("torch.cuda._sleep not available") + import cuda.bindings.driver as cuda_driver + + d, ratio, coff = 128, 4, 2 + w = coff * d + lens = [4096] + kv, score, ape, cu, cuc, total_comp, _ = _make_inputs(lens, d, ratio, coff) + total = kv.shape[0] + kv2, score2 = kv.view(total, w).contiguous(), score.view(total, w).contiguous() + + # Ground truth on the default stream (from private clones). + expected = compressor.csa_compressor_forward_wrapper( + kv2.clone(), score2.clone(), ape.clone(), cu.clone(), cuc.clone(), ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp + )["out"] + torch.cuda.synchronize() + + side = torch.cuda.Stream() + ext = cuda_driver.CUstream(side.cuda_stream) + with torch.cuda.stream(side): + torch.cuda._sleep(int(5e8)) # block the side stream so the kernel stays pending + out = compressor.csa_compressor_forward_wrapper(kv2, score2, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp, stream=ext)["out"] + # Drop every caller reference while the kernel is still queued behind the sleep, + # then try hard to get the freed storages reallocated and scribbled on the default + # (idle) stream. + del kv, score, kv2, score2, ape, cu, cuc + junk = [torch.full((total, w), 7.0, device="cuda", dtype=torch.bfloat16) for _ in range(4)] + junk.append(torch.full((ratio, w), 7.0, device="cuda", dtype=torch.float32)) + junk.append(torch.full((64,), 7, device="cuda", dtype=torch.int32)) + torch.cuda.synchronize() + assert torch.equal(out, expected), "explicit-stream inputs were recycled before the kernel consumed them" + + +@pytest.mark.L0 +def test_multi_device_launch(): + """Tensors on a non-current device produce correct results (device anchoring).""" + _require_sm100() + compressor = _import_compressor() + if torch.cuda.device_count() < 2: + pytest.skip("needs >= 2 visible GPUs") + if torch.cuda.get_device_capability(1) != (10, 0): + pytest.skip("second GPU is not CC 10.0") + d, ratio, coff = 128, 4, 2 + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1024], d, ratio, coff, device="cuda:1") + total = kv.shape[0] + kv2, score2, go2 = kv.view(total, -1), score.view(total, -1), go.view(total_comp, d) + + assert torch.cuda.current_device() == 0 # launch with a FOREIGN current device + out_foreign = compressor.csa_compressor_forward_wrapper(kv2, score2, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp)["out"] + grads_foreign = compressor.csa_compressor_backward_wrapper(kv2, score2, ape, cu, cuc, go2, ratio=ratio, head_dim=d, coff=coff) + torch.cuda.synchronize(torch.device("cuda", 1)) + + with torch.cuda.device(1): + out_native = compressor.csa_compressor_forward_wrapper(kv2, score2, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp)["out"] + grads_native = compressor.csa_compressor_backward_wrapper(kv2, score2, ape, cu, cuc, go2, ratio=ratio, head_dim=d, coff=coff) + torch.cuda.synchronize() + + assert out_foreign.device == torch.device("cuda", 1) + assert out_foreign.abs().sum().item() > 0 # the historic failure mode was all-zeros + assert torch.equal(out_foreign, out_native) + assert torch.equal(grads_foreign["grad_kv"], grads_native["grad_kv"]) + assert torch.equal(grads_foreign["grad_score"], grads_native["grad_score"])