[perf] Fuse VSA64 tile layout for SM100 inference (#1900)

For Wan VSA inference with 64-token tiles, each attention layer scattered Q, K,
V and the gate into a tile buffer and then made four transposed copies. On
SM100, one Triton kernel (tile_to_bhsd) now does the scatter and the transposes
in one pass. Training, 256-token tiles, other GPUs and builds without the kernel
use the previous path. FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT=1 turns the fused
path off.

On GB200 the layout and attention step goes from 15.7 ms to 8.6 ms. FastWan 1.3B
frames are bit-exact with the previous path at sequence parallel 1 and 2.

docs/attention/vsa/index.md

@@ -50,6 +50,28 @@ tile and head, including dense overrides. `achieved_sparsity` is the matching
 mean video-tile sparsity; in `compete` mode it measures actual selections,
 not the requested top-k budget. These are tile counts, not token-level FLOPs.
 
+## SM100 fused 64-token layout (Wan inference)
+
+On SM100 (B200) with 64-token tiles, BF16, and head dimension 128, Wan VSA
+inference fuses tile scatter and the four BSHD→BHSD transposes into one Triton
+kernel in `fastvideo-kernel`. Set `FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT=1` to
+force the legacy scatter/transpose path. Training, other tile sizes, and non-SM100
+hardware always use the original route. The gate matches compute capability
+10.0 exactly; other SM10x parts (e.g. SM103) are not yet validated and use the
+original route.
+
+The kernel ships in the first `fastvideo-kernel` release after 0.3.5. Standard
+installs pin `fastvideo-kernel==0.3.5`, which does not include it, so they
+silently use the original route. Until that release, build `fastvideo-kernel`
+from source (`cd fastvideo-kernel && ./build.sh`) to get the fused layout.
+
+To compare the two routes on a Wan 480p token grid (add `--include_attn` to
+also time the sparse attention call):
+
+```bash
+python fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py
+```
+
 ## Usage
 
 ```python

docs/contributing/env_vars.md

@@ -201,6 +201,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_DISABLE_ATTENTION_COMPILE`                      | bool                           | `1`                                                   | attention   | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile.                                            |
 | `FASTVIDEO_MLX_WINDOW`                                     | int                            | `0`                                                   | attention   | MLX FastWan windowed attention size in tokens. 0 uses full attention.                                                                                                                                                                                      |
 | `FASTVIDEO_MLX_WINDOW_SINK`                                | int                            | `0`                                                   | attention   | Number of sink tokens that MLX windowed attention always attends to.                                                                                                                                                                                       |
+| `FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT`                     | bool                           | `0`                                                   | attention   | VIDEO_SPARSE_ATTN keeps the original tile scatter and BSHD->BHSD transposes instead of the fused Triton layout kernel that no-grad SM100 BF16 head_dim-128 forwards with 64-token tiles use by default.                                                 |
 | `FASTVIDEO_INFERENCE_TORCH_COMPILE`                        | bool                           | `0`                                                   | performance | Compile each DiT transformer block with fullgraph torch.compile at inference. Same as FastVideoArgs.inference_torch_compile=True.                                                                                                                          |
 | `FASTVIDEO_VAE_PARALLEL_DECODE`                            | bool                           | `0`                                                   | performance | MiniMax-H3 VAE decode splits its temporal chunks across the sequence-parallel ranks instead of running serially on the output rank. Same as FastVideoArgs.vae_parallel_decode=True.                                                                        |
 | `FASTVIDEO_VAE_PARALLEL_ENCODE`                            | bool                           | `0`                                                   | performance | MiniMax-H3 reference-video VAE encode splits its temporal chunks across the sequence-parallel ranks. Same as FastVideoArgs.vae_parallel_encode=True.                                                                                                       |

fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""
+Benchmark the SM100 fused VSA64 tile layout against the legacy scatter/transpose route.
+
+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
+  - fused:  one Triton kernel writing tiled BHSD directly
+
+With ``--include_attn`` the timed region also covers ``forward`` (the
+``video_sparse_attn`` call). Outputs of both routes are checked for bitwise
+equality before timing. The fused route only activates on SM100.
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+from collections.abc import Callable
+
+import numpy as np
+import torch
+
+try:
+    from triton.testing import do_bench
+except Exception as e:  # pragma: no cover
+    raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
+
+import fastvideo.envs as envs
+from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionImpl,
+                                                            VideoSparseAttentionMetadataBuilder)
+
+
+def set_seed(seed: int = 42) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def parse_arguments() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Benchmark the SM100 fused VSA64 tile layout")
+    p.add_argument("--dit_shape",
+                   type=int,
+                   nargs=3,
+                   default=[21, 30, 52],
+                   metavar=("T", "H", "W"),
+                   help="Post-patchify token grid (default: Wan 480p, 81 frames)")
+    p.add_argument("--batch_size", type=int, default=1)
+    p.add_argument("--num_heads", type=int, default=40, help="Heads per rank")
+    p.add_argument("--sparsity", type=float, default=0.9)
+    p.add_argument("--include_attn", action="store_true", help="Also time the video_sparse_attn forward")
+    p.add_argument("--warmup", type=int, default=200, help="Warmup time in ms (triton do_bench)")
+    p.add_argument("--rep", type=int, default=2000, help="Measurement time in ms (triton do_bench)")
+    p.add_argument("--seed", type=int, default=42)
+    return p.parse_args()
+
+
+def make_step(impl: VideoSparseAttentionImpl, qkvg: torch.Tensor, builder: VideoSparseAttentionMetadataBuilder,
+              build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], torch.Tensor]:
+    metadata = builder.build(**build_kwargs)
+
+    def step() -> 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)
+        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)])
+
+    with torch.no_grad():
+        step()
+    if fused and not metadata.fused_layout_active:
+        raise RuntimeError("Fused layout did not activate (requires SM100, bf16, head_dim 128, "
+                           "and fastvideo_kernel.triton_kernels.vsa_tile_layout).")
+    return step
+
+
+def main() -> None:
+    args = parse_arguments()
+    set_seed(args.seed)
+
+    impl = object.__new__(VideoSparseAttentionImpl)
+    builder = VideoSparseAttentionMetadataBuilder()
+    build_kwargs = dict(current_timestep=0,
+                        raw_latent_shape=tuple(args.dit_shape),
+                        patch_size=(1, 1, 1),
+                        VSA_sparsity=args.sparsity,
+                        device=torch.device("cuda"),
+                        cache_tile_buf=True)
+    sequence = int(np.prod(args.dit_shape))
+    qkvg = torch.randn((4 * args.batch_size, sequence, args.num_heads, 128), device="cuda", dtype=torch.bfloat16)
+
+    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()):
+            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")
+
+    region = "layout+attn" if args.include_attn else "layout"
+    print(f"GPU: {torch.cuda.get_device_name()} | dit_shape={tuple(args.dit_shape)} seq={sequence} "
+          f"B={args.batch_size} H={args.num_heads} D=128 sparsity={args.sparsity}")
+    print(f"{region} (median): legacy {legacy_ms:.3f} ms | fused {fused_ms:.3f} ms | "
+          f"speedup {legacy_ms / fused_ms:.2f}x")
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark.")
+    main()

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/vsa_tile_layout.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused VSA64 tile scatter and BSHD->BHSD layout conversion for SM100 inference."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _tile_to_bhsd_kernel(
+    x_ptr,
+    source_index_ptr,
+    output_ptr,
+    sequence,
+    padded_sequence,
+    heads,
+    head_dim: tl.constexpr,
+    stride_batch,
+    stride_sequence,
+    stride_head,
+    stride_dim,
+    block: tl.constexpr,
+):
+    pid0 = tl.program_id(0).to(tl.int64)
+    batch_head = tl.program_id(1).to(tl.int64)
+    linear = pid0 * block + tl.arange(0, block).to(tl.int64)
+    batch = batch_head // heads
+    head = batch_head % heads
+    padded_pos = linear // head_dim
+    dim = linear % head_dim
+    valid = linear < padded_sequence * head_dim
+    source_pos = tl.load(source_index_ptr + padded_pos, mask=valid, other=-1).to(tl.int64)
+    x_offset = (batch * stride_batch + source_pos * stride_sequence + head * stride_head + dim * stride_dim)
+    value = tl.load(x_ptr + x_offset, mask=valid & (source_pos >= 0) & (source_pos < sequence), other=0)
+    out_offset = batch_head * padded_sequence * head_dim + linear
+    tl.store(output_ptr + out_offset, value, mask=valid)
+
+
+def tile_to_bhsd(
+    x: torch.Tensor,
+    source_index: torch.Tensor,
+    out: torch.Tensor,
+    sequence: int,
+    padded_sequence: int,
+    heads: int,
+    head_dim: int,
+) -> torch.Tensor:
+    """Scatter-tile ``x`` [B, S, H, D] into ``out`` [B, H, padded_S, D] in BHSD order."""
+    block = 1024
+    grid = (triton.cdiv(padded_sequence * head_dim, block), x.shape[0] * heads)
+    _tile_to_bhsd_kernel[grid](
+        x,
+        source_index,
+        out,
+        sequence,
+        padded_sequence,
+        heads,
+        head_dim,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        x.stride(3),
+        block,
+        num_warps=4,
+    )
+    return out

