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

BaseLayerWithLoRA.forward was decorated with torch.compile, so wrapping a linear also compiled its base layer.
Compiling a packed NVFP4/FP8 linear on its own changes its numerics, so a wrapped layer at strength 0 could differ
from the unwrapped layer. forward now adds no compiled region: the base layer runs exactly as it would unwrapped
(eagerly, or traced with a caller's compiled block), and only the adapter delta is compiled.

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))

fastvideo/layers/lora/linear.py

@@ -155,37 +155,43 @@ class BaseLayerWithLoRA(nn.Module):
         """The wrapped layer's bias, or ``None`` when it has none."""
         return getattr(self.base_layer, "bias", None)
 
-    @torch.compile()
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        lora_A = self.lora_A
-        lora_B = self.lora_B
-        if isinstance(self.lora_B, DTensor):
-            lora_B = self.lora_B.to_local()
-            lora_A = self.lora_A.to_local()
-
-        if self.side_branch and (self.lora_A is None or self._strength_is_zero):
-            # No adapter, or strength 0, is exactly the unadapted layer, down to the output dtype.
+        # The wrapper adds no compiled region around the base layer, so it runs exactly as it would unwrapped: eagerly,
+        # or inside a caller's compiled block. Compiling a packed NVFP4/FP8 linear on its own changes its numerics, and
+        # a wrapped layer with no adapter or strength 0 must equal the unwrapped layer exactly.
+        if self.lora_A is None or (self.side_branch and self._strength_is_zero):
+            # No adapter, or strength 0 on a side branch, is exactly the unadapted layer, down to the output dtype.
+            # A merge-mode wrapper reaches this branch while its first adapter is still loading or after a load
+            # failed: the pipeline wraps every linear before it reads the weights, so an unadapted request must
+            # not depend on an adapter that is not there.
             return self.base_layer(x)
 
-        if not self.merged and not self.disable_lora:
-            lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
-            lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
-            alpha_over_rank = 1.0
-            if self.lora_alpha != self.lora_rank:
-                alpha_over_rank = float(self.lora_alpha) / float(self.lora_rank)  # type: ignore[arg-type]
-            out, output_bias = self.base_layer(x)
-            if self.side_branch:
-                # Keep chunk temporaries opaque to Inductor so fullgraph capture cannot retain a full-size delta.
-                _add_low_rank_inplace(out, x, lora_A_sliced, lora_B_sliced, self._side_branch_strength, alpha_over_rank)
-                return out, output_bias
-            scale = float(self.lora_strength) * alpha_over_rank
-            delta = x @ lora_A_sliced.transpose(0, 1) @ lora_B_sliced.transpose(0, 1)
-            if scale != 1.0:
-                delta = delta * scale
-            return out + delta, output_bias
-        else:
-            out, output_bias = self.base_layer(x)
+        out, output_bias = self.base_layer(x)
+        if self.merged or self.disable_lora:
             return out.to(x), output_bias
