[feat] VSA-256 fastpath on Blackwell via FA4 CuTe block-sparse attention (#1354)

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",

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

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)

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

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