From be548a78eab190c4d03e48d5a284c78e06b69990 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Tue, 26 May 2026 12:58:29 -0700 Subject: [PATCH] [feat] VSA-256 fastpath on Blackwell via FA4 CuTe block-sparse attention (#1354) --- fastvideo-kernel/README.md | 23 ++ .../python/fastvideo_kernel/__init__.py | 2 + .../fastvideo_kernel/block_sparse_attn.py | 37 ++- .../fastvideo_kernel/block_sparse_attn_256.py | 170 +++++++++++ .../block_sparse_attn_cute_fwd.py | 267 ++++++++++++++++++ .../python/fastvideo_kernel/ops.py | 166 +++++++++-- fastvideo-kernel/tests/test_vsa256_forward.py | 131 +++++++++ .../tests/test_vsa256_forward_cross.py | 139 +++++++++ .../tests/test_vsa256_forward_vbs.py | 129 +++++++++ fastvideo-kernel/tests/test_vsa256_triton.py | 150 ++++++++++ .../attention/backends/video_sparse_attn.py | 41 ++- 11 files changed, 1211 insertions(+), 44 deletions(-) create mode 100644 fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py create mode 100644 fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py create mode 100644 fastvideo-kernel/tests/test_vsa256_forward.py create mode 100644 fastvideo-kernel/tests/test_vsa256_forward_cross.py create mode 100644 fastvideo-kernel/tests/test_vsa256_forward_vbs.py create mode 100644 fastvideo-kernel/tests/test_vsa256_triton.py diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 82d3fe81d..6df12db89 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -28,6 +28,29 @@ cd fastvideo-kernel ./build.sh --rocm ``` +### Optional: FA4 CuTe block-sparse backend (VSA-256 fastpath) + +The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes to the +FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is +an **optional** dependency: it is imported lazily, and `video_sparse_attn` +transparently falls back to the Triton backend when it is absent (so the package is +fully usable without it). + +The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`, +`flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by +[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to +commit `c19cd20e`: the wrapper targets that revision's `_flash_attn_fwd` signature +(`m_block_size` / `n_block_size`); later upstream revisions reshaped it into a +`tile_mn` tuple and are not drop-in compatible. + +```bash +pip install "nvidia-cutlass-dsl>=4.5.0" torchvision +pip install "git+https://github.com/Dao-AILab/flash-attention.git@c19cd20e#subdirectory=flash_attn/cute" +``` + +The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against +`tests/test_vsa256_forward*.py`. + ## Usage ### Sliding Tile Attention (STA) & Video Sparse Attention (VSA) diff --git a/fastvideo-kernel/python/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py index 0e6b8da6c..e3a7bbfce 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/__init__.py +++ b/fastvideo-kernel/python/fastvideo_kernel/__init__.py @@ -3,6 +3,7 @@ from .version import __version__ from fastvideo_kernel.ops import ( sliding_tile_attention, video_sparse_attn, + video_sparse_attn_bshd, ) from fastvideo_kernel.block_sparse_attn import ( @@ -27,6 +28,7 @@ from fastvideo_kernel.turbodiffusion_ops import ( __all__ = [ "sliding_tile_attention", "video_sparse_attn", + "video_sparse_attn_bshd", "block_sparse_attn", "block_sparse_attn_from_indices", "moba_attn_varlen", diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index b5a1217b3..caad8b672 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -32,7 +32,24 @@ def _is_sm90() -> bool: def _force_triton() -> bool: - return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1" + """True iff Triton is explicitly forced via env var. + + Honors both the new ``FASTVIDEO_VSA_TRITON`` and the legacy + ``FASTVIDEO_KERNEL_VSA_FORCE_TRITON`` for backward compatibility. + """ + return ( + os.environ.get("FASTVIDEO_VSA_TRITON", "0") == "1" + or os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1" + ) + + +def _force_tk() -> bool: + """True iff the sm_90 (ThunderKittens) backend is explicitly requested. + + Only effective on sm_90 with the compiled extension present; otherwise + falls back to the default selection. + """ + return os.environ.get("FASTVIDEO_VSA_TK", "0") == "1" # --------------------------------------------------------------------------- @@ -368,12 +385,24 @@ def block_sparse_attn_from_indices( variable_block_sizes = _as_int32_contig(variable_block_sizes, "variable_block_sizes") block_sparse_fwd, block_sparse_bwd = _get_sm90_ops() - use_sm90 = ( - (not _force_triton()) - and _is_sm90() + sm90_available = ( + _is_sm90() and block_sparse_fwd is not None and block_sparse_bwd is not None ) + + # Backend resolution: + # - FASTVIDEO_VSA_TRITON forces Triton everywhere. + # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually + # available (else falls through to the default below). + # - Otherwise: TK on sm_90 if available, else Triton. + if _force_triton(): + use_sm90 = False + elif _force_tk(): + use_sm90 = sm90_available + else: + use_sm90 = sm90_available + if use_sm90: return block_sparse_attn_sm90(q, k, v, q2k_idx, q2k_num, variable_block_sizes) # Triton path: supports q_seq_len != kv_seq_len as long as both are padded diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py new file mode 100644 index 000000000..dc95a5fe6 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py @@ -0,0 +1,170 @@ +"""VSA-256 block-sparse attention wrapper. + +The default 256-block path is Triton: it expands the logical 256-block map +to the existing 64-block Triton kernel via a dense 4x4 expansion per logical +edge ("route A"), and requires no optional dependencies. + +The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is +*opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to +:mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates +on 128-token KV blocks (this wrapper expands the logical 256-block map / +sizes into that physical 128-block representation). The CuTe kernel +(``flash_attn.cute`` with block-sparsity) is an optional dependency, +imported lazily only when this fastpath is selected. + +``FASTVIDEO_VSA_TRITON=1`` (or the legacy +``FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1``) forces Triton explicitly. +""" + +from __future__ import annotations + +import os +from typing import Tuple + +import torch + +from .block_sparse_attn import block_sparse_attn_triton, _force_triton + +# NOTE: ``block_sparse_attn_cute_fwd`` is imported lazily inside the CuTe +# branches below. Importing it at module load would pull in the optional +# FA4 CuTe build (``flash_attn.cute``) and make it a hard dependency of the +# default Triton path. + + +_KV_BLOCK_PHYS = 128 # FA4 CuTe BSA forward uses 128-token KV blocks. +_KV_BLOCK_TRITON = 64 # Existing Triton path uses 64-token KV blocks. + + +def _resolve_backend() -> str: + """Pick the backend for the 256-block VSA path. + + Default is Triton (no optional deps). The FA4 CuTe fastpath is opt-in + via ``FASTVIDEO_VSA_CUTEDSL=1`` and requires the optional FA4 CuTe + build. ``FASTVIDEO_VSA_TRITON=1`` / the legacy force-triton flag force + Triton explicitly and take precedence over the CuTe opt-in. + """ + if _force_triton(): + return "triton" + if os.environ.get("FASTVIDEO_VSA_CUTEDSL", "0") == "1": + return "cutedsl" + return "triton" + + +def _expand_mask_and_sizes_256_to_128( + logical_mask_256: torch.Tensor, + logical_kv_sizes_256: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Expand a [B, H, Qb256, KVb256] map to [B, H, Qb256, KVb128]. + + Each logical 256-token KV block splits into two physical 128-token + children. Each child inherits the logical mask edge; its valid-token + count is the logical count clamped into the child's window. + """ + expanded_mask = logical_mask_256.repeat_interleave(2, dim=3) + + sizes_i32 = logical_kv_sizes_256.to(torch.int32) + child0 = torch.clamp(sizes_i32, min=0, max=_KV_BLOCK_PHYS) + child1 = torch.clamp(sizes_i32 - _KV_BLOCK_PHYS, min=0, max=_KV_BLOCK_PHYS) + expanded_sizes = torch.empty( + (sizes_i32.numel() * 2,), + dtype=torch.int32, + device=sizes_i32.device, + ) + expanded_sizes[0::2] = child0 + expanded_sizes[1::2] = child1 + return expanded_mask, expanded_sizes + + +def _expand_mask_and_sizes_256_to_64( + logical_mask_256: torch.Tensor, + logical_kv_sizes_256: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Expand a [B, H, Qb256, KVb256] map to [B, H, Qb64, KVb64] (route A). + + Each logical 256-token tile splits 4-ways along both Q and KV. Sizes + are computed by chopping the logical count into 64-token strides. + """ + expanded_mask = logical_mask_256.repeat_interleave(4, dim=2).repeat_interleave(4, dim=3) + sizes_i32 = logical_kv_sizes_256.to(torch.int32) + offsets = torch.tensor( + [0, _KV_BLOCK_TRITON, 2 * _KV_BLOCK_TRITON, 3 * _KV_BLOCK_TRITON], + dtype=torch.int32, + device=sizes_i32.device, + ) + expanded_sizes = torch.clamp( + sizes_i32[:, None] - offsets[None, :], + min=0, + max=_KV_BLOCK_TRITON, + ).reshape(-1) + return expanded_mask, expanded_sizes + + +def _triton_via_route_a( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_mask_256: torch.Tensor, + logical_kv_sizes_256: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + from .triton_kernels.index import map_to_index as triton_map_to_index + + mask_64, sizes_64 = _expand_mask_and_sizes_256_to_64( + logical_mask_256, logical_kv_sizes_256 + ) + q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool)) + return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64) + + +def block_sparse_attn_256( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_block_map_256: torch.Tensor, + logical_variable_block_sizes_256: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """VSA-256 sparse-branch entrypoint for [B, H, S, D] inputs.""" + if logical_block_map_256.dim() == 3: + logical_block_map_256 = logical_block_map_256.unsqueeze(0) + + if _resolve_backend() == "triton": + return _triton_via_route_a( + q, k, v, logical_block_map_256, logical_variable_block_sizes_256 + ) + + mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128( + logical_block_map_256, logical_variable_block_sizes_256 + ) + from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd + return block_sparse_attn_cute_fwd(q, k, v, mask_128, sizes_128) + + +def block_sparse_attn_256_bshd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_block_map_256: torch.Tensor, + logical_variable_block_sizes_256: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """VSA-256 sparse-branch entrypoint for [B, S, H, D] inputs. + + Default CuTe path consumes BSHD directly; Triton fallback transposes + to BHSD as the legacy path expects. + """ + if logical_block_map_256.dim() == 3: + logical_block_map_256 = logical_block_map_256.unsqueeze(0) + + if _resolve_backend() == "triton": + out_bhsd, aux = _triton_via_route_a( + q.transpose(1, 2).contiguous(), + k.transpose(1, 2).contiguous(), + v.transpose(1, 2).contiguous(), + logical_block_map_256, + logical_variable_block_sizes_256, + ) + return out_bhsd.transpose(1, 2).contiguous(), aux + + mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128( + logical_block_map_256, logical_variable_block_sizes_256 + ) + from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd + return block_sparse_attn_cute_fwd_bshd(q, k, v, mask_128, sizes_128) diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py new file mode 100644 index 000000000..6d7fa5926 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py @@ -0,0 +1,267 @@ +"""CuTe-DSL block-sparse attention forward kernel. + +Thin wrapper around `flash_attn.cute.interface._flash_attn_fwd` that adapts +VSA's `(block_map, variable_block_sizes)` inputs into FA4's +`BlockSparseTensorsTorch` representation and the per-KV-block validity mask. + +Both [B, H, S, D] (BHSD) and [B, S, H, D] (BSHD) entrypoints are provided. +The BSHD variant is preferred from VSA-256 callers to avoid layout +round-trips on the hot path. + +The FA4 CuTe block-sparse kernel (``flash_attn.cute`` with +``block_sparsity``) is an *optional* dependency: it is imported lazily and +only exercised when the VSA-256 CuTe fastpath is explicitly selected +(``FASTVIDEO_VSA_CUTEDSL=1``). The default VSA-256 path is Triton and does +not require it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``. +""" + +from __future__ import annotations + +import functools +from typing import Tuple + +import torch + +_FA4_IMPORT_HINT = ( + "VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that " + "provides `flash_attn.cute` with block-sparsity support (plus " + "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional " + "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe " + "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath." +) + + +def _load_fa4_cute(): + """Lazily import the optional FA4 CuTe block-sparse symbols. + + Raising a clear, actionable error here keeps the optional FA4 CuTe + build from being a hard import-time dependency of this module (and of + the default Triton VSA-256 path). + """ + try: + from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch + from flash_attn.cute.interface import _flash_attn_fwd + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError(_FA4_IMPORT_HINT) from exc + return BlockSparseTensorsTorch, _flash_attn_fwd + + +# FA4 BSA fwd uses (m_block_size, n_block_size); m_block_size=128 is the +# Q-side tile, kv_block_size comes from the caller's VSA logical KV block. +_M_BLOCK_SIZE_DEFAULT = 128 + + +def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + if block_map.dim() == 3: + block_map = block_map.unsqueeze(0) + if block_map.dim() != 4: + raise ValueError( + f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), " + f"got shape={tuple(block_map.shape)}" + ) + if block_map.dtype != torch.bool: + block_map = block_map.to(torch.bool) + if not block_map.is_cuda: + raise RuntimeError("block_map must be a CUDA tensor.") + from fastvideo_kernel.triton_kernels.index import ( + map_to_index as triton_map_to_index, + ) + + return triton_map_to_index(block_map) + + +def _choose_q_sparse_block_size(q_len: int, m_block_size: int = _M_BLOCK_SIZE_DEFAULT) -> int: + # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > m_block_size. + major, _ = torch.cuda.get_device_capability() + if major >= 10 and q_len > m_block_size: + return 2 * m_block_size + return m_block_size + + +def _aggregate_q_block_map( + block_map: torch.Tensor, + q_sparse_block_size: int, + q_block_size: int, +) -> torch.Tensor: + factor = q_sparse_block_size // q_block_size + if factor <= 0 or q_sparse_block_size % q_block_size != 0: + raise ValueError( + f"q_sparse_block_size must be a positive multiple of " + f"q_block_size ({q_block_size}), got {q_sparse_block_size}" + ) + bsz, nhead, q_blocks, kv_blocks = block_map.shape + q_blocks_sparse = (q_blocks + factor - 1) // factor + pad_q = q_blocks_sparse * factor - q_blocks + if pad_q > 0: + pad = torch.zeros( + bsz, + nhead, + pad_q, + kv_blocks, + dtype=torch.bool, + device=block_map.device, + ) + block_map = torch.cat([block_map, pad], dim=2) + block_map = block_map.view(bsz, nhead, q_blocks_sparse, factor, kv_blocks) + return block_map.any(dim=3) + + +@functools.lru_cache(maxsize=4) +def _build_vbs_mask_mod(kv_block_size: int): + """Build a CuTe mask_mod that trims per-KV-block valid tokens. + + aux_tensors[0] must be an int32 tensor of shape [kv_blocks] giving the + valid token count in [0, kv_block_size] for each KV block. + """ + import cutlass + import cutlass.cute as cute + from flash_attn.cute import utils + from flash_attn.cute.block_sparsity import fast_sampling + + kv_block_size_const = int(kv_block_size) + + @fast_sampling + @cute.jit + def _vbs_mask_mod( + batch: cute.TensorSSA, + head: cute.TensorSSA, + m_idx: cute.TensorSSA, + n_idx: cute.TensorSSA, + seqlen_info, + aux_tensors, + ) -> cute.TensorSSA: + del batch, head, m_idx, seqlen_info + block_size_ssa = utils.scalar_to_ssa(kv_block_size_const, cutlass.Int32) + zero_ssa = utils.scalar_to_ssa(0, cutlass.Int32) + kv_blk = n_idx // block_size_ssa + kv_off = n_idx % block_size_ssa + kv_sizes = aux_tensors[0] + valid = utils.scalar_to_ssa(kv_sizes[kv_blk[0]], cutlass.Int32) + return (valid > zero_ssa) & (kv_off < valid) + + return _vbs_mask_mod + + +def _cute_forward( + q_bshd: torch.Tensor, + k_bshd: torch.Tensor, + v_bshd: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + *, + q_block_size: int, + kv_block_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Internal: FA4 CuTe BSA fwd with BSHD inputs.""" + BlockSparseTensorsTorch, _flash_attn_fwd = _load_fa4_cute() + q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1]) + q_sparse_block_size = max( + q_block_size, + ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size, + ) + sparse_map = _aggregate_q_block_map( + block_map, + q_sparse_block_size=q_sparse_block_size, + q_block_size=q_block_size, + ) + kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1) + kv_partial = ( + (variable_block_sizes > 0) & (variable_block_sizes < kv_block_size) + ).view(1, 1, 1, -1) + full_map = sparse_map & kv_full + mask_map = sparse_map & kv_partial + + full_block_idx, full_block_cnt = _map_to_index(full_map) + mask_block_idx, mask_block_cnt = _map_to_index(mask_map) + + sparse_tensors = BlockSparseTensorsTorch( + full_block_cnt=full_block_cnt.to(torch.int32).contiguous(), + full_block_idx=full_block_idx.to(torch.int32).contiguous(), + mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(), + mask_block_idx=mask_block_idx.to(torch.int32).contiguous(), + block_size=(q_sparse_block_size, kv_block_size), + ) + + out, lse = _flash_attn_fwd( + q_bshd, + k_bshd, + v_bshd, + m_block_size=_M_BLOCK_SIZE_DEFAULT, + n_block_size=kv_block_size, + mask_mod=_build_vbs_mask_mod(kv_block_size), + block_sparse_tensors=sparse_tensors, + aux_tensors=[variable_block_sizes], + causal=False, + return_lse=True, + ) + return out, lse + + +def block_sparse_attn_cute_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """CuTe forward-only block-sparse attention with [B, H, S, D] inputs.""" + if block_map.dim() == 3: + block_map = block_map.unsqueeze(0) + q_block_size = q.shape[2] // block_map.shape[2] + kv_block_size = k.shape[2] // block_map.shape[3] + + q_bshd = q.transpose(1, 2).contiguous() + k_bshd = k.transpose(1, 2).contiguous() + v_bshd = v.transpose(1, 2).contiguous() + out_bshd, lse_bshd = _cute_forward( + q_bshd, + k_bshd, + v_bshd, + block_map, + variable_block_sizes, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + ) + out = out_bshd.transpose(1, 2).contiguous() + if lse_bshd is None: + lse = torch.empty( + (q.shape[0], q.shape[1], q.shape[2]), + dtype=torch.float32, + device=q.device, + ) + else: + lse = lse_bshd.transpose(1, 2).contiguous() + return out, lse + + +def block_sparse_attn_cute_fwd_bshd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """CuTe forward-only block-sparse attention with [B, S, H, D] inputs.""" + if block_map.dim() == 3: + block_map = block_map.unsqueeze(0) + q_block_size = q.shape[1] // block_map.shape[2] + kv_block_size = k.shape[1] // block_map.shape[3] + + out, lse_bshd = _cute_forward( + q, + k, + v, + block_map, + variable_block_sizes, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + ) + if lse_bshd is None: + lse = torch.empty( + (q.shape[0], q.shape[2], q.shape[1]), + dtype=torch.float32, + device=q.device, + ) + else: + lse = lse_bshd.transpose(1, 2).contiguous() + return out, lse diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index 398a5c130..0f5458856 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -1,6 +1,10 @@ import math import torch from .block_sparse_attn import block_sparse_attn, block_sparse_attn_from_indices +from .block_sparse_attn_256 import ( + block_sparse_attn_256, + block_sparse_attn_256_bshd, +) from .triton_kernels.st_attn_triton import sliding_tile_attention_triton # Try to load the C++ extension @@ -71,19 +75,30 @@ def video_sparse_attn( block_size: int | tuple = 64, compress_attn_weight: torch.Tensor = None, ) -> torch.Tensor: + """VSA entrypoint for [B, H, S, D] tensors. + + Dispatches the sparse branch by ``block_elements = prod(block_size)``: + - 64 -> existing TK/Triton path (see ``block_sparse_attn_from_indices``). + - 256 -> CuTe FA4 block-sparse attention (see ``block_sparse_attn_256``). + + Backend overrides: + - ``FASTVIDEO_VSA_TRITON=1`` forces Triton in either path. + - ``FASTVIDEO_VSA_TK=1`` prefers the sm_90 TK kernel in the 64-block path. + - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 256-block path. + """ if isinstance(block_size, int): block_size = (block_size, block_size, block_size) - block_elements = block_size[0] * block_size[1] * block_size[2] + batch, heads, q_seq_len, dim = q.shape kv_seq_len = k.shape[2] + if k.shape[0] != batch or v.shape[0] != batch or k.shape[1] != heads or v.shape[1] != heads: + raise ValueError("Expected q/k/v to have the same batch and head dimensions.") if v.shape[2] != kv_seq_len: raise ValueError( f"Expected k and v to have the same sequence length, got " f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}" ) - if k.shape[0] != batch or v.shape[0] != batch or k.shape[1] != heads or v.shape[1] != heads: - raise ValueError("Expected q/k/v to have the same batch and head dimensions.") if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0: raise ValueError( @@ -92,52 +107,147 @@ def video_sparse_attn( ) q_num_blocks = q_seq_len // block_elements kv_num_blocks = kv_seq_len // block_elements - if variable_block_sizes.numel() != kv_num_blocks: raise ValueError( f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " f"got {variable_block_sizes.numel()}" ) - if q_variable_block_sizes.numel() != q_num_blocks: raise ValueError( f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " f"got {q_variable_block_sizes.numel()}" ) - # Compression branch + # Compression branch (token-level average per block + dense block-level attn). q_c = q.view(batch, heads, q_num_blocks, block_elements, dim) k_c = k.view(batch, heads, kv_num_blocks, block_elements, dim) v_c = v.view(batch, heads, kv_num_blocks, block_elements, dim) + q_c = (q_c.float().sum(dim=3) / q_variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype) - q_c = (q_c.float().sum(dim=3) / q_variable_block_sizes.view(1, 1, -1, 1)).to( - q.dtype) - k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to( - k.dtype) - v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to( - v.dtype) - - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) - out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) - out_c = out_c.repeat(1, 1, 1, block_elements, - 1).view(batch, heads, q_seq_len, dim) + out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim) - # Sparse branch: feed top-k indices directly, skipping the bool-mask round-trip. + # Sparse branch. topk_idx = torch.topk(scores, topk, dim=-1).indices - q2k_idx = topk_idx.to(torch.int32).contiguous() - q2k_num = torch.full( - (batch, heads, q_num_blocks), - topk, - dtype=torch.int32, - device=q.device, - ) - out_s = block_sparse_attn_from_indices( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - )[0] + + if block_elements == 256: + # CuTe path consumes a bool mask (full/partial split inside the wrapper). + mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) + out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0] + else: + # Index-native path for 64-block (TK/Triton). + q2k_idx = topk_idx.to(torch.int32).contiguous() + q2k_num = torch.full( + (batch, heads, q_num_blocks), + topk, + dtype=torch.int32, + device=q.device, + ) + out_s = block_sparse_attn_from_indices( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + )[0] if compress_attn_weight is not None: return out_c * compress_attn_weight + out_s return out_c + out_s + + +def video_sparse_attn_bshd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + variable_block_sizes: torch.Tensor, + q_variable_block_sizes: torch.Tensor, + topk: int, + block_size: int | tuple = 64, + compress_attn_weight: torch.Tensor = None, +) -> torch.Tensor: + """VSA entrypoint for [B, S, H, D] tensors. + + Avoids the BHSD<->BSHD round-trip that ``video_sparse_attn`` performs on + the CuTe 256-block path; the 64-block path still expects BHSD and is not + supported here. + """ + if isinstance(block_size, int): + block_size = (block_size, block_size, block_size) + block_elements = block_size[0] * block_size[1] * block_size[2] + if block_elements != 256: + raise ValueError( + "video_sparse_attn_bshd is only defined for block_elements=256 " + f"(got {block_elements}); use video_sparse_attn for the 64-block path." + ) + + batch, q_seq_len, heads, dim = q.shape + kv_seq_len = k.shape[1] + if k.shape[0] != batch or v.shape[0] != batch or k.shape[2] != heads or v.shape[2] != heads: + raise ValueError("Expected q/k/v to have the same batch and head dimensions.") + if v.shape[1] != kv_seq_len: + raise ValueError( + f"Expected k and v to have the same sequence length, got " + f"k.shape[1]={kv_seq_len}, v.shape[1]={v.shape[1]}" + ) + if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0: + raise ValueError( + f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, " + f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}" + ) + q_num_blocks = q_seq_len // block_elements + kv_num_blocks = kv_seq_len // block_elements + if variable_block_sizes.numel() != kv_num_blocks: + raise ValueError( + f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " + f"got {variable_block_sizes.numel()}" + ) + if q_variable_block_sizes.numel() != q_num_blocks: + raise ValueError( + f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " + f"got {q_variable_block_sizes.numel()}" + ) + + # Compression branch (BSHD-native: mean over the 256-token axis). + token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32) + q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1, 1)).view( + 1, q_num_blocks, block_elements, 1, 1 + ) + kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1, 1)).view( + 1, kv_num_blocks, block_elements, 1, 1 + ) + + q_c = q.view(batch, q_num_blocks, block_elements, heads, dim) + k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim) + v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim) + q_c = ((q_c.float() * q_token_valid).sum(dim=2) + / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype) + k_c = ((k_c.float() * kv_token_valid).sum(dim=2) + / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype) + v_c = ((v_c.float() * kv_token_valid).sum(dim=2) + / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype) + q_ch = q_c.permute(0, 2, 1, 3).contiguous() + k_ch = k_c.permute(0, 2, 1, 3).contiguous() + v_ch = v_c.permute(0, 2, 1, 3).contiguous() + + scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim ** 0.5) + attn = torch.softmax(scores, dim=-1) + out_c_ch = torch.matmul(attn, v_ch) + out_c_blk = out_c_ch.permute(0, 2, 1, 3).contiguous() + + # Sparse branch (CuTe BSHD). + topk_idx = torch.topk(scores, topk, dim=-1).indices + mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) + out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes) + + out = out_s + out_view = out.view(batch, q_num_blocks, block_elements, heads, dim) + if compress_attn_weight is not None: + gate_view = compress_attn_weight.view( + batch, q_num_blocks, block_elements, heads, dim + ) + out_view.add_(out_c_blk.unsqueeze(2) * gate_view) + else: + out_view.add_(out_c_blk.unsqueeze(2)) + return out diff --git a/fastvideo-kernel/tests/test_vsa256_forward.py b/fastvideo-kernel/tests/test_vsa256_forward.py new file mode 100644 index 000000000..666eb4303 --- /dev/null +++ b/fastvideo-kernel/tests/test_vsa256_forward.py @@ -0,0 +1,131 @@ +"""VSA-256 CuTe forward correctness vs. a dense torch reference.""" + +from __future__ import annotations + +import math +from typing import Tuple + +import pytest +import torch + +from fastvideo_kernel import video_sparse_attn + +from .utils import create_full_mask_from_block_mask + + +@pytest.fixture(autouse=True) +def _require_cute_backend(monkeypatch): + """Exercise the FA4 CuTe fastpath; skip if the optional build is absent. + + The VSA-256 default backend is Triton (see ``block_sparse_attn_256``); + these tests target the CuTe path, so they explicitly opt in and skip + cleanly when the optional ``flash_attn.cute`` build is not installed. + """ + pytest.importorskip( + "flash_attn.cute.block_sparsity", + reason="optional FA4 CuTe build (flash_attn.cute) not installed", + ) + monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1") + + +def _dense_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + full_mask: torch.Tensor, +) -> torch.Tensor: + # q,k,v: [B,H,S,D]; full_mask: [H,Sq,Skv]. + qf = q.float() + kf = k.float() + vf = v.float() + attn = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(q.shape[-1]) + attn = attn.masked_fill(~full_mask.unsqueeze(0), float("-inf")) + prob = torch.softmax(attn, dim=-1) + return torch.matmul(prob, vf).to(q.dtype) + + +def _run_case( + heads: int, + head_dim: int, + q_blocks_256: int, + kv_blocks_256: int, + topk_logical: int, +) -> Tuple[float, float]: + assert torch.cuda.is_available() + # Default dispatch routes (4,8,8) -> CuTe; no env var needed. + + device = torch.device("cuda") + dtype = torch.bfloat16 + batch = 1 + q_block = 256 + kv_block_logical = 256 + + sq = q_blocks_256 * q_block + skv = kv_blocks_256 * kv_block_logical + q = torch.randn(batch, heads, sq, head_dim, device=device, dtype=dtype) + k = torch.randn(batch, heads, skv, head_dim, device=device, dtype=dtype) + v = torch.randn(batch, heads, skv, head_dim, device=device, dtype=dtype) + + q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) + kv_var = torch.full((kv_blocks_256,), kv_block_logical, dtype=torch.int32, device=device) + + # Reproduce the compression branch's per-block average so we can pre-compute + # the top-k mask that the kernel will see and feed it to the dense reference. + q_c = q.view(batch, heads, q_blocks_256, q_block, head_dim) + k_c = k.view(batch, heads, kv_blocks_256, kv_block_logical, head_dim) + v_c = v.view(batch, heads, kv_blocks_256, kv_block_logical, head_dim) + q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype) + + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (head_dim ** 0.5) + attn = torch.softmax(scores, dim=-1) + out_c = torch.matmul(attn, v_c) + out_c = ( + out_c.view(batch, heads, q_blocks_256, 1, head_dim) + .repeat(1, 1, 1, q_block, 1) + .view(batch, heads, sq, head_dim) + ) + topk_idx = torch.topk(scores, topk_logical, dim=-1).indices + mask_256 = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)[0] + full_mask = create_full_mask_from_block_mask(mask_256, q_var, kv_var, device=device) + out_ref = out_c + _dense_reference(q, k, v, full_mask) + + out = video_sparse_attn( + q, k, v, + kv_var, + q_var, + topk_logical, + block_size=(4, 8, 8), + compress_attn_weight=None, + ) + + assert torch.isfinite(out).all().item(), "NaN/Inf in kernel output" + diff = (out_ref - out).abs() + avg_abs = diff.mean().item() + max_rel = (diff.max() / (out_ref.abs().mean() + 1e-6)).item() + return avg_abs, max_rel + + +@pytest.mark.cuda +def test_vsa256_forward_qk_equal() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + avg_abs, max_rel = _run_case( + heads=8, head_dim=128, q_blocks_256=8, kv_blocks_256=8, topk_logical=2, + ) + print(f"[vsa256 qk_equal] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") + assert avg_abs < 5e-2 + assert max_rel < 2.0 + + +@pytest.mark.cuda +def test_vsa256_forward_qk_diff() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + avg_abs, max_rel = _run_case( + heads=8, head_dim=128, q_blocks_256=8, kv_blocks_256=12, topk_logical=2, + ) + print(f"[vsa256 qk_diff] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") + assert avg_abs < 5e-2 + assert max_rel < 2.0 diff --git a/fastvideo-kernel/tests/test_vsa256_forward_cross.py b/fastvideo-kernel/tests/test_vsa256_forward_cross.py new file mode 100644 index 000000000..6527ab37b --- /dev/null +++ b/fastvideo-kernel/tests/test_vsa256_forward_cross.py @@ -0,0 +1,139 @@ +"""Three-way parity for the VSA-256 forward: torch reference vs CuTe vs Triton.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from fastvideo_kernel import video_sparse_attn + + +@pytest.fixture(autouse=True) +def _require_cute_backend(monkeypatch): + """3-way parity needs the CuTe leg; skip if the optional build is absent. + + VSA-256 defaults to Triton, so opt the CuTe leg in explicitly. The test + later forces Triton via ``FASTVIDEO_VSA_TRITON=1`` (which takes + precedence in ``_resolve_backend``) for the Triton leg. + """ + pytest.importorskip( + "flash_attn.cute.block_sparsity", + reason="optional FA4 CuTe build (flash_attn.cute) not installed", + ) + monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1") + + +def _torch_vsa256_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_var: torch.Tensor, + kv_var: torch.Tensor, + topk_logical: int, +) -> torch.Tensor: + bsz, heads, _sq, dim = q.shape + q_blocks = q_var.numel() + kv_blocks = kv_var.numel() + q_block = int(q_var[0].item()) + kv_block = int(kv_var[0].item()) + + q_c = q.view(bsz, heads, q_blocks, q_block, dim) + k_c = k.view(bsz, heads, kv_blocks, kv_block, dim) + v_c = v.view(bsz, heads, kv_blocks, kv_block, dim) + q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype) + + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim) + attn = torch.softmax(scores, dim=-1) + out_c = torch.matmul(attn, v_c) + out_c = ( + out_c.view(bsz, heads, q_blocks, 1, dim) + .repeat(1, 1, 1, q_block, 1) + .view_as(q) + ) + + topk_idx = torch.topk(scores, topk_logical, dim=-1).indices + block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) + token_mask = ( + block_mask + .repeat_interleave(q_block, dim=2) + .repeat_interleave(kv_block, dim=3) + .to(torch.bool) + ) + + qf, kf, vf = q.float(), k.float(), v.float() + logits = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(dim) + logits = logits.masked_fill(~token_mask, float("-inf")) + prob = torch.softmax(logits, dim=-1) + out_s = torch.matmul(prob, vf).to(q.dtype) + return out_c + out_s + + +def _metrics(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float]: + diff = (a - b).abs() + avg_abs = diff.mean().item() + max_rel = (diff.max() / (a.abs().mean() + 1e-6)).item() + return avg_abs, max_rel + + +@pytest.mark.cuda +def test_vsa256_forward_cross_torch_cute_triton(monkeypatch) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + + bsz, heads, dim = 1, 8, 128 + q_blocks_256, kv_blocks_256 = 8, 12 + topk_logical = 2 + q_block = 256 + kv_block = 256 + + sq = q_blocks_256 * q_block + skv = kv_blocks_256 * kv_block + q = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype) + k = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) + v = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) + q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) + kv_var = torch.full((kv_blocks_256,), kv_block, dtype=torch.int32, device=device) + + out_torch = _torch_vsa256_reference(q, k, v, q_var, kv_var, topk_logical) + + # CuTe (default). + out_cute = video_sparse_attn( + q, k, v, + kv_var, q_var, topk_logical, + block_size=(4, 8, 8), + compress_attn_weight=None, + ) + + # Triton via route-A 256->64 expansion. + monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") + out_triton = video_sparse_attn( + q, k, v, + kv_var, q_var, topk_logical, + block_size=(4, 8, 8), + compress_attn_weight=None, + ) + + for t in (out_torch, out_cute, out_triton): + assert torch.isfinite(t).all().item() + + torch_vs_cute = _metrics(out_torch, out_cute) + torch_vs_triton = _metrics(out_torch, out_triton) + cute_vs_triton = _metrics(out_cute, out_triton) + print( + "[cross-forward] " + f"torch_vs_cute(avg_abs={torch_vs_cute[0]:.6e}, max_rel={torch_vs_cute[1]:.6e}), " + f"torch_vs_triton(avg_abs={torch_vs_triton[0]:.6e}, max_rel={torch_vs_triton[1]:.6e}), " + f"cute_vs_triton(avg_abs={cute_vs_triton[0]:.6e}, max_rel={cute_vs_triton[1]:.6e})" + ) + + assert torch_vs_cute[0] < 1e-3 and torch_vs_cute[1] < 0.2 + assert torch_vs_triton[0] < 1e-3 and torch_vs_triton[1] < 0.2 + assert cute_vs_triton[0] < 1e-3 and cute_vs_triton[1] < 0.2 diff --git a/fastvideo-kernel/tests/test_vsa256_forward_vbs.py b/fastvideo-kernel/tests/test_vsa256_forward_vbs.py new file mode 100644 index 000000000..e993f82cc --- /dev/null +++ b/fastvideo-kernel/tests/test_vsa256_forward_vbs.py @@ -0,0 +1,129 @@ +"""VSA-256 CuTe correctness with variable KV block sizes (<256).""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from fastvideo_kernel import video_sparse_attn + + +@pytest.fixture(autouse=True) +def _require_cute_backend(monkeypatch): + """Exercise the FA4 CuTe fastpath; skip if the optional build is absent. + + VSA-256 defaults to Triton; this file targets the CuTe path, so it + opts in explicitly and skips when ``flash_attn.cute`` is not installed. + """ + pytest.importorskip( + "flash_attn.cute.block_sparsity", + reason="optional FA4 CuTe build (flash_attn.cute) not installed", + ) + monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1") + + +def _torch_vsa256_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_var: torch.Tensor, + kv_var: torch.Tensor, + topk_logical: int, +) -> torch.Tensor: + bsz, heads, _sq, dim = q.shape + q_blocks = q_var.numel() + kv_blocks = kv_var.numel() + q_block = q.shape[2] // q_blocks + kv_block = k.shape[2] // kv_blocks + + q_c = q.view(bsz, heads, q_blocks, q_block, dim) + k_c = k.view(bsz, heads, kv_blocks, kv_block, dim) + v_c = v.view(bsz, heads, kv_blocks, kv_block, dim) + q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype) + + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim) + attn = torch.softmax(scores, dim=-1) + out_c = torch.matmul(attn, v_c) + out_c = ( + out_c.view(bsz, heads, q_blocks, 1, dim) + .repeat(1, 1, 1, q_block, 1) + .view_as(q) + ) + + with torch.no_grad(): + topk_idx = torch.topk(scores.detach(), topk_logical, dim=-1).indices + block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) + block_token_idx = torch.arange(kv_block, device=kv_var.device, dtype=torch.int32) + kv_token_valid_by_block = ( + block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1) + ).to(torch.bool) + kv_token_valid = kv_token_valid_by_block.reshape(1, 1, 1, kv_blocks * kv_block) + token_mask = ( + block_mask + .repeat_interleave(q_block, dim=2) + .repeat_interleave(kv_block, dim=3) + ) + token_mask = token_mask & kv_token_valid + + qf, kf, vf = q.float(), k.float(), v.float() + logits = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(dim) + logits = logits.masked_fill(~token_mask, float("-inf")) + prob = torch.softmax(logits, dim=-1) + out_s = torch.matmul(prob, vf).to(q.dtype) + return out_c + out_s + + +@pytest.mark.cuda +def test_vsa256_cute_variable_block_size_vs_torch_ref() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + + bsz, heads, dim = 1, 8, 128 + q_blocks_256, kv_blocks_256 = 8, 12 + q_block = 256 + kv_block = 256 + topk_logical = 2 + sq = q_blocks_256 * q_block + skv = kv_blocks_256 * kv_block + + q = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype) + k = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) + v = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) + + q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) + kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256,), dtype=torch.int32, device=device) + token_idx = torch.arange(kv_block, device=device, dtype=torch.int32) + kv_valid = token_idx.view(1, -1) < kv_var.view(-1, 1) + kv_valid = kv_valid.view(1, 1, kv_blocks_256, kv_block, 1) + kv_valid = kv_valid.expand(bsz, heads, kv_blocks_256, kv_block, dim).reshape(bsz, heads, skv, dim) + k = k * kv_valid.to(k.dtype) + v = v * kv_valid.to(v.dtype) + + out = video_sparse_attn( + q, k, v, + kv_var, + q_var, + topk_logical, + block_size=(4, 8, 8), + compress_attn_weight=None, + ) + out_ref = _torch_vsa256_reference(q, k, v, q_var, kv_var, topk_logical) + + diff = (out - out_ref).abs() + avg_abs = diff.mean().item() + max_rel = (diff.max() / (out_ref.abs().mean() + 1e-6)).item() + print( + f"[vsa256-cute-vbs] kv_var[min={int(kv_var.min().item())}, " + f"max={int(kv_var.max().item())}], " + f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}" + ) + + assert avg_abs < 1e-3 and max_rel < 0.2 diff --git a/fastvideo-kernel/tests/test_vsa256_triton.py b/fastvideo-kernel/tests/test_vsa256_triton.py new file mode 100644 index 000000000..400953785 --- /dev/null +++ b/fastvideo-kernel/tests/test_vsa256_triton.py @@ -0,0 +1,150 @@ +"""VSA-256 Triton route-A (256->64 expansion) forward + backward parity. + +Forces ``FASTVIDEO_VSA_TRITON=1`` so the 256 path falls through to the +existing 64-block Triton kernel. +""" + +from __future__ import annotations + +import math +from typing import Tuple + +import pytest +import torch + +from fastvideo_kernel import video_sparse_attn + + +def _torch_vsa256_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_var: torch.Tensor, + kv_var: torch.Tensor, + topk_logical: int, +) -> torch.Tensor: + bsz, heads, _sq, dim = q.shape + q_blocks = q_var.numel() + kv_blocks = kv_var.numel() + q_block = q.shape[2] // q_blocks + kv_block = k.shape[2] // kv_blocks + + q_c = q.view(bsz, heads, q_blocks, q_block, dim) + k_c = k.view(bsz, heads, kv_blocks, kv_block, dim) + v_c = v.view(bsz, heads, kv_blocks, kv_block, dim) + q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype) + + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim) + attn = torch.softmax(scores, dim=-1) + out_c = torch.matmul(attn, v_c) + out_c = ( + out_c.view(bsz, heads, q_blocks, 1, dim) + .repeat(1, 1, 1, q_block, 1) + .view_as(q) + ) + + with torch.no_grad(): + topk_idx = torch.topk(scores.detach(), topk_logical, dim=-1).indices + block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) + block_token_idx = torch.arange(kv_block, device=kv_var.device, dtype=torch.int32) + kv_token_valid_by_block = ( + block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1) + ).to(torch.bool) + kv_token_valid = kv_token_valid_by_block.reshape(1, 1, 1, kv_blocks * kv_block) + token_mask = ( + block_mask + .repeat_interleave(q_block, dim=2) + .repeat_interleave(kv_block, dim=3) + ) + token_mask = token_mask & kv_token_valid + + qf, kf, vf = q.float(), k.float(), v.float() + logits = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(dim) + logits = logits.masked_fill(~token_mask, float("-inf")) + prob = torch.softmax(logits, dim=-1) + out_s = torch.matmul(prob, vf).to(q.dtype) + return out_c + out_s + + +def _metrics(a: torch.Tensor, b: torch.Tensor) -> Tuple[float, float]: + diff = (a - b).abs() + avg_abs = diff.mean().item() + max_rel = (diff.max() / (a.abs().mean() + 1e-6)).item() + return avg_abs, max_rel + + +@pytest.mark.cuda +def test_vsa256_triton_forward_backward_vs_torch_ref(monkeypatch) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + + bsz, heads, dim = 1, 8, 128 + q_blocks_256, kv_blocks_256 = 8, 12 + q_block = 256 + kv_block = 256 + topk_logical = 2 + sq = q_blocks_256 * q_block + skv = kv_blocks_256 * kv_block + + q_base = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype) + k_base = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) + v_base = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) + grad_out = torch.randn_like(q_base) + + q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) + kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256,), dtype=torch.int32, device=device) + token_idx = torch.arange(kv_block, device=device, dtype=torch.int32) + kv_valid = token_idx.view(1, -1) < kv_var.view(-1, 1) + kv_valid = kv_valid.view(1, 1, kv_blocks_256, kv_block, 1) + kv_valid = kv_valid.expand(bsz, heads, kv_blocks_256, kv_block, dim).reshape(bsz, heads, skv, dim) + k_base = k_base * kv_valid.to(k_base.dtype) + v_base = v_base * kv_valid.to(v_base.dtype) + + q = q_base.detach().clone().requires_grad_(True) + k = k_base.detach().clone().requires_grad_(True) + v = v_base.detach().clone().requires_grad_(True) + out = video_sparse_attn( + q, k, v, + kv_var, + q_var, + topk_logical, + block_size=(4, 8, 8), + compress_attn_weight=None, + ) + (out * grad_out).sum().backward() + dq, dk, dv = q.grad.detach(), k.grad.detach(), v.grad.detach() + + q_ref = q_base.detach().clone().requires_grad_(True) + k_ref = k_base.detach().clone().requires_grad_(True) + v_ref = v_base.detach().clone().requires_grad_(True) + out_ref = _torch_vsa256_reference(q_ref, k_ref, v_ref, q_var, kv_var, topk_logical) + (out_ref * grad_out).sum().backward() + dq_ref, dk_ref, dv_ref = q_ref.grad.detach(), k_ref.grad.detach(), v_ref.grad.detach() + + for t in (out, out_ref, dq, dk, dv): + assert torch.isfinite(t).all().item() + + m_out = _metrics(out_ref, out) + m_dq = _metrics(dq_ref, dq) + m_dk = _metrics(dk_ref, dk) + m_dv = _metrics(dv_ref, dv) + print( + "[vsa256-triton] " + f"kv_var[min={int(kv_var.min().item())}, max={int(kv_var.max().item())}], " + f"out(avg_abs={m_out[0]:.6e}, max_rel={m_out[1]:.6e}), " + f"dq(avg_abs={m_dq[0]:.6e}, max_rel={m_dq[1]:.6e}), " + f"dk(avg_abs={m_dk[0]:.6e}, max_rel={m_dk[1]:.6e}), " + f"dv(avg_abs={m_dv[0]:.6e}, max_rel={m_dv[1]:.6e})" + ) + + assert m_out[0] < 1e-3 and m_out[1] < 0.2 + assert m_dq[0] < 2e-2 and m_dq[1] < 0.5 + assert m_dk[0] < 2e-2 and m_dk[1] < 0.5 + assert m_dv[0] < 2e-2 and m_dv[1] < 0.5 diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index 2d2fc0384..df44c9d94 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -9,6 +9,10 @@ try: from fastvideo_kernel import video_sparse_attn except ImportError: video_sparse_attn = None +try: + from fastvideo_kernel import video_sparse_attn_bshd +except ImportError: + video_sparse_attn_bshd = None from typing import Any @@ -18,6 +22,9 @@ from fastvideo.distributed import get_sp_group from fastvideo.logger import init_logger logger = init_logger(__name__) +# VSA tile shape. The tile volume picks the kernel path automatically in +# forward(): (4,4,4)=64 -> existing TK/Triton path (default, unchanged); +# (4,8,8)=256 -> FA4 CuTe block-sparse attention fastpath (Blackwell). VSA_TILE_SIZE = (4, 4, 4) @@ -275,24 +282,34 @@ class VideoSparseAttentionImpl(AttentionImpl): gate_compress: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: - 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).contiguous() - VSA_sparsity = attn_metadata.VSA_sparsity + block_elements = math.prod(VSA_TILE_SIZE) + cur_topk = math.ceil((1 - VSA_sparsity) * (attn_metadata.total_seq_length / block_elements)) - cur_topk = math.ceil((1 - VSA_sparsity) * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE))) - - if video_sparse_attn is None: - raise NotImplementedError("video_sparse_attn is not installed") - hidden_states = video_sparse_attn(query, + # 256-element tiles auto-route to the FA4 CuTe BSHD fastpath, which + # consumes [B, S, H, D] directly -- skip the transpose round-trip. + if block_elements == 256 and video_sparse_attn_bshd is not None: + return video_sparse_attn_bshd(query, key, value, attn_metadata.variable_block_sizes, attn_metadata.variable_block_sizes, cur_topk, block_size=VSA_TILE_SIZE, - compress_attn_weight=gate_compress).transpose(1, 2) + compress_attn_weight=gate_compress) - return hidden_states + if video_sparse_attn is None: + raise NotImplementedError("video_sparse_attn is not installed") + # Default 64-element-tile path (unchanged): BHSD round-trip. + 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).contiguous() + return video_sparse_attn(query, + key, + value, + attn_metadata.variable_block_sizes, + attn_metadata.variable_block_sizes, + cur_topk, + block_size=VSA_TILE_SIZE, + compress_attn_weight=gate_compress).transpose(1, 2)