+        return self._add_adapter(out, x), output_bias
+
+    @torch.compile()
+    def _add_adapter(self, out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        lora_A = self.lora_A
+        lora_B = self.lora_B
+        if isinstance(lora_B, DTensor):
+            lora_B = lora_B.to_local()
+            lora_A = lora_A.to_local()
+        lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
+        lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
+        alpha_over_rank = 1.0
+        if self.lora_alpha != self.lora_rank:
+            alpha_over_rank = float(self.lora_alpha) / float(self.lora_rank)  # type: ignore[arg-type]
+        if self.side_branch:
+            # Keep chunk temporaries opaque to Inductor so fullgraph capture cannot retain a full-size delta.
+            _add_low_rank_inplace(out, x, lora_A_sliced, lora_B_sliced, self._side_branch_strength, alpha_over_rank)
+            return out
+        scale = float(self.lora_strength) * alpha_over_rank
+        delta = x @ lora_A_sliced.transpose(0, 1) @ lora_B_sliced.transpose(0, 1)
+        if scale != 1.0:
+            delta = delta * scale
+        return out + delta
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A

fastvideo/tests/layers/test_lora_side_branch.py

@@ -54,6 +54,89 @@ def _packed_linear(seed: int = 0) -> ReplicatedLinear:
     return layer
 
 
+def test_a_wrapped_layer_without_an_adapter_runs_the_base_layer() -> None:
+    """A dense wrapper with no adapter is the base layer.
+
+    The pipeline wraps every linear when the first runtime adapter is set. A render that then runs without one, or a
+    load that failed part-way, must still produce the base model's output instead of failing on the missing factors.
+    """
+    base = _dense_linear()
+    wrapped = get_lora_layer(base)
+    assert not wrapped.side_branch and wrapped.lora_A is None
+    x = torch.randn(2, IN, dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert torch.equal(wrapped(x)[0], base(x)[0])
+
+
+@pytest.mark.parametrize("packed", [False, True])
+def test_wrapped_layer_at_strength_zero_is_bitwise_the_base_layer(packed: bool) -> None:
+    base = _packed_linear() if packed else _dense_linear()
+    x = _x()
+    with torch.no_grad():
+        # Taken before wrapping, so a merge that wrongly rewrites the shared base weight cannot also change it.
+        expected = base(x)[0].clone()
+    layer = get_lora_layer(base)
+    layer.set_lora_weights(*_factors(), lora_alpha=8, strength=0.0)
+    assert layer.lora_A is not None and layer.lora_strength == 0.0
+    with torch.no_grad():
+        assert torch.equal(layer(x)[0], expected)
+
+
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+@pytest.mark.parametrize("adapter", [False, True])
+def test_wrapping_does_not_change_a_compiled_block(monkeypatch: pytest.MonkeyPatch, backend: str,
+                                                   adapter: bool) -> None:
+    """Under a caller's compile the base layer is traced with the block, as it is unwrapped; the wrapper adds no
+    region of its own, so the compiled block is unchanged with no adapter or at strength 0."""
+    if backend == "inductor" and not torch.cuda.is_available():
+        pytest.skip("Inductor execution requires a CUDA GPU for this regression")
+    device = torch.device("cuda" if backend == "inductor" else "cpu")
+    monkeypatch.setattr(torch._dynamo.config, "disable", False)
+    torch._dynamo.reset()
+
+    class Block(torch.nn.Module):
+
+        def __init__(self, proj: torch.nn.Module) -> None:
+            super().__init__()
+            self.proj = proj
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.proj(x)[0] + x[..., :OUT]
+
+    x = _x().to(device)
+    try:
+        with torch.no_grad():
+            expected = torch.compile(Block(_packed_linear().to(device)), backend=backend, fullgraph=True)(x)
+            torch._dynamo.reset()
+            layer = get_lora_layer(_packed_linear().to(device))
+            if adapter:
+                layer.set_lora_weights(*(factor.to(device) for factor in _factors()), lora_alpha=8, strength=0.0)
+            actual = torch.compile(Block(layer), backend=backend, fullgraph=True)(x)
+        assert torch.equal(actual, expected)
+    finally:
+        torch._dynamo.reset()
+
+
+def test_wrapper_adds_no_compiled_region_around_the_base_layer(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Compiling a packed base layer on its own changes its numerics, so an uncompiled caller runs it eagerly."""
+    monkeypatch.setattr(torch._dynamo.config, "disable", False)
+    torch._dynamo.reset()
+    base = _packed_linear()
+    layer = get_lora_layer(base)
+    layer.set_lora_weights(*_factors(), lora_alpha=8, strength=0.5)
+    seen: list[bool] = []
+    original = base.forward
+
+    def spy(x: torch.Tensor):
+        seen.append(torch.compiler.is_compiling())
+        return original(x)
+
+    monkeypatch.setattr(base, "forward", spy)
+    with torch.no_grad():
+        layer(_x())
+    assert seen == [False]
+
+
 def _factors(seed: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
     generator = torch.Generator().manual_seed(seed)
     return torch.randn(RANK, IN, generator=generator), torch.randn(OUT, RANK, generator=generator)