fastvideo-kernel/tests/test_vsa_tile_layout.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bitwise parity for the fused VSA64 tile-to-BHSD Triton kernel.
+
+Production only routes SM100 here, but the kernel is plain Triton, so parity is
+checked on any CUDA device.
+"""
+
+import pytest
+import torch
+
+pytest.importorskip("triton")
+
+from fastvideo_kernel.triton_kernels.vsa_tile_layout import tile_to_bhsd  # noqa: E402
+
+
+@pytest.mark.parametrize("dit_shape", [(5, 5, 6), (16, 28, 52)])
+def test_tile_to_bhsd_matches_scatter_and_four_transposes(dit_shape):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA device required")
+
+    from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionImpl,
+                                                                VideoSparseAttentionMetadataBuilder)
+
+    impl = object.__new__(VideoSparseAttentionImpl)
+    metadata = VideoSparseAttentionMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=dit_shape,
+        patch_size=(1, 1, 1),
+        VSA_sparsity=0.8,
+        device=torch.device("cuda"),
+        cache_tile_buf=True,
+    )
+    sequence = metadata.total_seq_length
+    padded_sequence = metadata.variable_block_sizes.numel() * 64
+    source = torch.randn((4, sequence + 7, 2, 128), device="cuda", dtype=torch.bfloat16)
+    qkvg = source[:, :sequence]
+    expected = tuple(piece.transpose(1, 2).contiguous() for piece in impl.tile(qkvg, metadata).chunk(4))
+
+    source_index = torch.full((padded_sequence, ), -1, device="cuda", dtype=torch.int32)
+    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):
+        assert old.shape == new.shape, name
+        assert torch.equal(old, new), name

fastvideo/attention/backends/video_sparse_attn.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import functools
 import math
+from collections.abc import Callable
 from dataclasses import dataclass
 
 import torch
@@ -16,6 +17,7 @@ except ImportError:
 
 from typing import Any
 
+import fastvideo.envs as envs
 from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
                                                    AttentionMetadataBuilder)
 from fastvideo.distributed import get_sp_group
@@ -156,6 +158,8 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
     # can release the large tiled QKVG scratch tensor after each attention call.
     tile_buf: torch.Tensor | None = None
     cache_tile_buf: bool = True
+    fused_tile_source_index: torch.Tensor | None = None
+    fused_layout_active: bool = False
 
 
 def compute_topk(sparsity: float, num_blocks: int) -> int:
@@ -199,6 +203,27 @@ def scatter_into_tile_buf(
     return buf
 
 
+@functools.cache
+def _get_tile_to_bhsd() -> Callable[..., Any] | None:
+    try:
+        from fastvideo_kernel.triton_kernels.vsa_tile_layout import tile_to_bhsd
+    except ImportError:
+        return None
+    return tile_to_bhsd
+
+
+def _can_fuse_vsa64_layout(x: torch.Tensor, metadata: VideoSparseAttentionMetadata) -> bool:
+    """Restrict the fused layout to the validated SM100 BF16 inference route."""
+    if envs.FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT.get():
+        return False
+    if _get_tile_to_bhsd() is None:
+        return False
+    return (math.prod(VSA_TILE_SIZE) == 64 and metadata.cache_tile_buf and not torch.is_grad_enabled()
+            and not x.requires_grad and x.is_cuda and x.dtype == torch.bfloat16 and x.ndim == 4
+            and x.shape[1] == metadata.total_seq_length and x.shape[-1] == 128
+            and torch.cuda.get_device_capability(x.device) == (10, 0))
+
+
 class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
 
     def __init__(self) -> None:
@@ -268,11 +293,11 @@ class VideoSparseAttentionImpl(AttentionImpl):
         valid until the next ``tile()`` / ``preprocess_qkv`` call on the
         same ``attn_metadata``.  Callers must consume (or copy) the
         result before invoking another VSA layer with the same metadata.
-        Today ``forward()`` materializes copies of q/k/v via
-        ``.transpose(...).contiguous()`` and consumes the gate view
-        synchronously inside the same call (the coarse/sparse combine
-        reads it before returning), so the contract holds; future
-        callers must preserve it.
+        The default 64-token path materializes q/k/v copies in ``forward()``
+        and consumes its transposed gate view synchronously in the coarse/sparse
+        combine before returning. The SM100 inference path instead writes BHSD
+        directly to this buffer. Both paths consume the result before the next
+        layer reuses it; future callers must preserve this contract.
         """
         num_tiles = attn_metadata.num_tiles
         t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
