From e43059105499b59e1d9c4747b7fdb5d568619368 Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Wed, 2 Sep 2026 23:16:13 -0700 Subject: [PATCH 1/4] [perf] avoid full-sequence materialization in VSA coarse/sparse combine The compression branch computed a per-block coarse output and then expanded it across the full sequence with repeat() before combining it with the sparse output. The gated combine then allocated two more [B, H, S, D] tensors for the product and the sum. Keep the coarse result at block resolution, [B, H, n_blocks, 1, D], and let it broadcast over the intra-block axis during the combine. This mirrors the BSHD 128/256 path, which already broadcasts out_c_blk.unsqueeze(2) rather than repeating. Under no_grad the combine accumulates into the sparse output with addcmul_, removing the remaining two full-sequence temporaries; with grad enabled it stays out-of-place, for the same reason the BSHD path documents (the sparse output is saved by FA4's autograd node for backward). Ungated the result is bit-exact. Gated, addcmul_ fuses the multiply-add instead of rounding the intermediate product to bf16, so the rounding differs from the old path; in the added tests the fused result is the more accurate of the two when both are compared against fp32. Top-k routing, shapes and dtypes are unchanged, and 64/128/256 dispatch is unaffected. Isolated combine microbenchmark at an H3-like shape (bf16, 56 heads, dim 128, 15488 tokens, block 64), gated: 4.06 -> 2.12 ms and 848 -> 212 MiB peak allocated. This measures the combine alone, not an end-to-end workload. Adds tests/test_vsa_combine.py (17 cases, CI-sized, no model or pipeline dependency) and benchmarks/bench_vsa_combine.py. --- .../benchmarks/bench_vsa_combine.py | 102 +++++++ .../python/fastvideo_kernel/ops.py | 63 ++++- fastvideo-kernel/tests/test_vsa_combine.py | 263 ++++++++++++++++++ 3 files changed, 425 insertions(+), 3 deletions(-) create mode 100644 fastvideo-kernel/benchmarks/bench_vsa_combine.py create mode 100644 fastvideo-kernel/tests/test_vsa_combine.py diff --git a/fastvideo-kernel/benchmarks/bench_vsa_combine.py b/fastvideo-kernel/benchmarks/bench_vsa_combine.py new file mode 100644 index 0000000000..b9615a2ecd --- /dev/null +++ b/fastvideo-kernel/benchmarks/bench_vsa_combine.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark the coarse/sparse combine in the 64-block VSA path. + +Compares the previous implementation (coarse output expanded across the full +sequence with ``repeat()``, then an out-of-place multiply and add) against the +current one (coarse output kept at block resolution and broadcast). + + python benchmarks/bench_vsa_combine.py # CI-sized shapes + python benchmarks/bench_vsa_combine.py --large # adds an H3-like shape + +The large case is representative of MiniMax H3 inference: bf16, 56 heads, head +dim 128, ~15-16k tokens, block size 64, top-k ~20%. It needs roughly 3 GiB and +is opt-in so the default run stays cheap. +""" +from __future__ import annotations + +import argparse + +import torch + +from fastvideo_kernel.ops import _combine_coarse_sparse + +BE = 64 + + +def _reference_combine(out_c, out_s, weight, batch, heads, n_blocks, be, dim, seq): + out_c = out_c.repeat(1, 1, 1, be, 1).view(batch, heads, seq, dim) + if weight is not None: + return out_c * weight + out_s + return out_c + out_s + + +def _time(fn, iters: int, warmup: int = 5) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _measure(fn, iters): + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + base_a, base_r = torch.cuda.memory_allocated(), torch.cuda.memory_reserved() + ms = _time(fn, iters) + torch.cuda.synchronize() + return (ms, + (torch.cuda.max_memory_allocated() - base_a) / 2**20, + (torch.cuda.max_memory_reserved() - base_r) / 2**20) + + +def run_case(name, batch, heads, n_blocks, dim, gated, iters): + seq = n_blocks * BE + torch.manual_seed(0) + out_c = torch.randn(batch, heads, n_blocks, 1, dim, device="cuda", dtype=torch.bfloat16) + out_s = torch.randn(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16) + weight = torch.rand(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16) if gated else None + full_mib = batch * heads * seq * dim * 2 / 2**20 + + with torch.no_grad(): + ref_ms, ref_a, ref_r = _measure( + lambda: _reference_combine(out_c, out_s.clone(), weight, batch, heads, n_blocks, BE, dim, + seq), iters) + new_ms, new_a, new_r = _measure( + lambda: _combine_coarse_sparse(out_c, out_s.clone(), weight, batch, heads, n_blocks, BE, + dim, seq), iters) + + print(f"\n{name} B={batch} H={heads} S={seq} D={dim} gated={gated} " + f"(one full [B,H,S,D] bf16 tensor = {full_mib:.1f} MiB)") + print(f" {'':<8}{'latency ms':>12}{'peak alloc MiB':>17}{'peak reserved MiB':>20}") + print(f" {'before':<8}{ref_ms:>12.4f}{ref_a:>17.1f}{ref_r:>20.1f}") + print(f" {'after':<8}{new_ms:>12.4f}{new_a:>17.1f}{new_r:>20.1f}") + speed = (ref_ms / new_ms - 1.0) * 100.0 if new_ms else 0.0 + print(f" {'delta':<8}{speed:>11.1f}%{new_a - ref_a:>17.1f}{new_r - ref_r:>20.1f}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--large", action="store_true", help="add the H3-like ~15.4k-token case") + ap.add_argument("--iters", type=int, default=50) + args = ap.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA/ROCm device required") + print("device:", torch.cuda.get_device_name(0)) + + for gated in (False, True): + run_case("small ", 1, 8, 64, 128, gated, args.iters) + if args.large: + # 15,488 tokens = 242 blocks x 64, matching MiniMax H3 at 864x480/124f + for gated in (False, True): + run_case("H3-like", 1, 56, 242, 128, gated, max(10, args.iters // 5)) + + +if __name__ == "__main__": + main() diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index 474c8c889d..2638d2b41f 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -116,8 +116,10 @@ def video_sparse_attn( scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) + # Kept at block resolution: [B, H, q_num_blocks, 1, D]. The coarse result is + # constant within a block, so it can broadcast over the intra-block axis in + # the combine below instead of being materialized across the full sequence. out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) - out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim) # Sparse branch (fused Triton topk mask) mask = fused_topk_mask(scores, topk) @@ -128,9 +130,64 @@ def video_sparse_attn( else: out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] + return _combine_coarse_sparse(out_c, out_s, compress_attn_weight, batch, heads, q_num_blocks, + block_elements, dim, q_seq_len) + + +def _combine_coarse_sparse( + out_c: torch.Tensor, + out_s: torch.Tensor, + compress_attn_weight: torch.Tensor | None, + batch: int, + heads: int, + q_num_blocks: int, + block_elements: int, + dim: int, + q_seq_len: int, +) -> torch.Tensor: + """Combine the block-resolution coarse output with the sparse output. + + ``out_c`` is [B, H, q_num_blocks, 1, D] and broadcasts over the intra-block + axis, so the coarse branch is never expanded to [B, H, S, D]. This mirrors + the BSHD 128/256 combine, which already broadcasts ``out_c_blk.unsqueeze(2)`` + instead of repeating. + + That path notes it must stay out-of-place because the sparse output is saved + by FA4's autograd node for backward. The same applies whenever grad is + enabled, so the autograd branch below is out-of-place too. Under + ``no_grad``/``inference_mode`` no autograd node is attached to ``out_s`` and + nothing outside this function aliases it, so accumulating into it is safe and + removes the last two full-sequence temporaries. + + Ungated this is bit-exact with the previous ``out_c.repeat(...) + out_s``. + Gated, ``addcmul_`` fuses the multiply-add rather than rounding the + intermediate product to bf16, so the rounding differs from the old path. In + the added tests the fused result is the more accurate of the two when both + are compared against fp32. See ``tests/test_vsa_combine.py``. + """ + blocked = (batch, heads, q_num_blocks, block_elements, dim) + + def as_blocked(t: torch.Tensor) -> torch.Tensor: + return t.view(*blocked) if t.is_contiguous() else t.reshape(*blocked) + + if torch.is_grad_enabled(): + out_s_b = as_blocked(out_s) + if compress_attn_weight is not None: + combined = out_c * as_blocked(compress_attn_weight) + out_s_b + else: + combined = out_c + out_s_b + return combined.view(batch, heads, q_seq_len, dim) + + # Inference: accumulate in place so neither the broadcast product nor the sum + # allocates another full-sequence tensor. Requires a real view of out_s. + if not out_s.is_contiguous(): + out_s = out_s.contiguous() + out_s_b = out_s.view(*blocked) if compress_attn_weight is not None: - return out_c * compress_attn_weight + out_s - return out_c + out_s + out_s_b.addcmul_(out_c, as_blocked(compress_attn_weight)) + else: + out_s_b.add_(out_c) + return out_s def video_sparse_attn_bshd( diff --git a/fastvideo-kernel/tests/test_vsa_combine.py b/fastvideo-kernel/tests/test_vsa_combine.py new file mode 100644 index 0000000000..d9aa77b976 --- /dev/null +++ b/fastvideo-kernel/tests/test_vsa_combine.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Equivalence tests for the block-resolution coarse/sparse combine in VSA. + +The combine used to expand the coarse block output across the full sequence with +``repeat()`` before adding it to the sparse output. It now keeps the coarse +result at block resolution and broadcasts. These tests pin the numerics of that +change against a literal transcription of the previous implementation. + +Shapes are deliberately small so this runs in CI; no MiniMax H3 checkpoint, +pipeline or ComfyUI dependency is involved. +""" +import pytest +import torch + +from fastvideo_kernel.ops import _combine_coarse_sparse + +DEVICE = "cuda" +BLOCK_ELEMENTS = 64 + + +def _reference_combine(out_c_blocked, out_s, weight, batch, heads, n_blocks, be, dim, seq): + """The pre-change implementation, transcribed literally. + + ``out_c`` was expanded with repeat() to [B, H, S, D], then combined + out-of-place with a full-size multiply and add. + """ + out_c = out_c_blocked.repeat(1, 1, 1, be, 1).view(batch, heads, seq, dim) + if weight is not None: + return out_c * weight + out_s + return out_c + out_s + + +def _make(batch, heads, n_blocks, dim, gated, seed=0): + torch.manual_seed(seed) + be = BLOCK_ELEMENTS + seq = n_blocks * be + out_c = torch.randn(batch, heads, n_blocks, 1, dim, device=DEVICE, dtype=torch.bfloat16) + out_s = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) + weight = torch.rand(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) if gated else None + return out_c, out_s, weight, be, seq + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("batch,heads,n_blocks,dim", [(1, 4, 8, 64), (1, 8, 16, 128), (2, 4, 5, 128)]) +def test_ungated_is_bit_exact(batch, heads, n_blocks, dim): + """Ungated: pure broadcast add, so the result must be bit-identical.""" + out_c, out_s, _, be, seq = _make(batch, heads, n_blocks, dim, gated=False) + with torch.no_grad(): + ref = _reference_combine(out_c, out_s.clone(), None, batch, heads, n_blocks, be, dim, seq) + got = _combine_coarse_sparse(out_c, out_s.clone(), None, batch, heads, n_blocks, be, dim, seq) + assert got.shape == ref.shape == (batch, heads, seq, dim) + assert got.dtype == ref.dtype == torch.bfloat16 + assert torch.equal(got, ref), "ungated combine must be bit-exact" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("batch,heads,n_blocks,dim", [(1, 4, 8, 64), (1, 8, 16, 128), (2, 4, 5, 128)]) +def test_gated_is_no_less_accurate(batch, heads, n_blocks, dim): + """Gated: the combine must be at least as accurate as the old one. + + The old path rounded the coarse*gate product to bf16 before adding, while + ``addcmul_`` fuses the multiply-add. The rounding therefore differs, and the + two implementations disagree on a fraction of elements. The meaningful + assertion is accuracy against fp32 rather than agreement with the old + rounding, so that is what this checks. + """ + out_c, out_s, weight, be, seq = _make(batch, heads, n_blocks, dim, gated=True) + with torch.no_grad(): + ref = _reference_combine(out_c, out_s.clone(), weight, batch, heads, n_blocks, be, dim, seq) + got = _combine_coarse_sparse(out_c, out_s.clone(), weight, batch, heads, n_blocks, be, dim, seq) + truth = (out_c.float().repeat(1, 1, 1, be, 1).view(batch, heads, seq, dim) * weight.float() + + out_s.float()) + + assert got.shape == ref.shape and got.dtype == ref.dtype + + old_err, new_err = (ref.float() - truth).abs(), (got.float() - truth).abs() + assert new_err.mean() <= old_err.mean() * 1.02, "combine is less accurate than the old path" + assert new_err.max() <= old_err.max() * 1.02, "combine has a worse worst case than the old path" + + # And the two implementations still agree to bf16 rounding overall. + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel <= 5e-3, f"relative L2 vs old implementation {rel:.3e} too large" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("gated", [False, True]) +def test_autograd_path_is_out_of_place(gated): + """With grad enabled the combine must not mutate the sparse output. + + The BSHD 128/256 combine documents that the sparse output is saved by FA4's + autograd node; in-place mutation there would invalidate the graph. The + in-place fast path must therefore be confined to no-grad. + """ + out_c, out_s, weight, be, seq = _make(1, 4, 8, 64, gated=gated) + out_s = out_s.requires_grad_(True) + before = out_s.detach().clone() + got = _combine_coarse_sparse(out_c, out_s, weight, 1, 4, 8, be, 64, seq) + assert torch.equal(out_s.detach(), before), "combine mutated a grad-tracking tensor" + assert got.requires_grad + got.sum().backward() + assert out_s.grad is not None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_noncontiguous_sparse_output(): + """A non-contiguous sparse output must still combine correctly.""" + batch, heads, n_blocks, dim = 1, 4, 8, 64 + out_c, out_s, weight, be, seq = _make(batch, heads, n_blocks, dim, gated=True) + # transpose-then-restore yields an equal but non-contiguous tensor + noncontig = out_s.transpose(1, 2).transpose(1, 2) + if noncontig.is_contiguous(): + noncontig = out_s[:, :, torch.arange(seq, device=DEVICE)] + with torch.no_grad(): + ref = _reference_combine(out_c, out_s.clone(), weight, batch, heads, n_blocks, be, dim, seq) + got = _combine_coarse_sparse(out_c, noncontig.clone(), weight, batch, heads, n_blocks, be, dim, + seq) + assert got.shape == ref.shape + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel <= 5e-3, f"relative L2 {rel:.3e} too large for a non-contiguous input" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_no_full_sequence_coarse_materialization(): + """Peak allocation must not include a full-sequence copy of the coarse output.""" + batch, heads, n_blocks, dim = 1, 8, 128, 128 + out_c, out_s, weight, be, seq = _make(batch, heads, n_blocks, dim, gated=True) + full = batch * heads * seq * dim * out_s.element_size() + + torch.cuda.synchronize() + with torch.no_grad(): + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + _combine_coarse_sparse(out_c, out_s, weight, batch, heads, n_blocks, be, dim, seq) + torch.cuda.synchronize() + new_peak = torch.cuda.max_memory_allocated() - base + + # The in-place no-grad path should allocate well under one extra full tensor. + assert new_peak < full, (f"combine allocated {new_peak} bytes, at least one full-sequence " + f"tensor ({full} bytes)") + + +# -------------------------------------------------------------------------- +# End-to-end through video_sparse_attn (64-block path) +# -------------------------------------------------------------------------- + + +def _reference_video_sparse_attn(q, k, v, variable_block_sizes, q_variable_block_sizes, topk, + block_size, compress_attn_weight): + """Pre-change ``video_sparse_attn`` body for the 64-block path.""" + from fastvideo_kernel.block_sparse_attn import block_sparse_attn + from fastvideo_kernel.triton_kernels.fused_compress_topk import (fused_block_mean, + fused_topk_mask) + block_elements = block_size[0] * block_size[1] * block_size[2] + batch, heads, q_seq_len, dim = q.shape + q_num_blocks = q_seq_len // block_elements + + q_c = fused_block_mean(q, q_variable_block_sizes, block_elements) + k_c = fused_block_mean(k, variable_block_sizes, block_elements) + v_c = fused_block_mean(v, variable_block_sizes, block_elements) + + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5) + attn = torch.softmax(scores, dim=-1) + out_c = torch.matmul(attn, v_c) + out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) + out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim) + + mask = fused_topk_mask(scores, topk) + out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] + if compress_attn_weight is not None: + return out_c * compress_attn_weight + out_s, mask + return out_c + out_s, mask + + +def _vsa_inputs(heads, n_blocks, dim, gated, seed=0): + torch.manual_seed(seed) + be, batch = BLOCK_ELEMENTS, 1 + seq = n_blocks * be + q = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) + vbs = torch.full((n_blocks,), be, device=DEVICE, dtype=torch.int32) + w = torch.rand(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) if gated else None + return q, k, v, vbs, w, seq + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("gated", [False, True]) +@pytest.mark.parametrize("heads,n_blocks,dim,ratio", [(4, 8, 64, 0.5), (8, 16, 128, 0.2)]) +def test_video_sparse_attn_matches_reference(gated, heads, n_blocks, dim, ratio): + """Full 64-block VSA: output and top-k routing must match the old path.""" + from fastvideo_kernel.ops import video_sparse_attn + + q, k, v, vbs, w, seq = _vsa_inputs(heads, n_blocks, dim, gated) + topk = max(1, int(n_blocks * ratio)) + with torch.no_grad(): + ref, ref_mask = _reference_video_sparse_attn(q, k, v, vbs, vbs, topk, (4, 4, 4), w) + got = video_sparse_attn(q, k, v, vbs, vbs, topk, (4, 4, 4), compress_attn_weight=w) + + assert got.shape == ref.shape == (1, heads, seq, dim) + assert got.dtype == ref.dtype == torch.bfloat16 + # routing is untouched by the combine change + assert bool(ref_mask.sum(-1).eq(topk).all().item()), "top-k rows must select exactly topk blocks" + if gated: + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel <= 5e-3, f"relative L2 {rel:.3e} too large" + else: + assert torch.equal(got, ref), "ungated full path must be bit-exact" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_block_size_64_selects_triton_path(): + """block_size (4,4,4) must still route to the 64-block sparse kernel.""" + import fastvideo_kernel.ops as ops + + called = {} + original = ops.block_sparse_attn + + def spy(*args, **kwargs): + called["hit"] = True + return original(*args, **kwargs) + + q, k, v, vbs, w, _ = _vsa_inputs(4, 8, 64, gated=True) + ops.block_sparse_attn = spy + try: + with torch.no_grad(): + ops.video_sparse_attn(q, k, v, vbs, vbs, 4, (4, 4, 4), compress_attn_weight=w) + finally: + ops.block_sparse_attn = original + assert called.get("hit"), "64-block path did not call block_sparse_attn" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("block_elements", [128, 256]) +def test_128_256_paths_still_dispatch(block_elements): + """The 128/256 kernels must still be selected; the combine is shared.""" + import fastvideo_kernel.ops as ops + + name = "block_sparse_attn_128" if block_elements == 128 else "block_sparse_attn_256" + called = {} + original = getattr(ops, name) + + def spy(*args, **kwargs): + called["hit"] = True + return original(*args, **kwargs) + + be = block_elements + n_blocks, heads, dim = 4, 4, 64 + seq = n_blocks * be + torch.manual_seed(0) + q, k, v = (torch.randn(1, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) for _ in range(3)) + vbs = torch.full((n_blocks,), be, device=DEVICE, dtype=torch.int32) + bs = (be // 16, 4, 4) + assert bs[0] * bs[1] * bs[2] == be + + setattr(ops, name, spy) + try: + with torch.no_grad(): + ops.video_sparse_attn(q, k, v, vbs, vbs, 2, bs) + except Exception as exc: # kernel unavailable on this backend + pytest.skip(f"{name} unavailable here: {type(exc).__name__}: {exc}") + finally: + setattr(ops, name, original) + assert called.get("hit"), f"{name} was not dispatched" From f9ff387f6adcf694ff42cc6f95156ba9d9c6528e Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 5 Sep 2026 00:39:19 +0000 Subject: [PATCH 2/4] [perf]: harden the VSA coarse/sparse combine and share it across layouts Follow-up to the block-resolution combine in video_sparse_attn: - one helper for both layouts (seq_dim=2 for the BHSD entry, seq_dim=1 for video_sparse_attn_bshd), so the 128/256 BSHD path also stops allocating combine temporaries at inference - the grad branch uses out-of-place torch.addcmul: one temporary instead of two, and bit-identical to the in-place addcmul_ path, so results no longer depend on the caller's grad mode - the in-place path is keyed on out_s.requires_grad (the actual aliasing invariant) rather than the global grad mode, and skipped when dtype promotion would otherwise downcast into out_s - strict shape validation of out_s and the gate (same-numel wrong-layout inputs used to be accepted silently); the dead .contiguous() fallback is gone because splitting the sequence axis is always expressible as a view - traceable by torch.compile(fullgraph=True) Tests (GB200): 77 cases on the default Triton route and 44 on the FA4 CuTe route. Ungated results are bit-exact with the old path; every gated element is within half a bf16 ulp of the fp32 truth in all four grad/no-grad modes; gradients match the old path; exact peak-allocation pins (0 bytes in place, one full tensor out of place); end-to-end value, routing and gradient checks for 64/128/256 tiles in both layouts; a fullgraph compile check. The benchmark now uses triton.testing.do_bench, measures peak memory outside the timed window, adds end-to-end video_sparse_attn / video_sparse_attn_bshd cases, and drops the MiniMax H3 label (that backend does not call this combine). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NjFimeedTQWxgqSzP5xo4j --- .../benchmarks/bench_vsa_combine.py | 204 +++-- .../python/fastvideo_kernel/ops.py | 117 +-- fastvideo-kernel/tests/test_vsa_combine.py | 697 ++++++++++++++---- 3 files changed, 744 insertions(+), 274 deletions(-) diff --git a/fastvideo-kernel/benchmarks/bench_vsa_combine.py b/fastvideo-kernel/benchmarks/bench_vsa_combine.py index b9615a2ecd..9b67c25b0f 100644 --- a/fastvideo-kernel/benchmarks/bench_vsa_combine.py +++ b/fastvideo-kernel/benchmarks/bench_vsa_combine.py @@ -1,101 +1,187 @@ # SPDX-License-Identifier: Apache-2.0 -"""Benchmark the coarse/sparse combine in the 64-block VSA path. +"""Benchmark the coarse/sparse combine used by ``video_sparse_attn`` (BHSD entry, every tile size). Compares the previous implementation (coarse output expanded across the full -sequence with ``repeat()``, then an out-of-place multiply and add) against the -current one (coarse output kept at block resolution and broadcast). - - python benchmarks/bench_vsa_combine.py # CI-sized shapes - python benchmarks/bench_vsa_combine.py --large # adds an H3-like shape - -The large case is representative of MiniMax H3 inference: bf16, 56 heads, head -dim 128, ~15-16k tokens, block size 64, top-k ~20%. It needs roughly 3 GiB and -is opt-in so the default run stays cheap. +sequence with ``repeat()``, then an out-of-place multiply and add: three +full-sequence temporaries) against the current one (block-resolution broadcast, +fused ``addcmul``, in place whenever the sparse output is not in an autograd +graph, one temporary otherwise). + + python benchmarks/bench_vsa_combine.py # combine alone, CI-sized shape + python benchmarks/bench_vsa_combine.py --large # adds 56x15488x128 and a Wan-14B 480p-like shape + python benchmarks/bench_vsa_combine.py --e2e # full video_sparse_attn (64) / _bshd (256) calls + +Latency comes from ``triton.testing.do_bench`` (L2 flushed between runs, median). +Peak memory is measured on one call, separately from timing, with the operands +allocated beforehand, so the columns show the combine's own allocations only. """ from __future__ import annotations import argparse +from typing import Callable import torch +from triton.testing import do_bench +import fastvideo_kernel.ops as ops from fastvideo_kernel.ops import _combine_coarse_sparse BE = 64 - - -def _reference_combine(out_c, out_s, weight, batch, heads, n_blocks, be, dim, seq): - out_c = out_c.repeat(1, 1, 1, be, 1).view(batch, heads, seq, dim) +MIB = 2**20 + + +def _old_combine(out_c, out_s, weight, be, seq_dim): + """The pre-change combines, transcribed from the old ``video_sparse_attn`` / ``_bshd`` bodies.""" + if seq_dim == 2: # BHSD: repeat the coarse output to the full sequence, then out-of-place mul and add + batch, heads, n_blocks, dim = out_c.shape + out_c = out_c.unsqueeze(3).repeat(1, 1, 1, be, 1).view(batch, heads, n_blocks * be, dim) + if weight is not None: + return out_c * weight + out_s + return out_c + out_s + batch, n_blocks, heads, dim = out_c.shape # BSHD: broadcast multiply, out-of-place add + out_view = out_s.view(batch, n_blocks, be, heads, dim) if weight is not None: - return out_c * weight + out_s - return out_c + out_s - - -def _time(fn, iters: int, warmup: int = 5) -> float: - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start, end = torch.cuda.Event(True), torch.cuda.Event(True) - start.record() - for _ in range(iters): - fn() - end.record() - torch.cuda.synchronize() - return start.elapsed_time(end) / iters + out = out_view + out_c.unsqueeze(2) * weight.view(batch, n_blocks, be, heads, dim) + else: + out = out_view + out_c.unsqueeze(2) + return out.view(batch, n_blocks * be, heads, dim) -def _measure(fn, iters): +def _peak_mib(fn: Callable[[], torch.Tensor]) -> float: + fn() # warm up (Triton compilation and autotuning allocate scratch on the first call) torch.cuda.synchronize() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() - base_a, base_r = torch.cuda.memory_allocated(), torch.cuda.memory_reserved() - ms = _time(fn, iters) + base = torch.cuda.memory_allocated() + fn() torch.cuda.synchronize() - return (ms, - (torch.cuda.max_memory_allocated() - base_a) / 2**20, - (torch.cuda.max_memory_reserved() - base_r) / 2**20) + return (torch.cuda.max_memory_allocated() - base) / MIB + + +def _latency_us(fn: Callable[[], torch.Tensor]) -> float: + return do_bench(fn, warmup=25, rep=200, return_mode="median") * 1e3 -def run_case(name, batch, heads, n_blocks, dim, gated, iters): +def _row(label: str, fn: Callable[[], torch.Tensor], baseline_us: float | None = None) -> float: + peak = _peak_mib(fn) + us = _latency_us(fn) + speedup = f"{baseline_us / us:>7.2f}x" if baseline_us else f"{'':>8}" + print(f" {label:<38}{us:>12.1f}{speedup}{peak:>16.1f}") + return us + + +def _header(): + print(f" {'':<38}{'latency us':>12}{'vs before':>8}{'peak alloc MiB':>16}") + + +def bench_combine(name, batch, heads, n_blocks, dim, gated): seq = n_blocks * BE torch.manual_seed(0) - out_c = torch.randn(batch, heads, n_blocks, 1, dim, device="cuda", dtype=torch.bfloat16) + out_c = torch.randn(batch, heads, n_blocks, dim, device="cuda", dtype=torch.bfloat16) out_s = torch.randn(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16) weight = torch.rand(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16) if gated else None - full_mib = batch * heads * seq * dim * 2 / 2**20 + # The model's gate is BSHD; the BHSD caller used to copy it into BHSD layout. + weight_view = (torch.rand(batch, seq, heads, dim, device="cuda", dtype=torch.bfloat16).transpose(1, 2) + if gated else None) + scratch = out_s.clone() # accumulation target for the in-place variants (values drift; latency does not) + leaf = out_s.clone().requires_grad_(True) + + print(f"\n{name} B={batch} H={heads} S={seq} D={dim} gated={gated} " + f"(one full [B,H,S,D] bf16 tensor = {out_s.numel() * 2 / MIB:.1f} MiB)") + _header() + with torch.no_grad(): + before = _row("before (repeat, mul, add)", lambda: _old_combine(out_c, out_s, weight, BE, 2)) + _row("after, no grad (in place)", lambda: _combine_coarse_sparse(out_c, scratch, weight, BE, 2), before) + if gated: + _row("after, no grad, BSHD gate view", lambda: _combine_coarse_sparse(out_c, scratch, weight_view, BE, 2), + before) + with torch.enable_grad(): + _row("after, grad (out of place addcmul)", lambda: _combine_coarse_sparse(out_c, leaf, weight, BE, 2), before) + +def bench_e2e(name, heads, n_blocks, dim, gated, ratio=0.2): + seq = n_blocks * BE + torch.manual_seed(0) + q, k, v = (torch.randn(1, heads, seq, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + vbs = torch.full((n_blocks,), BE, device="cuda", dtype=torch.int32) + gate = torch.rand(1, seq, heads, dim, device="cuda", dtype=torch.bfloat16).transpose(1, 2) if gated else None + gate_contig = gate.contiguous() if gated else None + topk = max(1, int(n_blocks * ratio)) + + def call(g): + return ops.video_sparse_attn(q, k, v, vbs, vbs, topk, (4, 4, 4), compress_attn_weight=g) + + print(f"\n{name} video_sparse_attn B=1 H={heads} S={seq} D={dim} topk={topk}/{n_blocks} gated={gated} " + f"(one full tensor = {q.numel() * 2 / MIB:.1f} MiB)") + _header() + current = ops._combine_coarse_sparse + with torch.no_grad(): + ops._combine_coarse_sparse = _old_combine + try: + before = _row("before (old combine)", lambda: call(gate_contig)) + finally: + ops._combine_coarse_sparse = current + _row("after", lambda: call(gate_contig), before) + if gated: + _row("after, BSHD gate view (no copy)", lambda: call(gate), before) + + +def bench_e2e_bshd(name, heads, n_blocks, dim, gated, ratio=0.2): + """``video_sparse_attn_bshd`` at 256-token tiles (the production 256-tile route).""" + be = 256 + n_blocks = (n_blocks * BE) // be + seq = n_blocks * be + torch.manual_seed(0) + q, k, v = (torch.randn(1, seq, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + vbs = torch.full((n_blocks,), be, device="cuda", dtype=torch.int32) + gate = torch.rand(1, seq, heads, dim, device="cuda", dtype=torch.bfloat16) if gated else None + topk = max(1, int(n_blocks * ratio)) + + def call(): + return ops.video_sparse_attn_bshd(q, k, v, vbs, vbs, topk, (16, 4, 4), compress_attn_weight=gate) + + print(f"\n{name} video_sparse_attn_bshd B=1 H={heads} S={seq} D={dim} topk={topk}/{n_blocks} gated={gated} " + f"(one full tensor = {q.numel() * 2 / MIB:.1f} MiB)") + _header() + current = ops._combine_coarse_sparse with torch.no_grad(): - ref_ms, ref_a, ref_r = _measure( - lambda: _reference_combine(out_c, out_s.clone(), weight, batch, heads, n_blocks, BE, dim, - seq), iters) - new_ms, new_a, new_r = _measure( - lambda: _combine_coarse_sparse(out_c, out_s.clone(), weight, batch, heads, n_blocks, BE, - dim, seq), iters) - - print(f"\n{name} B={batch} H={heads} S={seq} D={dim} gated={gated} " - f"(one full [B,H,S,D] bf16 tensor = {full_mib:.1f} MiB)") - print(f" {'':<8}{'latency ms':>12}{'peak alloc MiB':>17}{'peak reserved MiB':>20}") - print(f" {'before':<8}{ref_ms:>12.4f}{ref_a:>17.1f}{ref_r:>20.1f}") - print(f" {'after':<8}{new_ms:>12.4f}{new_a:>17.1f}{new_r:>20.1f}") - speed = (ref_ms / new_ms - 1.0) * 100.0 if new_ms else 0.0 - print(f" {'delta':<8}{speed:>11.1f}%{new_a - ref_a:>17.1f}{new_r - ref_r:>20.1f}") + ops._combine_coarse_sparse = _old_combine + try: + before = _row("before (old combine)", call) + finally: + ops._combine_coarse_sparse = current + _row("after", call, before) def main(): ap = argparse.ArgumentParser() - ap.add_argument("--large", action="store_true", help="add the H3-like ~15.4k-token case") - ap.add_argument("--iters", type=int, default=50) + ap.add_argument("--large", action="store_true", help="add 56x15488x128 and a Wan-14B 480p-like 40x39936x128") + ap.add_argument("--e2e", action="store_true", + help="benchmark full video_sparse_attn (64-token tiles) and video_sparse_attn_bshd (256) calls") + ap.add_argument("--no-combine", action="store_true", help="skip the combine-only microbenchmark") args = ap.parse_args() if not torch.cuda.is_available(): raise SystemExit("CUDA/ROCm device required") print("device:", torch.cuda.get_device_name(0)) - for gated in (False, True): - run_case("small ", 1, 8, 64, 128, gated, args.iters) + cases = [("small", 1, 8, 64, 128)] if args.large: - # 15,488 tokens = 242 blocks x 64, matching MiniMax H3 at 864x480/124f - for gated in (False, True): - run_case("H3-like", 1, 56, 242, 128, gated, max(10, args.iters // 5)) + # 39,936 tokens = 624 tiles of 64: a 21x30x52 token grid (480p, 81 frames, patch 1x2x2) tiled 4x4x4. + cases += [("large-15k", 1, 56, 242, 128), ("wan14b-480p-like", 1, 40, 624, 128)] + if not args.no_combine: + for name, batch, heads, n_blocks, dim in cases: + for gated in (False, True): + bench_combine(name, batch, heads, n_blocks, dim, gated) + if args.e2e: + from fastvideo_kernel.block_sparse_attn_256 import _resolve_backend + print(f"\n128/256-tile sparse backend: {_resolve_backend()} (FASTVIDEO_VSA_CUTEDSL=1 selects the FA4 CuTe route)") + for name, batch, heads, n_blocks, dim in cases: + for gated in (False, True): + bench_e2e(name, heads, n_blocks, dim, gated) + for name, batch, heads, n_blocks, dim in cases: + for gated in (False, True): + bench_e2e_bshd(name, heads, n_blocks, dim, gated) if __name__ == "__main__": diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index 2638d2b41f..5115807043 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -115,11 +115,10 @@ def video_sparse_attn( scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) + # Kept at block resolution, [B, H, q_num_blocks, D]: the coarse result is + # constant within a block, so the combine broadcasts it over the intra-block + # axis instead of materializing it across the full sequence. out_c = torch.matmul(attn, v_c) - # Kept at block resolution: [B, H, q_num_blocks, 1, D]. The coarse result is - # constant within a block, so it can broadcast over the intra-block axis in - # the combine below instead of being materialized across the full sequence. - out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) # Sparse branch (fused Triton topk mask) mask = fused_topk_mask(scores, topk) @@ -130,64 +129,74 @@ def video_sparse_attn( else: out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] - return _combine_coarse_sparse(out_c, out_s, compress_attn_weight, batch, heads, q_num_blocks, - block_elements, dim, q_seq_len) + return _combine_coarse_sparse(out_c, out_s, compress_attn_weight, block_elements, seq_dim=2) def _combine_coarse_sparse( out_c: torch.Tensor, out_s: torch.Tensor, compress_attn_weight: torch.Tensor | None, - batch: int, - heads: int, - q_num_blocks: int, block_elements: int, - dim: int, - q_seq_len: int, + seq_dim: int, ) -> torch.Tensor: """Combine the block-resolution coarse output with the sparse output. - ``out_c`` is [B, H, q_num_blocks, 1, D] and broadcasts over the intra-block - axis, so the coarse branch is never expanded to [B, H, S, D]. This mirrors - the BSHD 128/256 combine, which already broadcasts ``out_c_blk.unsqueeze(2)`` - instead of repeating. - - That path notes it must stay out-of-place because the sparse output is saved - by FA4's autograd node for backward. The same applies whenever grad is - enabled, so the autograd branch below is out-of-place too. Under - ``no_grad``/``inference_mode`` no autograd node is attached to ``out_s`` and - nothing outside this function aliases it, so accumulating into it is safe and - removes the last two full-sequence temporaries. - - Ungated this is bit-exact with the previous ``out_c.repeat(...) + out_s``. - Gated, ``addcmul_`` fuses the multiply-add rather than rounding the - intermediate product to bf16, so the rounding differs from the old path. In - the added tests the fused result is the more accurate of the two when both - are compared against fp32. See ``tests/test_vsa_combine.py``. + ``out_s`` and the optional gate are full-sequence tensors whose sequence axis + is ``seq_dim`` (2 for [B, H, S, D], 1 for [B, S, H, D]); ``out_c`` has the + same layout with ``S // block_elements`` blocks on that axis. The coarse + result is constant within a block, so it broadcasts over the intra-block + axis of a view that splits the sequence axis into (blocks, block_elements) + and is never expanded to the full sequence. Splitting one dimension is + always expressible as a view, so no copy is made whatever the strides of + ``out_s`` or the gate (the BHSD caller may pass a transposed BSHD gate + directly). + + Numerics: ungated this is a plain broadcast add and bit-exact with the + previous ``out_c.repeat(...) + out_s``. Gated, both branches use + ``addcmul``, which multiplies and accumulates in fp32 and rounds once, so + the result is identical whether or not it runs in place, and it is at least + as accurate as the old two-rounding ``out_c * w + out_s`` (see + ``tests/test_vsa_combine.py``). + + In-place contract: when ``out_s`` does not require grad no autograd node has + saved it (every sparse kernel's node saves its output for backward, so a + grad-tracking ``out_s`` is never mutated), and within ``video_sparse_attn`` + nothing else aliases it. The combine then accumulates into ``out_s`` and + returns it, allocating nothing. The in-place path is skipped when the result + dtype would be promoted, since in place would silently downcast. Callers + must run the combine in the mode that produced ``out_s``: an inference + tensor combined outside ``inference_mode`` raises in the in-place update. + That is not guarded here because ``Tensor.is_inference`` is not traceable by + ``torch.compile`` and would split the graph. """ - blocked = (batch, heads, q_num_blocks, block_elements, dim) - - def as_blocked(t: torch.Tensor) -> torch.Tensor: - return t.view(*blocked) if t.is_contiguous() else t.reshape(*blocked) + full = tuple(out_s.shape) + q_num_blocks, remainder = divmod(full[seq_dim], block_elements) + coarse = full[:seq_dim] + (q_num_blocks, ) + full[seq_dim + 1:] + if remainder != 0 or tuple(out_c.shape) != coarse: + raise ValueError(f"expected out_c {list(coarse)} for out_s {list(full)} with block_elements=" + f"{block_elements} on dim {seq_dim}, got out_c {list(out_c.shape)}") + if compress_attn_weight is not None and tuple(compress_attn_weight.shape) != full: + raise ValueError(f"compress_attn_weight must match out_s {list(full)}, got " + f"{list(compress_attn_weight.shape)}") + + blocked = full[:seq_dim] + (q_num_blocks, block_elements) + full[seq_dim + 1:] + out_c = out_c.unsqueeze(seq_dim + 1) + out_s_b = out_s.view(*blocked) + gate_b = None if compress_attn_weight is None else compress_attn_weight.view(*blocked) - if torch.is_grad_enabled(): - out_s_b = as_blocked(out_s) - if compress_attn_weight is not None: - combined = out_c * as_blocked(compress_attn_weight) + out_s_b + same_dtype = out_c.dtype == out_s.dtype and (gate_b is None or gate_b.dtype == out_s.dtype) + if not out_s.requires_grad and same_dtype: + if gate_b is not None: + out_s_b.addcmul_(out_c, gate_b) else: - combined = out_c + out_s_b - return combined.view(batch, heads, q_seq_len, dim) + out_s_b.add_(out_c) + return out_s - # Inference: accumulate in place so neither the broadcast product nor the sum - # allocates another full-sequence tensor. Requires a real view of out_s. - if not out_s.is_contiguous(): - out_s = out_s.contiguous() - out_s_b = out_s.view(*blocked) - if compress_attn_weight is not None: - out_s_b.addcmul_(out_c, as_blocked(compress_attn_weight)) + if gate_b is not None: + combined = torch.addcmul(out_s_b, out_c, gate_b) else: - out_s_b.add_(out_c) - return out_s + combined = out_s_b + out_c + return combined.view(*full) def video_sparse_attn_bshd( @@ -248,19 +257,13 @@ def video_sparse_attn_bshd( scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) out_c_ch = torch.matmul(attn, v_ch) - out_c_blk = out_c_ch.permute(0, 2, 1, 3).contiguous() + out_c_blk = out_c_ch.permute(0, 2, 1, 3).contiguous() # [B, q_num_blocks, H, D] # Sparse branch (fused Triton topk mask + CuTe BSHD). mask = fused_topk_mask(scores, topk) attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd out_s, _ = attention(q, k, v, mask, variable_block_sizes) - # Out-of-place: ``out_s`` is the tensor FA4's autograd node saved for its - # backward, so mutating it in place invalidates the graph. - out_view = out_s.view(batch, q_num_blocks, block_elements, heads, dim) - if compress_attn_weight is not None: - gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim) - out = out_view + out_c_blk.unsqueeze(2) * gate_view - else: - out = out_view + out_c_blk.unsqueeze(2) - return out.view(batch, q_seq_len, heads, dim) + # Shared combine: out of place when ``out_s`` is saved by the kernel's + # autograd node (grad), in place otherwise. + return _combine_coarse_sparse(out_c_blk, out_s, compress_attn_weight, block_elements, seq_dim=1) diff --git a/fastvideo-kernel/tests/test_vsa_combine.py b/fastvideo-kernel/tests/test_vsa_combine.py index d9aa77b976..ec612431ce 100644 --- a/fastvideo-kernel/tests/test_vsa_combine.py +++ b/fastvideo-kernel/tests/test_vsa_combine.py @@ -1,153 +1,414 @@ # SPDX-License-Identifier: Apache-2.0 -"""Equivalence tests for the block-resolution coarse/sparse combine in VSA. +"""Tests for the block-resolution coarse/sparse combine in VSA. The combine used to expand the coarse block output across the full sequence with -``repeat()`` before adding it to the sparse output. It now keeps the coarse -result at block resolution and broadcasts. These tests pin the numerics of that -change against a literal transcription of the previous implementation. - -Shapes are deliberately small so this runs in CI; no MiniMax H3 checkpoint, -pipeline or ComfyUI dependency is involved. +``repeat()`` before multiplying by the gate and adding the sparse output, three +full-sequence temporaries in all. It now keeps the coarse result at block +resolution, broadcasts it over a blocked view, and accumulates into the sparse +output in place whenever that tensor is not part of an autograd graph. + +These tests pin the numerics against a literal transcription of the previous +implementation and against an fp32 oracle, and pin the allocation behaviour of +both branches. Shapes are deliberately small so this runs in CI. """ import pytest import torch +import fastvideo_kernel.ops as ops from fastvideo_kernel.ops import _combine_coarse_sparse +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") + DEVICE = "cuda" -BLOCK_ELEMENTS = 64 +BE = 64 +# bf16 keeps 8 significand bits, so a single round-to-nearest of an fp32 value +# is off by at most half an ulp, i.e. at most 2**-8 relative. +BF16_HALF_ULP_REL = 2.0**-8 + +# How the combine is invoked. "grad_leaf" is the training case (``out_s`` is +# saved by the sparse kernel's autograd node and must not be mutated); the other +# three are the cases in which accumulating into ``out_s`` is safe. +MODES = ["no_grad", "inference_mode", "grad_untracked", "grad_leaf"] -def _reference_combine(out_c_blocked, out_s, weight, batch, heads, n_blocks, be, dim, seq): - """The pre-change implementation, transcribed literally. +def _reference_combine(out_c, out_s, weight, be): + """The pre-change BHSD implementation, transcribed literally. - ``out_c`` was expanded with repeat() to [B, H, S, D], then combined - out-of-place with a full-size multiply and add. + ``out_c`` [B, H, n_blocks, D] was expanded with repeat() to [B, H, S, D], + then combined out-of-place with a full-size multiply and add. """ - out_c = out_c_blocked.repeat(1, 1, 1, be, 1).view(batch, heads, seq, dim) + batch, heads, n_blocks, dim = out_c.shape + out_c = out_c.unsqueeze(3).repeat(1, 1, 1, be, 1).view(batch, heads, n_blocks * be, dim) if weight is not None: return out_c * weight + out_s return out_c + out_s -def _make(batch, heads, n_blocks, dim, gated, seed=0): +def _reference_combine_bshd(out_c, out_s, weight, be): + """The pre-change BSHD implementation (``video_sparse_attn_bshd``), transcribed literally.""" + batch, n_blocks, heads, dim = out_c.shape + out_view = out_s.view(batch, n_blocks, be, heads, dim) + if weight is not None: + out = out_view + out_c.unsqueeze(2) * weight.view(batch, n_blocks, be, heads, dim) + else: + out = out_view + out_c.unsqueeze(2) + return out.view(batch, n_blocks * be, heads, dim) + + +def _fp32_truth(out_c, out_s, weight, be, seq_dim=2): + """The exact-product, once-rounded fp32 value of the combine.""" + out_c = out_c.float().unsqueeze(seq_dim + 1) + blocked = out_s.shape[:seq_dim] + (out_c.shape[seq_dim], be) + out_s.shape[seq_dim + 1:] + out_s = out_s.float().view(*blocked) + if weight is not None: + return (out_c * weight.float().view(*blocked) + out_s).view(weight.shape) + return (out_c + out_s).view(out_s.shape[:seq_dim] + (-1, ) + out_s.shape[seq_dim + 2:]) + + +def _make(batch, heads, n_blocks, dim, gated, be=BE, seed=0, dtype=torch.bfloat16, seq_dim=2): torch.manual_seed(seed) - be = BLOCK_ELEMENTS seq = n_blocks * be - out_c = torch.randn(batch, heads, n_blocks, 1, dim, device=DEVICE, dtype=torch.bfloat16) - out_s = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) - weight = torch.rand(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) if gated else None - return out_c, out_s, weight, be, seq + if seq_dim == 2: + out_c = torch.randn(batch, heads, n_blocks, dim, device=DEVICE, dtype=dtype) + full = (batch, heads, seq, dim) + else: + out_c = torch.randn(batch, n_blocks, heads, dim, device=DEVICE, dtype=dtype) + full = (batch, seq, heads, dim) + out_s = torch.randn(*full, device=DEVICE, dtype=dtype) + weight = torch.rand(*full, device=DEVICE, dtype=dtype) if gated else None + return out_c, out_s, weight + + +def _run(mode, out_c, out_s, weight, be=BE, seq_dim=2): + """Run the combine under ``mode`` on clones; returns (result, out_s clone).""" + out_s = out_s.clone() + if mode == "no_grad": + with torch.no_grad(): + return _combine_coarse_sparse(out_c, out_s, weight, be, seq_dim), out_s + if mode == "inference_mode": + with torch.inference_mode(): + return _combine_coarse_sparse(out_c, out_s, weight, be, seq_dim), out_s + if mode == "grad_untracked": + with torch.enable_grad(): + return _combine_coarse_sparse(out_c, out_s, weight, be, seq_dim), out_s + assert mode == "grad_leaf" + out_s.requires_grad_(True) + with torch.enable_grad(): + return _combine_coarse_sparse(out_c, out_s, weight, be, seq_dim), out_s + + +# -------------------------------------------------------------------------- +# Numerics +# -------------------------------------------------------------------------- -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("batch,heads,n_blocks,dim", [(1, 4, 8, 64), (1, 8, 16, 128), (2, 4, 5, 128)]) -def test_ungated_is_bit_exact(batch, heads, n_blocks, dim): - """Ungated: pure broadcast add, so the result must be bit-identical.""" - out_c, out_s, _, be, seq = _make(batch, heads, n_blocks, dim, gated=False) - with torch.no_grad(): - ref = _reference_combine(out_c, out_s.clone(), None, batch, heads, n_blocks, be, dim, seq) - got = _combine_coarse_sparse(out_c, out_s.clone(), None, batch, heads, n_blocks, be, dim, seq) - assert got.shape == ref.shape == (batch, heads, seq, dim) +def test_ungated_is_bit_exact(mode, batch, heads, n_blocks, dim): + """Ungated: pure broadcast add, so every branch must be bit-identical to the old path.""" + out_c, out_s, _ = _make(batch, heads, n_blocks, dim, gated=False) + ref = _reference_combine(out_c, out_s.clone(), None, BE) + got, _ = _run(mode, out_c, out_s, None) + assert got.shape == ref.shape == (batch, heads, n_blocks * BE, dim) assert got.dtype == ref.dtype == torch.bfloat16 - assert torch.equal(got, ref), "ungated combine must be bit-exact" + assert torch.equal(got.detach(), ref), "ungated combine must be bit-exact" -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") @pytest.mark.parametrize("batch,heads,n_blocks,dim", [(1, 4, 8, 64), (1, 8, 16, 128), (2, 4, 5, 128)]) -def test_gated_is_no_less_accurate(batch, heads, n_blocks, dim): - """Gated: the combine must be at least as accurate as the old one. - - The old path rounded the coarse*gate product to bf16 before adding, while - ``addcmul_`` fuses the multiply-add. The rounding therefore differs, and the - two implementations disagree on a fraction of elements. The meaningful - assertion is accuracy against fp32 rather than agreement with the old - rounding, so that is what this checks. +def test_gated_is_within_one_bf16_rounding_of_fp32(batch, heads, n_blocks, dim): + """Gated: every element is the once-rounded fp32 value, in every branch. + + ``addcmul`` computes ``out_s + out_c * w`` in fp32 and rounds once. The + product of two bf16 values is exact in fp32, so the fused result is within + half a bf16 ulp of the fp32 truth for every element. The old path rounded + the product to bf16 before adding and violates that bound on a noticeable + fraction of elements, so this assertion discriminates the two and would + catch a regression to an unfused implementation. """ - out_c, out_s, weight, be, seq = _make(batch, heads, n_blocks, dim, gated=True) - with torch.no_grad(): - ref = _reference_combine(out_c, out_s.clone(), weight, batch, heads, n_blocks, be, dim, seq) - got = _combine_coarse_sparse(out_c, out_s.clone(), weight, batch, heads, n_blocks, be, dim, seq) - truth = (out_c.float().repeat(1, 1, 1, be, 1).view(batch, heads, seq, dim) * weight.float() - + out_s.float()) - - assert got.shape == ref.shape and got.dtype == ref.dtype - - old_err, new_err = (ref.float() - truth).abs(), (got.float() - truth).abs() - assert new_err.mean() <= old_err.mean() * 1.02, "combine is less accurate than the old path" - assert new_err.max() <= old_err.max() * 1.02, "combine has a worse worst case than the old path" - + out_c, out_s, weight = _make(batch, heads, n_blocks, dim, gated=True) + truth = _fp32_truth(out_c, out_s, weight, BE) + ref = _reference_combine(out_c, out_s.clone(), weight, BE) + bound = BF16_HALF_ULP_REL * truth.abs() + + results = {mode: _run(mode, out_c, out_s, weight)[0].detach() for mode in MODES} + for mode, got in results.items(): + assert got.shape == ref.shape and got.dtype == ref.dtype + new_err = (got.float() - truth).abs() + assert bool((new_err <= bound).all()), f"[{mode}] combine is not within one bf16 rounding of fp32" + + # The in-place and out-of-place branches must agree bit for bit, so results + # do not depend on whether the caller runs under grad. + first = results[MODES[0]] + for mode, got in results.items(): + assert torch.equal(got, first), f"[{mode}] differs from [{MODES[0]}]" + + # The old path is never more accurate. + old_err = (ref.float() - truth).abs() + new_err = (first.float() - truth).abs() + assert new_err.mean() <= old_err.mean() + assert new_err.max() <= old_err.max() # And the two implementations still agree to bf16 rounding overall. - rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + rel = ((first.float() - ref.float()).norm() / ref.float().norm()).item() assert rel <= 5e-3, f"relative L2 vs old implementation {rel:.3e} too large" -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +# -------------------------------------------------------------------------- +# Autograd and aliasing +# -------------------------------------------------------------------------- + + @pytest.mark.parametrize("gated", [False, True]) def test_autograd_path_is_out_of_place(gated): - """With grad enabled the combine must not mutate the sparse output. - - The BSHD 128/256 combine documents that the sparse output is saved by FA4's - autograd node; in-place mutation there would invalidate the graph. The - in-place fast path must therefore be confined to no-grad. - """ - out_c, out_s, weight, be, seq = _make(1, 4, 8, 64, gated=gated) + """With ``out_s`` in a graph the combine must not mutate it, and grads match the old path.""" + out_c, out_s, weight = _make(1, 4, 8, 64, gated=gated) out_s = out_s.requires_grad_(True) + out_c = out_c.requires_grad_(True) + if gated: + weight = weight.requires_grad_(True) before = out_s.detach().clone() - got = _combine_coarse_sparse(out_c, out_s, weight, 1, 4, 8, be, 64, seq) + + got = _combine_coarse_sparse(out_c, out_s, weight, BE, 2) assert torch.equal(out_s.detach(), before), "combine mutated a grad-tracking tensor" + assert got.requires_grad and got.data_ptr() != out_s.data_ptr() + + torch.manual_seed(1) + grad_out = torch.randn_like(got) + got.backward(grad_out) + got_grads = [t.grad.clone() for t in (out_c, out_s, weight) if t is not None] + for t in (out_c, out_s, weight): + if t is not None: + t.grad = None + + ref = _reference_combine(out_c, out_s, weight, BE) + ref.backward(grad_out) + ref_grads = [t.grad for t in (out_c, out_s, weight) if t is not None] + for name, g, r in zip(("out_c", "out_s", "weight"), got_grads, ref_grads): + torch.testing.assert_close(g, r, rtol=2.0**-6, atol=0.0, msg=lambda m: f"grad of {name}: {m}") + + +@pytest.mark.parametrize("gated", [False, True]) +def test_untracked_sparse_output_accumulates_in_place_under_grad(gated): + """Grad enabled but ``out_s`` outside any graph: accumulate into it, keep grads for the rest. + + This is the case an eval/teacher forward or a frozen-attention finetune hits + when the surrounding code has not wrapped the call in ``no_grad``. In-place + is safe because no autograd node saved ``out_s`` and addcmul's backward only + needs the other two operands. + """ + out_c, out_s, weight = _make(1, 4, 8, 64, gated=gated) + out_c = out_c.requires_grad_(True) + if gated: + weight = weight.requires_grad_(True) + with torch.enable_grad(): + got = _combine_coarse_sparse(out_c, out_s, weight, BE, 2) + assert got.data_ptr() == out_s.data_ptr(), "untracked sparse output was not reused" assert got.requires_grad - got.sum().backward() - assert out_s.grad is not None + torch.manual_seed(1) + grad_out = torch.randn_like(got) + got.backward(grad_out) + + ref_c = out_c.detach().clone().requires_grad_(True) + ref_w = weight.detach().clone().requires_grad_(True) if gated else None + _reference_combine(ref_c, torch.zeros_like(out_s), ref_w, BE).backward(grad_out) + torch.testing.assert_close(out_c.grad, ref_c.grad, rtol=2.0**-6, atol=0.0) + if gated: + torch.testing.assert_close(weight.grad, ref_w.grad, rtol=2.0**-6, atol=0.0) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -def test_noncontiguous_sparse_output(): - """A non-contiguous sparse output must still combine correctly.""" - batch, heads, n_blocks, dim = 1, 4, 8, 64 - out_c, out_s, weight, be, seq = _make(batch, heads, n_blocks, dim, gated=True) - # transpose-then-restore yields an equal but non-contiguous tensor - noncontig = out_s.transpose(1, 2).transpose(1, 2) - if noncontig.is_contiguous(): - noncontig = out_s[:, :, torch.arange(seq, device=DEVICE)] +def test_in_place_requires_no_grad_tracking_on_sparse_output_only(): + """``out_c`` requiring grad under no_grad does not stop the in-place path.""" + out_c, out_s, weight = _make(1, 4, 8, 64, gated=True) + out_c.requires_grad_(True) with torch.no_grad(): - ref = _reference_combine(out_c, out_s.clone(), weight, batch, heads, n_blocks, be, dim, seq) - got = _combine_coarse_sparse(out_c, noncontig.clone(), weight, batch, heads, n_blocks, be, dim, - seq) - assert got.shape == ref.shape - rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() - assert rel <= 5e-3, f"relative L2 {rel:.3e} too large for a non-contiguous input" + got = _combine_coarse_sparse(out_c, out_s, weight, BE, 2) + assert got.data_ptr() == out_s.data_ptr() + assert not got.requires_grad + + +def test_inference_tensor_contract(): + """An inference-mode ``out_s`` combines in place inside inference mode and raises outside it. + + The helper does not guard this (``Tensor.is_inference`` is untraceable by + ``torch.compile``); ``video_sparse_attn`` always combines in the mode that + produced ``out_s``, and this pins the documented behaviour for other callers. + """ + out_c, out_s, weight = _make(1, 4, 8, 64, gated=True) + expected = _run("no_grad", out_c, out_s, weight)[0] + with torch.inference_mode(): + out_s_inf = out_s.clone() + got = _combine_coarse_sparse(out_c, out_s_inf, weight, BE, 2) + assert got.is_inference() and got.data_ptr() == out_s_inf.data_ptr() + assert torch.equal(got, expected) + + with torch.inference_mode(): + out_s_inf = out_s.clone() + with torch.no_grad(), pytest.raises(RuntimeError, match="[Ii]nference"): + _combine_coarse_sparse(out_c, out_s_inf, weight, BE, 2) + + +# -------------------------------------------------------------------------- +# Layout, shape and dtype contract +# -------------------------------------------------------------------------- + +@pytest.mark.parametrize("mode", ["no_grad", "grad_leaf"]) +def test_noncontiguous_sparse_output_and_gate(mode): + """Non-contiguous ``out_s`` and a transposed BSHD gate combine without a copy. + + Splitting the sequence axis is always a view, so neither operand needs to + be contiguous: the production caller can hand over the transposed BSHD gate + without ``.contiguous()``. + """ + batch, heads, n_blocks, dim = 1, 4, 8, 64 + seq = n_blocks * BE + out_c, out_s, weight = _make(batch, heads, n_blocks, dim, gated=True) + noncontig_s = out_s.transpose(2, 3).contiguous().transpose(2, 3) + assert not noncontig_s.is_contiguous() and torch.equal(noncontig_s, out_s) + torch.manual_seed(3) + gate_bshd = torch.rand(batch, seq, heads, dim, device=DEVICE, dtype=torch.bfloat16) + gate = gate_bshd.transpose(1, 2) + assert not gate.is_contiguous() + + contiguous_result, _ = _run(mode, out_c, out_s, gate.contiguous()) + got, arg = _run(mode, out_c, noncontig_s, gate) + assert not arg.is_contiguous(), "fixture lost its strides" + assert got.shape == (batch, heads, seq, dim) + assert torch.equal(got.detach(), contiguous_result.detach()), "strides changed the arithmetic" + if mode == "no_grad": + assert got.data_ptr() == arg.data_ptr(), "non-contiguous out_s was copied instead of reused" + + +def test_shape_mismatch_raises(): + """A same-numel gate or sparse output in the wrong layout must raise, not compute garbage.""" + batch, heads, n_blocks, dim = 1, 4, 8, 64 + seq = n_blocks * BE + out_c, out_s, weight = _make(batch, heads, n_blocks, dim, gated=True) + bshd_gate = torch.rand(batch, seq, heads, dim, device=DEVICE, dtype=torch.bfloat16) + for mode in ("no_grad", "grad_leaf"): + with pytest.raises(ValueError, match="compress_attn_weight"): + _run(mode, out_c, out_s, bshd_gate) + with pytest.raises(ValueError, match="out_s"): + _run(mode, out_c, out_s.view(batch, seq, heads, dim), weight) + with pytest.raises(ValueError): + _run(mode, out_c, out_s, weight[:, :, :, :1]) # broadcasting gates are not supported + with pytest.raises(ValueError): + _run(mode, out_c, out_s, weight, be=BE // 2) + + +@pytest.mark.parametrize("mode", MODES) +def test_dtype_promotion_is_never_downcast(mode): + """A wider gate or coarse output promotes the result like the old path did.""" + out_c, out_s, weight = _make(1, 4, 8, 64, gated=True) + ref = _reference_combine(out_c, out_s.clone(), weight.float(), BE) + got, arg = _run(mode, out_c, out_s, weight.float()) + assert got.dtype == ref.dtype == torch.float32 + assert got.data_ptr() != arg.data_ptr() + assert torch.equal(arg.detach(), out_s), "in-place path would have downcast into out_s" + torch.testing.assert_close(got.detach(), ref, rtol=2.0**-7, atol=0.0) + + ref = _reference_combine(out_c.float(), out_s.clone(), None, BE) + got, arg = _run(mode, out_c.float(), out_s, None) + assert got.dtype == ref.dtype == torch.float32 + assert torch.equal(got.detach(), ref) + + +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("gated", [False, True]) +def test_bshd_layout_matches_old_bshd_combine(mode, gated): + """seq_dim=1 (the 128/256 BSHD entry): bit-exact ungated, once-rounded gated, in place when allowed.""" + out_c, out_s, weight = _make(2, 4, 6, 64, gated=gated, seq_dim=1) + ref = _reference_combine_bshd(out_c, out_s.clone(), weight, BE) + got, arg = _run(mode, out_c, out_s, weight, seq_dim=1) + assert got.shape == ref.shape == (2, 6 * BE, 4, 64) + if gated: + truth = _fp32_truth(out_c, out_s, weight, BE, seq_dim=1) + assert bool(((got.detach().float() - truth).abs() <= BF16_HALF_ULP_REL * truth.abs()).all()) + assert torch.equal(got.detach(), _run("no_grad", out_c, out_s, weight, seq_dim=1)[0]) + else: + assert torch.equal(got.detach(), ref) + assert (got.data_ptr() == arg.data_ptr()) == (mode != "grad_leaf") -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -def test_no_full_sequence_coarse_materialization(): - """Peak allocation must not include a full-sequence copy of the coarse output.""" - batch, heads, n_blocks, dim = 1, 8, 128, 128 - out_c, out_s, weight, be, seq = _make(batch, heads, n_blocks, dim, gated=True) - full = batch * heads * seq * dim * out_s.element_size() +# -------------------------------------------------------------------------- +# Allocation +# -------------------------------------------------------------------------- + + +def _peak_bytes(fn): torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + out = fn() + torch.cuda.synchronize() + return torch.cuda.max_memory_allocated() - base, out + + +@pytest.mark.parametrize("seq_dim", [2, 1]) +@pytest.mark.parametrize("gated", [False, True]) +def test_no_grad_combine_allocates_nothing(gated, seq_dim): + """The in-place branch must not allocate at all, in either layout.""" + out_c, out_s, weight = _make(1, 8, 128, 128, gated=gated, seq_dim=seq_dim) with torch.no_grad(): - torch.cuda.reset_peak_memory_stats() - base = torch.cuda.memory_allocated() - _combine_coarse_sparse(out_c, out_s, weight, batch, heads, n_blocks, be, dim, seq) - torch.cuda.synchronize() - new_peak = torch.cuda.max_memory_allocated() - base + peak, out = _peak_bytes(lambda: _combine_coarse_sparse(out_c, out_s, weight, BE, seq_dim)) + assert peak == 0, f"no-grad combine allocated {peak} bytes" + assert out.data_ptr() == out_s.data_ptr() - # The in-place no-grad path should allocate well under one extra full tensor. - assert new_peak < full, (f"combine allocated {new_peak} bytes, at least one full-sequence " - f"tensor ({full} bytes)") + +@pytest.mark.parametrize("seq_dim", [2, 1]) +@pytest.mark.parametrize("gated", [False, True]) +def test_grad_combine_allocates_exactly_one_output(gated, seq_dim): + """The out-of-place branch allocates its output and nothing else. + + The old BHSD path peaked at three full-sequence tensors (the repeated + coarse output, the product and the sum) and the old BSHD path at two; + ``addcmul`` fuses the product away. + """ + out_c, out_s, weight = _make(1, 8, 128, 128, gated=gated, seq_dim=seq_dim) + out_s.requires_grad_(True) + full = out_s.numel() * out_s.element_size() + with torch.enable_grad(): + peak, out = _peak_bytes(lambda: _combine_coarse_sparse(out_c, out_s, weight, BE, seq_dim)) + assert peak == full, f"grad-mode combine allocated {peak} bytes, expected one full tensor ({full})" + assert out.data_ptr() != out_s.data_ptr() # -------------------------------------------------------------------------- -# End-to-end through video_sparse_attn (64-block path) +# torch.compile +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["no_grad", "grad_leaf"]) +def test_combine_compiles_fullgraph(mode): + """The branch predicate must be traceable so a compiled model keeps one graph.""" + torch._dynamo.reset() + compiled = torch.compile(_combine_coarse_sparse, fullgraph=True, dynamic=False) + out_c, out_s, weight = _make(1, 4, 8, 64, gated=True) + eager, _ = _run(mode, out_c, out_s, weight) + out_s = out_s.clone() + if mode == "grad_leaf": + out_s.requires_grad_(True) + with torch.enable_grad(): + got = compiled(out_c, out_s, weight, BE, 2) + else: + with torch.no_grad(): + got = compiled(out_c, out_s, weight, BE, 2) + assert torch.equal(got.detach(), eager.detach()) + + +# -------------------------------------------------------------------------- +# End-to-end through video_sparse_attn # -------------------------------------------------------------------------- def _reference_video_sparse_attn(q, k, v, variable_block_sizes, q_variable_block_sizes, topk, block_size, compress_attn_weight): - """Pre-change ``video_sparse_attn`` body for the 64-block path.""" + """Pre-change ``video_sparse_attn`` body, dispatching the sparse branch like ``ops`` does. + + The kernels are imported from their modules rather than looked up on + ``ops`` so the spies installed on ``ops`` see only the code under test. + """ from fastvideo_kernel.block_sparse_attn import block_sparse_attn + from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128, block_sparse_attn_256 from fastvideo_kernel.triton_kernels.fused_compress_topk import (fused_block_mean, fused_topk_mask) block_elements = block_size[0] * block_size[1] * block_size[2] @@ -165,99 +426,219 @@ def _reference_video_sparse_attn(q, k, v, variable_block_sizes, q_variable_block out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim) mask = fused_topk_mask(scores, topk) - out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] + if block_elements in (128, 256): + attention = block_sparse_attn_128 if block_elements == 128 else block_sparse_attn_256 + else: + attention = block_sparse_attn + out_s = attention(q, k, v, mask, variable_block_sizes)[0] if compress_attn_weight is not None: return out_c * compress_attn_weight + out_s, mask return out_c + out_s, mask -def _vsa_inputs(heads, n_blocks, dim, gated, seed=0): +def _vsa_inputs(heads, n_blocks, dim, be, seed=0): torch.manual_seed(seed) - be, batch = BLOCK_ELEMENTS, 1 + batch = 1 seq = n_blocks * be - q = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) - k = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) - v = torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) + q, k, v = (torch.randn(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) for _ in range(3)) vbs = torch.full((n_blocks,), be, device=DEVICE, dtype=torch.int32) - w = torch.rand(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) if gated else None - return q, k, v, vbs, w, seq + return q, k, v, vbs, seq -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -@pytest.mark.parametrize("gated", [False, True]) -@pytest.mark.parametrize("heads,n_blocks,dim,ratio", [(4, 8, 64, 0.5), (8, 16, 128, 0.2)]) -def test_video_sparse_attn_matches_reference(gated, heads, n_blocks, dim, ratio): - """Full 64-block VSA: output and top-k routing must match the old path.""" - from fastvideo_kernel.ops import video_sparse_attn +def _gate(layout, batch, heads, seq, dim, seed=1): + torch.manual_seed(seed) + if layout == "none": + return None + if layout == "bhsd": + return torch.rand(batch, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) + assert layout == "bshd_transposed" # the model's native gate layout, handed over as a view + gate = torch.rand(batch, seq, heads, dim, device=DEVICE, dtype=torch.bfloat16).transpose(1, 2) + assert not gate.is_contiguous() + return gate + - q, k, v, vbs, w, seq = _vsa_inputs(heads, n_blocks, dim, gated) +def _kernel_name(block_elements): + return {64: "block_sparse_attn", 128: "block_sparse_attn_128", 256: "block_sparse_attn_256"}[block_elements] + + +@pytest.mark.parametrize("gate_layout", ["none", "bhsd", "bshd_transposed"]) +@pytest.mark.parametrize("block_elements", [64, 128, 256]) +@pytest.mark.parametrize("heads,n_blocks,dim,ratio", [(4, 8, 64, 0.5), (8, 16, 128, 0.2)]) +def test_video_sparse_attn_matches_reference(monkeypatch, gate_layout, block_elements, heads, n_blocks, dim, + ratio): + """Inference VSA at every tile size: same kernel, same top-k routing, same output as the old path.""" + q, k, v, vbs, seq = _vsa_inputs(heads, n_blocks, dim, block_elements) + gate = _gate(gate_layout, 1, heads, seq, dim) topk = max(1, int(n_blocks * ratio)) + block_size = (block_elements // 16, 4, 4) + + # Spy on the kernel dispatch and capture the routing mask of the code under test. + name = _kernel_name(block_elements) + original = getattr(ops, name) + calls = [] + monkeypatch.setattr(ops, name, lambda *a, **kw: calls.append(True) or original(*a, **kw)) + masks = [] + original_topk = ops.fused_topk_mask + monkeypatch.setattr(ops, "fused_topk_mask", lambda *a, **kw: masks.append(original_topk(*a, **kw)) or masks[-1]) + with torch.no_grad(): - ref, ref_mask = _reference_video_sparse_attn(q, k, v, vbs, vbs, topk, (4, 4, 4), w) - got = video_sparse_attn(q, k, v, vbs, vbs, topk, (4, 4, 4), compress_attn_weight=w) + ref, ref_mask = _reference_video_sparse_attn(q, k, v, vbs, vbs, topk, block_size, gate) + got = ops.video_sparse_attn(q, k, v, vbs, vbs, topk, block_size, compress_attn_weight=gate) + assert len(calls) == 1, f"{name} was dispatched {len(calls)} times, expected once" + assert len(masks) == 1 and torch.equal(masks[0], ref_mask), "top-k routing changed" assert got.shape == ref.shape == (1, heads, seq, dim) assert got.dtype == ref.dtype == torch.bfloat16 - # routing is untouched by the combine change - assert bool(ref_mask.sum(-1).eq(topk).all().item()), "top-k rows must select exactly topk blocks" - if gated: + if gate is None: + assert torch.equal(got, ref), "ungated full path must be bit-exact" + else: rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() assert rel <= 5e-3, f"relative L2 {rel:.3e} too large" - else: - assert torch.equal(got, ref), "ungated full path must be bit-exact" -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -def test_block_size_64_selects_triton_path(): - """block_size (4,4,4) must still route to the 64-block sparse kernel.""" - import fastvideo_kernel.ops as ops +def test_video_sparse_attn_inference_allocates_no_combine_temporaries(): + """Peak memory of a 64-tile inference call must drop by at least two full-sequence tensors. - called = {} - original = ops.block_sparse_attn + Only the 64-tile Triton kernel keeps the combine's temporaries at the top of + the peak: the 128/256 entries copy q/k/v into the kernel's layout first, so + their peak is set inside the kernel on both the Triton and CuTe routes. The + combine's own allocations on those paths are pinned by the combine-level + tests above. + """ + heads, n_blocks, dim, block_elements = 8, 16, 128, 64 + q, k, v, vbs, seq = _vsa_inputs(heads, n_blocks, dim, block_elements) + gate = _gate("bhsd", 1, heads, seq, dim) + block_size = (4, 4, 4) + full = q.numel() * q.element_size() + with torch.no_grad(): + # warm up Triton compilation outside the measurement + ops.video_sparse_attn(q, k, v, vbs, vbs, 4, block_size, compress_attn_weight=gate) + _reference_video_sparse_attn(q, k, v, vbs, vbs, 4, block_size, gate) + old_peak, _ = _peak_bytes(lambda: _reference_video_sparse_attn(q, k, v, vbs, vbs, 4, block_size, gate)) + new_peak, _ = _peak_bytes(lambda: ops.video_sparse_attn(q, k, v, vbs, vbs, 4, block_size, + compress_attn_weight=gate)) + assert old_peak - new_peak >= 2 * full, (f"expected at least two fewer full tensors ({2 * full} bytes), " + f"got old={old_peak} new={new_peak}") - def spy(*args, **kwargs): - called["hit"] = True - return original(*args, **kwargs) - q, k, v, vbs, w, _ = _vsa_inputs(4, 8, 64, gated=True) - ops.block_sparse_attn = spy - try: - with torch.no_grad(): - ops.video_sparse_attn(q, k, v, vbs, vbs, 4, (4, 4, 4), compress_attn_weight=w) - finally: - ops.block_sparse_attn = original - assert called.get("hit"), "64-block path did not call block_sparse_attn" +@pytest.mark.parametrize("gated", [False, True]) +def test_video_sparse_attn_training_matches_reference(gated): + """Grad-enabled 64-tile VSA: output and input gradients match the old path.""" + heads, n_blocks, dim = 4, 8, 64 + q, k, v, vbs, seq = _vsa_inputs(heads, n_blocks, dim, 64) + gate = _gate("bhsd" if gated else "none", 1, heads, seq, dim) + leaves_ref = [t.clone().requires_grad_(True) for t in (q, k, v) + ((gate, ) if gated else ())] + leaves_got = [t.clone().requires_grad_(True) for t in (q, k, v) + ((gate, ) if gated else ())] + + ref, _ = _reference_video_sparse_attn(*leaves_ref[:3], vbs, vbs, 4, (4, 4, 4), leaves_ref[3] if gated else None) + got = ops.video_sparse_attn(*leaves_got[:3], vbs, vbs, 4, (4, 4, 4), + compress_attn_weight=leaves_got[3] if gated else None) + torch.manual_seed(2) + grad_out = torch.randn_like(ref) + ref.backward(grad_out) + got.backward(grad_out) + if gated: + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel <= 5e-3 + else: + assert torch.equal(got, ref) + for name, a, b in zip(("q", "k", "v", "gate"), leaves_got, leaves_ref): + assert a.grad is not None and b.grad is not None + torch.testing.assert_close(a.grad, b.grad, rtol=2.0**-5, atol=2.0**-5, msg=lambda m: f"grad of {name}: {m}") -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -@pytest.mark.parametrize("block_elements", [128, 256]) -def test_128_256_paths_still_dispatch(block_elements): - """The 128/256 kernels must still be selected; the combine is shared.""" - import fastvideo_kernel.ops as ops +# -------------------------------------------------------------------------- +# End-to-end through video_sparse_attn_bshd (128/256-token tiles, [B, S, H, D]) +# -------------------------------------------------------------------------- - name = "block_sparse_attn_128" if block_elements == 128 else "block_sparse_attn_256" - called = {} - original = getattr(ops, name) - def spy(*args, **kwargs): - called["hit"] = True - return original(*args, **kwargs) +def _reference_video_sparse_attn_bshd(q, k, v, variable_block_sizes, q_variable_block_sizes, topk, block_size, + compress_attn_weight): + """Pre-change ``video_sparse_attn_bshd`` body.""" + from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128_bshd, block_sparse_attn_256_bshd + from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_topk_mask + block_elements = block_size[0] * block_size[1] * block_size[2] + batch, q_seq_len, heads, dim = q.shape + q_num_blocks = q_seq_len // block_elements + kv_num_blocks = k.shape[1] // block_elements - be = block_elements - n_blocks, heads, dim = 4, 4, 64 + q_c = q.view(batch, q_num_blocks, block_elements, heads, dim) + k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim) + v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim) + q_c = (q_c.float().sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype) + q_ch, k_ch, v_ch = (t.permute(0, 2, 1, 3).contiguous() for t in (q_c, k_c, v_c)) + + scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim**0.5) + attn = torch.softmax(scores, dim=-1) + out_c_blk = torch.matmul(attn, v_ch).permute(0, 2, 1, 3).contiguous() + + mask = fused_topk_mask(scores, topk) + attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd + out_s, _ = attention(q, k, v, mask, variable_block_sizes) + return _reference_combine_bshd(out_c_blk, out_s, compress_attn_weight, block_elements), mask + + +def _vsa_inputs_bshd(heads, n_blocks, dim, be, seed=0): + torch.manual_seed(seed) seq = n_blocks * be - torch.manual_seed(0) - q, k, v = (torch.randn(1, heads, seq, dim, device=DEVICE, dtype=torch.bfloat16) for _ in range(3)) + q, k, v = (torch.randn(1, seq, heads, dim, device=DEVICE, dtype=torch.bfloat16) for _ in range(3)) vbs = torch.full((n_blocks,), be, device=DEVICE, dtype=torch.int32) - bs = (be // 16, 4, 4) - assert bs[0] * bs[1] * bs[2] == be + return q, k, v, vbs, seq - setattr(ops, name, spy) - try: - with torch.no_grad(): - ops.video_sparse_attn(q, k, v, vbs, vbs, 2, bs) - except Exception as exc: # kernel unavailable on this backend - pytest.skip(f"{name} unavailable here: {type(exc).__name__}: {exc}") - finally: - setattr(ops, name, original) - assert called.get("hit"), f"{name} was not dispatched" + +@pytest.mark.parametrize("gated", [False, True]) +@pytest.mark.parametrize("block_elements", [128, 256]) +@pytest.mark.parametrize("heads,n_blocks,dim,ratio", [(4, 8, 64, 0.5), (8, 16, 128, 0.2)]) +def test_video_sparse_attn_bshd_matches_reference(monkeypatch, gated, block_elements, heads, n_blocks, dim, ratio): + """Inference BSHD VSA: same routing mask and same output as the old path.""" + q, k, v, vbs, seq = _vsa_inputs_bshd(heads, n_blocks, dim, block_elements) + torch.manual_seed(1) + gate = torch.rand(1, seq, heads, dim, device=DEVICE, dtype=torch.bfloat16) if gated else None + topk = max(1, int(n_blocks * ratio)) + block_size = (block_elements // 16, 4, 4) + + masks = [] + original_topk = ops.fused_topk_mask + monkeypatch.setattr(ops, "fused_topk_mask", lambda *a, **kw: masks.append(original_topk(*a, **kw)) or masks[-1]) + with torch.no_grad(): + ref, ref_mask = _reference_video_sparse_attn_bshd(q, k, v, vbs, vbs, topk, block_size, gate) + got = ops.video_sparse_attn_bshd(q, k, v, vbs, vbs, topk, block_size, compress_attn_weight=gate) + + assert len(masks) == 1 and torch.equal(masks[0], ref_mask), "top-k routing changed" + assert got.shape == ref.shape == (1, seq, heads, dim) and got.dtype == torch.bfloat16 + if gated: + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel <= 5e-3, f"relative L2 {rel:.3e} too large" + else: + assert torch.equal(got, ref), "ungated BSHD path must be bit-exact" + + +@pytest.mark.parametrize("gated", [False, True]) +def test_video_sparse_attn_bshd_training_matches_reference(gated): + """Grad-enabled BSHD VSA (256 tiles): output and input gradients match the old path.""" + heads, n_blocks, dim, be = 4, 8, 64, 256 + q, k, v, vbs, seq = _vsa_inputs_bshd(heads, n_blocks, dim, be) + torch.manual_seed(1) + gate = torch.rand(1, seq, heads, dim, device=DEVICE, dtype=torch.bfloat16) if gated else None + tensors = (q, k, v) + ((gate, ) if gated else ()) + leaves_ref = [t.clone().requires_grad_(True) for t in tensors] + leaves_got = [t.clone().requires_grad_(True) for t in tensors] + + ref, _ = _reference_video_sparse_attn_bshd(*leaves_ref[:3], vbs, vbs, 4, (16, 4, 4), + leaves_ref[3] if gated else None) + got = ops.video_sparse_attn_bshd(*leaves_got[:3], vbs, vbs, 4, (16, 4, 4), + compress_attn_weight=leaves_got[3] if gated else None) + torch.manual_seed(2) + grad_out = torch.randn_like(ref) + ref.backward(grad_out) + got.backward(grad_out) + if gated: + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel <= 5e-3 + else: + assert torch.equal(got, ref) + for name, a, b in zip(("q", "k", "v", "gate"), leaves_got, leaves_ref): + assert a.grad is not None and b.grad is not None + torch.testing.assert_close(a.grad, b.grad, rtol=2.0**-5, atol=2.0**-5, msg=lambda m: f"grad of {name}: {m}") From 28c68b9d43ded3ff7e026807b832d7738bfc55eb Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 5 Sep 2026 00:39:19 +0000 Subject: [PATCH 3/4] [perf]: hand the VSA gate to video_sparse_attn as a transposed view The 64-tile BHSD path copied gate_compress into a contiguous BHSD tensor on every call. The coarse/sparse combine views the gate at block resolution, which never needs contiguity, so the transposed BSHD view is passed as is. This saves one full-sequence copy per attention layer with bit-identical output. On a GB200 at 1x39936x12x128 (Wan 1.3B, 480p) the per-call peak drops from 639.8 to 521.8 MiB and latency from 11.64 to 11.29 ms. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NjFimeedTQWxgqSzP5xo4j --- fastvideo/attention/backends/video_sparse_attn.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index f90e70e542..91a7fbfb2b 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -327,11 +327,14 @@ def forward( # type: ignore[override] if video_sparse_attn is None: raise NotImplementedError("video_sparse_attn is not installed") - # Default 64-element-tile path (unchanged): BHSD round-trip. + # Default 64-element-tile path (unchanged): BHSD round-trip. The gate is + # only read elementwise by the coarse/sparse combine, which views it at + # block resolution without requiring contiguity, so the transposed view + # is handed over as is (no full-sequence copy). query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() value = value.transpose(1, 2).contiguous() - gate_compress = gate_compress.transpose(1, 2).contiguous() + gate_compress = gate_compress.transpose(1, 2) return video_sparse_attn(query, key, value, From 38412b2841fa2372a66e9ffcd9adc02172ba1e92 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 5 Sep 2026 00:39:19 +0000 Subject: [PATCH 4/4] [perf]: fuse the MiniMax H3 VSA gate combine, accumulate in place at inference torch.addcmul keeps one full-sequence temporary under grad instead of two, and when the sparse output is not tracked by autograd the gate branch accumulates into it directly, allocating nothing. The product is no longer rounded to bf16 before the add, so gated values move by at most half a bf16 ulp (towards the fp32 value). test_vsa_h3_backward.py passes on the Triton and FA4 CuTe backends on a GB200. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NjFimeedTQWxgqSzP5xo4j --- .../attention/backends/video_sparse_attn_h3.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index dcfd54f3a9..844ea9f366 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -763,11 +763,18 @@ def forward( # type: ignore[override] out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled) # [B, H, n_tiles, D] out_c = out_c.permute(0, 2, 1, 3).to(out.dtype) # [B, n_tiles, H, D] batch, seq_len, heads, dim = out.shape - # Out-of-place: on the CuTe backend ``out`` is the tensor FA4's - # autograd node saved for its backward, so an in-place add here - # bumps its version counter and backward dies with "one of the - # variables needed for gradient computation has been modified". out_tiled = out.view(batch, n_tiles, tile_elems, heads, dim) gate_tiled = logical_gate.view(batch, n_tiles, tile_elems, heads, dim) - out = (out_tiled + out_c.unsqueeze(2) * gate_tiled).view(batch, seq_len, heads, dim) + out_c = out_c.unsqueeze(2) + if out.requires_grad or gate_tiled.dtype != out.dtype: + # Out-of-place: ``out`` is the tensor the attention kernel's + # autograd node saved for its backward, so an in-place add here + # bumps its version counter and backward dies with "one of the + # variables needed for gradient computation has been modified". + # Fused addcmul keeps one full-sequence temporary instead of two. + out = torch.addcmul(out_tiled, out_c, gate_tiled).view(batch, seq_len, heads, dim) + else: + # Inference: ``out`` is a fresh kernel output nobody else holds, + # so accumulate into it and allocate nothing. + out_tiled.addcmul_(out_c, gate_tiled) return out