[bugfix] Fix extra work in VSA64 layout benchmark (#2027)

fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py

@@ -5,7 +5,7 @@ Benchmark the SM100 fused VSA64 tile layout against the legacy scatter/transpose
 Times Wan VSA inference preprocessing through the backend
 (``VideoSparseAttentionImpl.preprocess_qkv``), toggling
 ``FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT``:
-  - legacy: tile scatter into BSHD + four BSHD->BHSD transposes
+  - legacy: tile scatter into BSHD + three contiguous Q/K/V transposes + a gate view
   - fused:  one Triton kernel writing tiled BHSD directly
 
 With ``--include_attn`` the timed region also covers ``forward`` (the
@@ -59,20 +59,20 @@ def parse_arguments() -> argparse.Namespace:
 
 
 def make_step(impl: VideoSparseAttentionImpl, qkvg: torch.Tensor, builder: VideoSparseAttentionMetadataBuilder,
-              build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], torch.Tensor]:
+              build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], tuple[torch.Tensor, ...]]:
     metadata = builder.build(**build_kwargs)
 
-    def step() -> torch.Tensor:
+    def step() -> tuple[torch.Tensor, ...]:
         # The routing gate reads the switch on every call.
         with envs.FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT.override(not fused):
             tiled = impl.preprocess_qkv(qkvg, metadata)
         q, k, v, gate = tiled.chunk(4)
         if include_attn:
-            return impl.forward(q, k, v, gate, metadata)
+            return (impl.forward(q, k, v, gate, metadata), )
         if metadata.fused_layout_active:
-            return tiled
-        # Legacy route: the transposes forward() applies before video_sparse_attn.
-        return torch.cat([t.transpose(1, 2).contiguous() for t in (q, k, v, gate)])
+            return q, k, v, gate
+        return (q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(),
+                gate.transpose(1, 2))
 
     with torch.no_grad():
         step()
@@ -100,7 +100,7 @@ def main() -> None:
     legacy = make_step(impl, qkvg, builder, build_kwargs, fused=False, include_attn=args.include_attn)
     fused = make_step(impl, qkvg, builder, build_kwargs, fused=True, include_attn=args.include_attn)
     with torch.no_grad():
-        if not torch.equal(legacy(), fused()):
+        if not all(torch.equal(old, new) for old, new in zip(legacy(), fused(), strict=True)):
             raise RuntimeError("Fused and legacy outputs differ.")
         legacy_ms = do_bench(legacy, warmup=args.warmup, rep=args.rep, return_mode="median")
         fused_ms = do_bench(fused, warmup=args.warmup, rep=args.rep, return_mode="median")

fastvideo-kernel/tests/test_vsa_tile_layout.py

@@ -5,6 +5,9 @@ Production only routes SM100 here, but the kernel is plain Triton, so parity is
 checked on any CUDA device.
 """
 
+import importlib.util
+from pathlib import Path
+
 import pytest
 import torch
 
@@ -40,6 +43,63 @@ def test_tile_to_bhsd_matches_scatter_and_four_transposes(dit_shape):
     source_index[metadata.non_pad_index] = metadata.tile_partition_indices.to(torch.int32)
     out = torch.empty((4, 2, padded_sequence, 128), device="cuda", dtype=torch.bfloat16)
     actual = tile_to_bhsd(qkvg, source_index, out, sequence, padded_sequence, 2, 128).chunk(4)
-    for name, old, new in zip(("q", "k", "v", "gate"), expected, actual):
+    for name, old, new in zip(("q", "k", "v", "gate"), expected, actual, strict=True):
         assert old.shape == new.shape, name
         assert torch.equal(old, new), name
+
+
+@pytest.mark.parametrize("fused", [False, True])
+def test_layout_benchmark_matches_forward_preprocessing(monkeypatch, fused):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA device required")
+
+    from fastvideo.attention.backends import video_sparse_attn as vsa
+
+    path = Path(__file__).resolve().parents[1] / "benchmarks/bench_vsa_tile_layout.py"
+    spec = importlib.util.spec_from_file_location("bench_vsa_tile_layout", path)
+    benchmark = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(benchmark)
+    original_gate = vsa._can_fuse_vsa64_layout
+
+    def allow_test_device(x, metadata):
+        with monkeypatch.context() as gate_patch:
+            gate_patch.setattr(torch.cuda, "get_device_capability", lambda *args: (10, 0))
+            return original_gate(x, metadata)
+
+    monkeypatch.setattr(vsa, "_can_fuse_vsa64_layout", allow_test_device)
+    impl = object.__new__(vsa.VideoSparseAttentionImpl)
+    builder = vsa.VideoSparseAttentionMetadataBuilder()
+    kwargs = dict(current_timestep=0,
+                  raw_latent_shape=(5, 5, 6),
+                  patch_size=(1, 1, 1),
+                  VSA_sparsity=0.8,
+                  device=torch.device("cuda"),
+                  cache_tile_buf=True)
+    qkvg = torch.randn((4, 157, 2, 128), device="cuda", dtype=torch.bfloat16)[:, :150]
+    captured = []
+
+    def capture_attention_inputs(q, k, v, *args, **kwargs):
+        captured[:] = [q, k, v, kwargs["compress_attn_weight"]]
+        return q
+
+    monkeypatch.setattr(vsa, "video_sparse_attn", capture_attention_inputs)
+    with torch.no_grad():
+        layout_step = benchmark.make_step(impl, qkvg, builder, kwargs, fused=fused, include_attn=False)
+        attention_step = benchmark.make_step(impl, qkvg, builder, kwargs, fused=fused, include_attn=True)
+
+        def reject_cat(*args, **kwargs):
+            pytest.fail("The timed layout step must not concatenate its outputs")
+
+        with monkeypatch.context() as timed_patch:
+            timed_patch.setattr(torch, "cat", reject_cat)
+            layout = layout_step()
+            attention = attention_step()
+
+    assert isinstance(layout, tuple) and len(layout) == 4
+    for actual, expected in zip(layout, captured, strict=True):
+        assert torch.equal(actual, expected)
+        assert actual.stride() == expected.stride()
+    assert all(t.is_contiguous() for t in layout[:3])
+    assert layout[3].is_contiguous() is fused
+    assert isinstance(attention, tuple) and len(attention) == 1
+    assert torch.equal(attention[0], captured[0].transpose(1, 2))