From 9e83ba630c2ad337090365c81e18412e054f1110 Mon Sep 17 00:00:00 2001 From: sumyyyyy <46865688+freemty@users.noreply.github.com> Date: Wed, 1 Jul 2026 03:19:25 +0800 Subject: [PATCH] [kernel] Extract VSA utility functions into fastvideo_kernel (#1408) Co-authored-by: SolitaryThinker --- .../python/fastvideo_kernel/__init__.py | 15 + .../python/fastvideo_kernel/vsa_utils.py | 160 ++++++++++ fastvideo-kernel/tests/test_vsa_utils.py | 276 ++++++++++++++++++ 3 files changed, 451 insertions(+) create mode 100644 fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py create mode 100644 fastvideo-kernel/tests/test_vsa_utils.py diff --git a/fastvideo-kernel/python/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py index cf12e6ab6..d8ff068c3 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/__init__.py +++ b/fastvideo-kernel/python/fastvideo_kernel/__init__.py @@ -29,6 +29,15 @@ from fastvideo_kernel.block_sparse_attn_varlen import ( block_sparse_attn_varlen, ) +from fastvideo_kernel.vsa_utils import ( + VSA_TILE_SIZE, + get_tile_partition_indices, + get_reverse_tile_partition_indices, + construct_variable_block_sizes, + get_non_pad_index, + build_vsa_metadata, +) + __all__ = [ "sliding_tile_attention", "video_sparse_attn", @@ -44,5 +53,11 @@ __all__ = [ "FastLayerNorm", "int8_linear", "int8_quant", + "VSA_TILE_SIZE", + "get_tile_partition_indices", + "get_reverse_tile_partition_indices", + "construct_variable_block_sizes", + "get_non_pad_index", + "build_vsa_metadata", "__version__", ] diff --git a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py new file mode 100644 index 000000000..d0dc1b840 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py @@ -0,0 +1,160 @@ +"""VSA metadata utilities — standalone, no fastvideo framework dependency. + +Provides the tile-partition index helpers and variable-block-size +computations that are needed to call `video_sparse_attn` or +`block_sparse_attn_from_indices` without depending on the full +fastvideo package. +""" + +from __future__ import annotations + +import functools +import math + +import torch + +VSA_TILE_SIZE = (4, 4, 4) +_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 256) + + +def _canonicalize_device(device: torch.device | str) -> torch.device: + """Resolve an indexless CUDA device before it is used as a cache key.""" + device = torch.device(device) + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +@functools.lru_cache(maxsize=10) +def get_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + """Map raster-order token indices to tile-contiguous order. + + Groups spatially adjacent tokens into (ts_t x ts_h x ts_w) tiles + so that each tile's tokens are contiguous in the output. + """ + T, H, W = dit_seq_shape + ts, hs, ws = tile_size + indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W) + ls = [] + for t in range(math.ceil(T / ts)): + for h in range(math.ceil(H / hs)): + for w in range(math.ceil(W / ws)): + ls.append(indices[ + t * ts:min(t * ts + ts, T), + h * hs:min(h * hs + hs, H), + w * ws:min(w * ws + ws, W), + ].flatten()) + return torch.cat(ls, dim=0) + + +@functools.lru_cache(maxsize=10) +def get_reverse_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + """Inverse of get_tile_partition_indices: tile order back to raster.""" + return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device)) + + +@functools.lru_cache(maxsize=10) +def construct_variable_block_sizes( + dit_seq_shape: tuple[int, int, int], + num_tiles: tuple[int, int, int], + device: torch.device, + tile_size: tuple[int, int, int] = VSA_TILE_SIZE, +) -> torch.LongTensor: + """Compute the number of valid tokens in each tile. + + Tiles at the boundary of each dimension may contain fewer tokens + when the video shape is not evenly divisible by tile_size. + """ + t, h, w = dit_seq_shape + ts_t, ts_h, ts_w = tile_size + n_t, n_h, n_w = num_tiles + + def _sizes(dim_len: int, tile: int, n: int) -> torch.LongTensor: + sizes = torch.full((n,), tile, dtype=torch.int, device=device) + remainder = dim_len - (n - 1) * tile + sizes[-1] = remainder if remainder > 0 else tile + return sizes + + t_sizes = _sizes(t, ts_t, n_t) + h_sizes = _sizes(h, ts_h, n_h) + w_sizes = _sizes(w, ts_w, n_w) + + return ( + t_sizes[:, None, None] + * h_sizes[None, :, None] + * w_sizes[None, None, :] + ).reshape(-1) + + +def get_non_pad_index( + variable_block_sizes: torch.LongTensor, + max_block_size: int, +) -> torch.LongTensor: + """Find positions of real tokens within a block-padded layout. + + Each block occupies max_block_size slots. This returns the flat + indices of the valid (non-padding) positions. + """ + n_win = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts_pad = torch.arange(n_win, device=device) * max_block_size + index_pad = starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :] + index_mask = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None] + return index_pad[index_mask] + + +def build_vsa_metadata( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int] = VSA_TILE_SIZE, + device: torch.device | str = "cuda", +) -> dict: + """Build all VSA metadata from a video latent shape in one call. + + Args: + dit_seq_shape: (T, H, W) — temporal frames, spatial height, width. + tile_size: (ts_t, ts_h, ts_w) — tokens per tile in each dimension. + The resulting tile volume must be supported by the VSA kernels. + device: Target device for index tensors. + + Returns: + Dict with keys: tile_partition_indices, reverse_tile_partition_indices, + variable_block_sizes, non_pad_index, num_tiles, max_block_size. + """ + device = _canonicalize_device(device) + + T, H, W = dit_seq_shape + ts_t, ts_h, ts_w = tile_size + max_block_size = math.prod(tile_size) + if max_block_size not in _SUPPORTED_VSA_BLOCK_VOLUMES: + raise ValueError( + f"Unsupported VSA tile volume {max_block_size} for tile_size={tile_size}; " + f"supported volumes are {_SUPPORTED_VSA_BLOCK_VOLUMES}." + ) + + num_tiles = ( + math.ceil(T / ts_t), + math.ceil(H / ts_h), + math.ceil(W / ts_w), + ) + + tile_indices = get_tile_partition_indices(dit_seq_shape, tile_size, device) + reverse_tile_indices = get_reverse_tile_partition_indices(dit_seq_shape, tile_size, device) + vbs = construct_variable_block_sizes(dit_seq_shape, num_tiles, device, tile_size) + npi = get_non_pad_index(vbs, max_block_size) + + return { + "tile_partition_indices": tile_indices, + "reverse_tile_partition_indices": reverse_tile_indices, + "variable_block_sizes": vbs, + "non_pad_index": npi, + "num_tiles": num_tiles, + "max_block_size": max_block_size, + } diff --git a/fastvideo-kernel/tests/test_vsa_utils.py b/fastvideo-kernel/tests/test_vsa_utils.py new file mode 100644 index 000000000..35e6b21cf --- /dev/null +++ b/fastvideo-kernel/tests/test_vsa_utils.py @@ -0,0 +1,276 @@ +"""Tests for vsa_utils — standalone VSA metadata utilities. + +These tests use CPU index computation except for an optional multi-GPU +regression covering cache isolation between CUDA devices. +""" + +import math +import pytest +import torch + +from fastvideo_kernel.vsa_utils import ( + VSA_TILE_SIZE, + _canonicalize_device, + get_tile_partition_indices, + get_reverse_tile_partition_indices, + construct_variable_block_sizes, + get_non_pad_index, + build_vsa_metadata, +) + + +class TestDeviceCanonicalization: + + def test_unindexed_cuda_resolves_current_device(self, monkeypatch): + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + assert _canonicalize_device("cuda") == torch.device("cuda:3") + assert _canonicalize_device(torch.device("cuda")) == torch.device("cuda:3") + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") + def test_metadata_cache_isolated_between_cuda_devices(self): + shape = (7, 8, 8) + tensor_keys = ( + "tile_partition_indices", + "reverse_tile_partition_indices", + "variable_block_sizes", + "non_pad_index", + ) + + with torch.cuda.device(0): + metadata_0 = build_vsa_metadata(shape, device="cuda") + with torch.cuda.device(1): + metadata_1 = build_vsa_metadata(shape, device="cuda") + + for key in tensor_keys: + assert metadata_0[key].device == torch.device("cuda:0") + assert metadata_1[key].device == torch.device("cuda:1") + + +class TestGetTilePartitionIndices: + + @pytest.mark.parametrize("dit_seq_shape,tile_size", [ + ((8, 16, 16), (4, 4, 4)), + ((4, 8, 8), (4, 4, 4)), + ((9, 10, 7), (4, 4, 4)), + ]) + def test_is_valid_permutation(self, dit_seq_shape, tile_size): + """Output must be a permutation of 0..N-1.""" + device = torch.device("cpu") + idx = get_tile_partition_indices(dit_seq_shape, tile_size, device) + n = math.prod(dit_seq_shape) + assert idx.shape == (n,) + assert idx.dtype == torch.long + assert set(idx.tolist()) == set(range(n)) + + def test_exact_values_small(self): + """Manually verify a small case: (2,2,2) with tile (2,2,2) = 1 tile.""" + device = torch.device("cpu") + idx = get_tile_partition_indices((2, 2, 2), (2, 2, 2), device) + assert idx.tolist() == list(range(8)) + + def test_non_divisible_shape(self): + """When shape doesn't divide evenly by tile_size, all tokens still covered.""" + device = torch.device("cpu") + shape = (5, 7, 3) + idx = get_tile_partition_indices(shape, (4, 4, 4), device) + assert idx.shape == (5 * 7 * 3,) + assert set(idx.tolist()) == set(range(5 * 7 * 3)) + + +class TestGetReverseTilePartitionIndices: + + @pytest.mark.parametrize("dit_seq_shape", [ + (8, 16, 16), + (9, 10, 7), + ]) + def test_inverse_of_forward(self, dit_seq_shape): + """reverse[forward[i]] == i for all i.""" + device = torch.device("cpu") + tile_size = (4, 4, 4) + fwd = get_tile_partition_indices(dit_seq_shape, tile_size, device) + rev = get_reverse_tile_partition_indices(dit_seq_shape, tile_size, device) + n = math.prod(dit_seq_shape) + identity = torch.arange(n, device=device) + assert torch.equal(rev[fwd], identity) + assert torch.equal(fwd[rev], identity) + + +class TestConstructVariableBlockSizes: + + def test_sum_equals_total_tokens(self): + """Sum of block sizes must equal T*H*W.""" + device = torch.device("cpu") + shape = (8, 16, 16) + tile_size = (4, 4, 4) + num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size)) + vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + assert vbs.sum().item() == math.prod(shape) + + def test_max_block_size(self): + """No block can exceed tile volume.""" + device = torch.device("cpu") + shape = (9, 10, 7) + tile_size = (4, 4, 4) + num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size)) + vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + assert vbs.max().item() <= math.prod(tile_size) + + def test_num_blocks(self): + """Number of blocks = product of num_tiles.""" + device = torch.device("cpu") + shape = (8, 16, 16) + tile_size = (4, 4, 4) + num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size)) + vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + assert vbs.shape[0] == math.prod(num_tiles) + + def test_exact_divisible(self): + """When perfectly divisible, all blocks have the same size.""" + device = torch.device("cpu") + shape = (8, 8, 8) + tile_size = (4, 4, 4) + num_tiles = (2, 2, 2) + vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + assert (vbs == 64).all() + + def test_non_divisible_last_tile_smaller(self): + """When not divisible, at least one block is smaller than max.""" + device = torch.device("cpu") + shape = (9, 8, 8) + tile_size = (4, 4, 4) + num_tiles = (3, 2, 2) + vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + assert vbs.min().item() < math.prod(tile_size) + + def test_custom_tile_size(self): + """tile_size parameter overrides default VSA_TILE_SIZE.""" + device = torch.device("cpu") + shape = (6, 8, 16) + tile_size = (2, 4, 8) + num_tiles = (3, 2, 2) + vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + assert vbs.sum().item() == math.prod(shape) + assert (vbs == 64).all() + + +class TestGetNonPadIndex: + + def test_length_equals_sum_block_sizes(self): + """Output length must equal sum of variable_block_sizes.""" + vbs = torch.tensor([32, 48, 64], dtype=torch.long) + idx = get_non_pad_index(vbs, 64) + assert idx.shape[0] == 32 + 48 + 64 + + def test_indices_in_valid_range(self): + """All indices must be in [0, num_blocks * max_block_size).""" + vbs = torch.tensor([32, 48], dtype=torch.long) + idx = get_non_pad_index(vbs, 64) + assert idx.min().item() >= 0 + assert idx.max().item() < 2 * 64 + + def test_block_boundary_alignment(self): + """First token of block i starts at i * max_block_size.""" + vbs = torch.tensor([20, 40], dtype=torch.long) + idx = get_non_pad_index(vbs, 64) + assert idx[0].item() == 0 + assert idx[20].item() == 64 + + def test_full_blocks(self): + """When all blocks are full, output is just 0..N-1.""" + vbs = torch.tensor([64, 64], dtype=torch.long) + idx = get_non_pad_index(vbs, 64) + assert torch.equal(idx, torch.arange(128)) + + +class TestBuildVsaMetadata: + + def test_all_keys_present(self): + """build_vsa_metadata returns all expected keys.""" + meta = build_vsa_metadata((8, 16, 16), device="cpu") + expected_keys = { + "tile_partition_indices", "reverse_tile_partition_indices", + "variable_block_sizes", "non_pad_index", + "num_tiles", "max_block_size", + } + assert set(meta.keys()) == expected_keys + + def test_types(self): + """Return types are correct.""" + meta = build_vsa_metadata((8, 16, 16), device="cpu") + assert isinstance(meta["tile_partition_indices"], torch.Tensor) + assert isinstance(meta["reverse_tile_partition_indices"], torch.Tensor) + assert isinstance(meta["variable_block_sizes"], torch.Tensor) + assert isinstance(meta["non_pad_index"], torch.Tensor) + assert isinstance(meta["num_tiles"], tuple) + assert isinstance(meta["max_block_size"], int) + + def test_num_tiles_correct(self): + meta = build_vsa_metadata((9, 10, 7), tile_size=(4, 4, 4), device="cpu") + assert meta["num_tiles"] == (3, 3, 2) + assert meta["max_block_size"] == 64 + + @pytest.mark.parametrize("tile_size,expected_num_tiles,expected_block_size", [ + ((2, 4, 8), (3, 2, 2), 64), + ((4, 8, 8), (2, 1, 2), 256), + ]) + def test_supported_custom_tile_size(self, tile_size, expected_num_tiles, expected_block_size): + meta = build_vsa_metadata((6, 8, 16), tile_size=tile_size, device="cpu") + assert meta["num_tiles"] == expected_num_tiles + assert meta["max_block_size"] == expected_block_size + + def test_unsupported_tile_volume(self): + with pytest.raises(ValueError, match="Unsupported VSA tile volume 27"): + build_vsa_metadata((6, 6, 6), tile_size=(3, 3, 3), device="cpu") + + def test_consistency(self): + """All components are internally consistent.""" + shape = (8, 16, 16) + meta = build_vsa_metadata(shape, device="cpu") + n = math.prod(shape) + assert meta["tile_partition_indices"].shape == (n,) + assert meta["reverse_tile_partition_indices"].shape == (n,) + assert meta["variable_block_sizes"].sum().item() == n + assert meta["non_pad_index"].shape[0] == n + + +class TestConsistencyWithFramework: + """Verify vsa_utils matches the framework-level functions exactly. + + Only runs if fastvideo is importable (skip otherwise). + """ + + @pytest.fixture(autouse=True) + def _skip_if_no_framework(self): + try: + from fastvideo.attention.backends.video_sparse_attn import ( + get_tile_partition_indices as fw_get_tile, + ) + except ImportError: + pytest.skip("fastvideo framework not installed") + + @pytest.mark.parametrize("shape", [(8, 16, 16), (9, 10, 7)]) + def test_tile_indices_match(self, shape): + from fastvideo.attention.backends.video_sparse_attn import ( + get_tile_partition_indices as fw_get_tile, + ) + device = torch.device("cpu") + tile_size = (4, 4, 4) + ours = get_tile_partition_indices(shape, tile_size, device) + theirs = fw_get_tile(shape, tile_size, device) + assert torch.equal(ours, theirs) + + @pytest.mark.parametrize("shape", [(8, 16, 16), (9, 10, 7)]) + def test_variable_block_sizes_match(self, shape): + from fastvideo.attention.backends.video_sparse_attn import ( + construct_variable_block_sizes as fw_construct_vbs, + ) + device = torch.device("cpu") + tile_size = (4, 4, 4) + num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size)) + ours = construct_variable_block_sizes(shape, num_tiles, device, tile_size) + theirs = fw_construct_vbs(shape, num_tiles, device) + assert torch.equal(ours, theirs) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"])