diff --git a/docs/attention/vsa/index.md b/docs/attention/vsa/index.md index 3db0c812b..c66385597 100644 --- a/docs/attention/vsa/index.md +++ b/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 diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 7d449eb87..1f12a9f15 100644 --- a/docs/contributing/env_vars.md +++ b/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. | diff --git a/fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py b/fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py new file mode 100644 index 000000000..cd3e605ab --- /dev/null +++ b/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() diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/vsa_tile_layout.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/vsa_tile_layout.py new file mode 100644 index 000000000..18df37965 --- /dev/null +++ b/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 diff --git a/fastvideo-kernel/tests/test_vsa_tile_layout.py b/fastvideo-kernel/tests/test_vsa_tile_layout.py new file mode 100644 index 000000000..13988232c --- /dev/null +++ b/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 diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index 686c99355..cf77e13d7 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/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, diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 25f53d48b..d1cfb62ac 100644 --- a/fastvideo/envs.py +++ b/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 ================== diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 9ee5e45b5..2e4bfb36a 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/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)): diff --git a/fastvideo/tests/attention/test_vsa64_fused_layout.py b/fastvideo/tests/attention/test_vsa64_fused_layout.py new file mode 100644 index 000000000..e7f6ed56e --- /dev/null +++ b/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)