@@ -305,7 +330,33 @@ class VideoSparseAttentionImpl(AttentionImpl):
         attn_metadata: VideoSparseAttentionMetadata,
     ) -> torch.Tensor:
         """Tile QKV; aliasing contract: see ``tile()``."""
-        return self.tile(qkv, attn_metadata)
+        attn_metadata.fused_layout_active = _can_fuse_vsa64_layout(qkv, attn_metadata)
+        if not attn_metadata.fused_layout_active:
+            return self.tile(qkv, attn_metadata)
+
+        tile_to_bhsd = _get_tile_to_bhsd()
+        assert tile_to_bhsd is not None
+
+        padded_sequence = attn_metadata.variable_block_sizes.numel() * math.prod(VSA_TILE_SIZE)
+        if attn_metadata.fused_tile_source_index is None:
+            source_index = torch.full((padded_sequence, ), -1, device=qkv.device, dtype=torch.int32)
+            source_index[attn_metadata.non_pad_index] = attn_metadata.tile_partition_indices.to(torch.int32)
+            attn_metadata.fused_tile_source_index = source_index
+        target_shape = (qkv.shape[0], qkv.shape[2], padded_sequence, qkv.shape[3])
+        buffer = attn_metadata.tile_buf
+        if buffer is None or buffer.shape != target_shape or buffer.dtype != qkv.dtype or buffer.device != qkv.device:
+            buffer = torch.empty(target_shape, device=qkv.device, dtype=qkv.dtype)
+            attn_metadata.tile_buf = buffer
+        tile_to_bhsd(
+            qkv,
+            attn_metadata.fused_tile_source_index,
+            buffer,
+            qkv.shape[1],
+            padded_sequence,
+            qkv.shape[2],
+            qkv.shape[3],
+        )
+        return buffer
 
     def postprocess_output(
         self,
@@ -339,14 +390,17 @@ class VideoSparseAttentionImpl(AttentionImpl):
 
         if video_sparse_attn is None:
             raise NotImplementedError("video_sparse_attn is not installed")
-        # 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)
+        # The SM100 inference preprocessing path already emits BHSD. Other
+        # shapes, architectures, training, and the benchmark disable switch
+        # retain the original BSHD scatter/transpose path.
+        padded_sequence = attn_metadata.variable_block_sizes.numel() * block_elements
+        if not (attn_metadata.fused_layout_active and query.shape[2] == padded_sequence):
+            query = query.transpose(1, 2).contiguous()
+            key = key.transpose(1, 2).contiguous()
+            value = value.transpose(1, 2).contiguous()
+            # The coarse/sparse combine only reads the gate elementwise at
+            # block resolution, so its transposed view needs no contiguous copy.
+            gate_compress = gate_compress.transpose(1, 2)
         return video_sparse_attn(query,
                                  key,
                                  value,

fastvideo/envs.py

@@ -378,6 +378,11 @@ FASTVIDEO_MLX_WINDOW = EnvInt(0,
 FASTVIDEO_MLX_WINDOW_SINK = EnvInt(0,
                                    category="attention",
                                    doc="Number of sink tokens that MLX windowed attention always attends to.")
+FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT = EnvBool(
+    False,
+    category="attention",
+    doc="VIDEO_SPARSE_ATTN keeps the original tile scatter and BSHD->BHSD transposes instead of the fused Triton "
+    "layout kernel that no-grad SM100 BF16 head_dim-128 forwards with 64-token tiles use by default.")
 
 # ================== Performance ==================
 

fastvideo/pipelines/stages/denoising.py

@@ -1231,7 +1231,7 @@ class Cosmos25DistilledT2WDenoisingStage(Cosmos25DenoisingStage):
         else:
             fps_tensor = torch.as_tensor(batch.fps, device=latents.device, dtype=target_dtype).reshape(-1)
             if fps_tensor.numel() != 1:
-                # The Cosmos25 RoPE builds one batch-independent (T*H*W, D) table
+                # The Cosmos25 RoPE builds one batch-independent (T * H * W, D) table
                 # (see cosmos2_5.py), so the DiT consumes a single shared fps
                 # value per forward pass, never one value per sample.
                 if not torch.equal(fps_tensor, fps_tensor[:1].expand_as(fps_tensor)):
@@ -1355,7 +1355,7 @@ class Cosmos25DFDV2WDenoisingStage(Cosmos25DenoisingStage):
         fps_value = 24 if batch.fps is None else batch.fps
         fps_tensor = torch.as_tensor(fps_value, device=state.device, dtype=torch.float32).reshape(-1)
         if fps_tensor.numel() != 1:
-            # The Cosmos25 RoPE builds one batch-independent (T*H*W, D) table
+            # The Cosmos25 RoPE builds one batch-independent (T * H * W, D) table
             # (see cosmos2_5.py), so the DiT consumes a single shared fps
             # value per forward pass, never one value per sample.
             if not torch.equal(fps_tensor, fps_tensor[:1].expand_as(fps_tensor)):

fastvideo/tests/attention/test_vsa64_fused_layout.py

@@ -0,0 +1,164 @@
+"""VSA64 fused layout routing, fallback, and forward parity.
+
+The fused route is gated on SM100, but the Triton layout kernel itself is
+architecture-agnostic. Routing and layout tests therefore run on any CUDA GPU
+(the unit CI runs on L40S) with the capability check pinned to SM100; only the
+end-to-end ``video_sparse_attn`` parity test needs a real SM100 device.
+"""
+
+import pytest
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.attention.backends import video_sparse_attn as vsa_mod
+from fastvideo.attention.backends.video_sparse_attn import VideoSparseAttentionImpl, VideoSparseAttentionMetadataBuilder
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device required")
+
+requires_sm100 = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0),
+    reason="SM100 CUDA device required",
+)
+
+
+@pytest.fixture
+def as_sm100(monkeypatch):
+    """Pin the capability gate to SM100 so fallbacks are caused by the condition under test."""
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
+
+
+def _fused_layout(enabled: bool):
+    return envs.FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT.override(not enabled)
+
+
+@pytest.fixture
+def tile_kernel():
+    """Require the in-tree layout kernel instead of skipping without it.
+
+    CI builds fastvideo-kernel from source, so a missing module means the build
+    is broken or stale (e.g. the pinned PyPI wheel won), and a skip would let
+    the lane pass without exercising the fused route.
+    """
+    pytest.importorskip("triton")
+    try:
+        import fastvideo_kernel.triton_kernels.vsa_tile_layout  # noqa: F401
+    except ImportError as exc:
+        pytest.fail("fastvideo_kernel.triton_kernels.vsa_tile_layout is missing; build fastvideo-kernel from "
+                    f"source (cd fastvideo-kernel && ./build.sh): {exc}")
+
+
+def _build(dit_shape: tuple[int, int, int]):
+    impl = object.__new__(VideoSparseAttentionImpl)
+    builder = VideoSparseAttentionMetadataBuilder()
+    kwargs = dict(
+        current_timestep=0,
+        raw_latent_shape=dit_shape,
+        patch_size=(1, 1, 1),
+        VSA_sparsity=0.8,
+        device=torch.device("cuda"),
+        cache_tile_buf=True,
+    )
+    return impl, builder, kwargs
+
+
+@pytest.mark.parametrize("dit_shape", [(5, 5, 6), (16, 28, 52)])
+def test_vsa64_fused_layout_matches_legacy_tiles(as_sm100, tile_kernel, dit_shape):
+    impl, builder, kwargs = _build(dit_shape)
+    fused_metadata = builder.build(**kwargs)
+    legacy_metadata = builder.build(**kwargs)
+    sequence = fused_metadata.total_seq_length
+    # Non-contiguous input, as produced by the SP-padding trim in the attention layer.
+    qkvg = torch.randn((4, sequence + 7, 2, 128), device="cuda", dtype=torch.bfloat16)[:, :sequence]
+
+    with _fused_layout(True), torch.no_grad():
+        fused = impl.preprocess_qkv(qkvg, fused_metadata)
+    assert fused_metadata.fused_layout_active
+
+    with _fused_layout(False), torch.no_grad():
+        legacy = impl.preprocess_qkv(qkvg, legacy_metadata).transpose(1, 2).contiguous()
+    assert not legacy_metadata.fused_layout_active
+
+    assert fused.shape == legacy.shape
+    assert torch.equal(fused, legacy)
+
+
+@pytest.mark.parametrize("dit_shape", [(5, 5, 6), (16, 28, 52)])
+def test_vsa64_fused_layout_can_be_disabled(as_sm100, dit_shape):
+    impl, builder, kwargs = _build(dit_shape)
+    metadata = builder.build(**kwargs)
+    sequence = metadata.total_seq_length
+    qkvg = torch.randn((4, sequence, 2, 128), device="cuda", dtype=torch.bfloat16)
+    with _fused_layout(False), torch.no_grad():
+        disabled = impl.preprocess_qkv(qkvg, metadata)
+    assert not metadata.fused_layout_active
+    assert disabled.ndim == 4 and disabled.shape[1] == metadata.variable_block_sizes.numel() * 64
+
+
+def test_vsa64_fused_layout_training_falls_back(as_sm100):
+    impl, builder, kwargs = _build((5, 5, 6))
+    metadata = builder.build(**kwargs)
+    qkvg = torch.randn((4, 150, 2, 128), device="cuda", dtype=torch.bfloat16)
+    with _fused_layout(True):
+        training = impl.preprocess_qkv(qkvg.requires_grad_(), metadata)
+    assert not metadata.fused_layout_active
+    assert training.shape == (4, 512, 2, 128)
+    assert training.requires_grad
+
+
+def test_vsa64_mismatched_metadata_uses_original_path(monkeypatch, as_sm100):
+    impl, builder, kwargs = _build((5, 5, 6))
+    metadata = builder.build(**kwargs)
+    qkvg = torch.empty((4, metadata.total_seq_length - 1, 2, 128), device="cuda", dtype=torch.bfloat16)
+    fallback = object()
+    calls = []
+
+    def tile(x, meta):
+        calls.append((x, meta))
+        return fallback
+
+    monkeypatch.setattr(impl, "tile", tile)
+    with _fused_layout(True), torch.no_grad():
+        result = impl.preprocess_qkv(qkvg, metadata)
+    assert result is fallback
+    assert len(calls) == 1
+    assert calls[0][0] is qkvg and calls[0][1] is metadata
+    assert not metadata.fused_layout_active
+
+
+def test_vsa64_fused_layout_falls_back_when_kernel_missing(monkeypatch, as_sm100):
+    impl, builder, kwargs = _build((5, 5, 6))
+    metadata = builder.build(**kwargs)
+    qkvg = torch.randn((4, metadata.total_seq_length, 2, 128), device="cuda", dtype=torch.bfloat16)
+    monkeypatch.setattr(vsa_mod, "_get_tile_to_bhsd", lambda: None)
+    with _fused_layout(True), torch.no_grad():
+        result = impl.preprocess_qkv(qkvg, metadata)
+    assert not metadata.fused_layout_active
+    assert result.shape[1] == metadata.variable_block_sizes.numel() * 64
+
+
+@requires_sm100
+@pytest.mark.parametrize("dit_shape", [(5, 5, 6), (16, 28, 52)])
+def test_vsa64_fused_forward_matches_legacy(dit_shape):
+    try:
+        from fastvideo_kernel import video_sparse_attn  # noqa: F401
+    except ImportError:
+        pytest.skip("video_sparse_attn is not installed")
+
+    impl, builder, kwargs = _build(dit_shape)
+    fused_metadata = builder.build(**kwargs)
+    legacy_metadata = builder.build(**kwargs)
+    sequence = fused_metadata.total_seq_length
+    qkvg = torch.randn((4, sequence, 2, 128), device="cuda", dtype=torch.bfloat16)
+
+    with _fused_layout(True), torch.no_grad():
+        fused_tiled = impl.preprocess_qkv(qkvg, fused_metadata)
+        assert fused_metadata.fused_layout_active
+        q, k, v, gate = fused_tiled.chunk(4)
+        fused_out = impl.forward(q, k, v, gate, fused_metadata)
+
+    with _fused_layout(False), torch.no_grad():
+        legacy_tiled = impl.preprocess_qkv(qkvg, legacy_metadata)
+        lq, lk, lv, lg = legacy_tiled.chunk(4)
+        legacy_out = impl.forward(lq, lk, lv, lg, legacy_metadata)
+
+    assert torch.equal(fused_out, legacy_out)