[perf] avoid full-sequence materialization in VSA coarse/sparse combine

fastvideo-kernel/benchmarks/bench_vsa_combine.py

@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: Apache-2.0
+"""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: 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
+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:
+        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 _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 = torch.cuda.memory_allocated()
+    fn()
+    torch.cuda.synchronize()
+    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 _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, 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
+    # 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():
+        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 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))
+
+    cases = [("small", 1, 8, 64, 128)]
+    if args.large:
+        # 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__":
+    main()

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -115,9 +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)
-    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 +129,74 @@ def video_sparse_attn(
     else:
         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
-    return out_c + out_s
+    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,
+    block_elements: int,
+    seq_dim: int,
+) -> torch.Tensor:
+    """Combine the block-resolution coarse output with the sparse output.
+
+    ``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.
+    """
+    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)
+
+    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:
+            out_s_b.add_(out_c)
+        return out_s
+
+    if gate_b is not None:
+        combined = torch.addcmul(out_s_b, out_c, gate_b)
+    else:
+        combined = out_s_b + out_c
+    return combined.view(*full)
 
 
 def video_sparse_attn_bshd(
@@ -191,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)