210 lines
7.8 KiB
Python
210 lines
7.8 KiB
Python
"""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`. By default this wrapper
|
|
expands logical KV256 blocks to the historical physical KV128 tiles. Set
|
|
``FASTVIDEO_VSA_FA4_BLOCK_SHAPE=128x64`` for the fine-grained FA4 schedule.
|
|
The production ``64x64`` selection pairs adjacent Q64 children into physical
|
|
Q128 tiles: their KV lists are identical because this wrapper derives them
|
|
from one original Q256 metadata row. Direct Q64 callers with arbitrary masks
|
|
continue to use the native Q64/KV64 path. 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.
|
|
|
|
_LOGICAL_BLOCK_SIZE = 256
|
|
_FA4_BLOCK_SHAPES = {
|
|
# The public name describes the original VSA mask. FA4 internally splits
|
|
# its logical KV256 edge into two physical KV128 tiles.
|
|
"256x256": (256, 128),
|
|
"128x64": (128, 64),
|
|
# This pairing is safe only at this original-Q256 wrapper boundary: all
|
|
# four Q64 children inherit exactly the same KV list from their parent.
|
|
# Keep arbitrary/distinct Q64 maps on block_sparse_attn_cute_fwd's native
|
|
# Q64/KV64 path, where they are not coalesced.
|
|
"64x64": (128, 64),
|
|
}
|
|
_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 _resolve_fa4_block_shape() -> Tuple[int, int]:
|
|
requested = os.environ.get("FASTVIDEO_VSA_FA4_BLOCK_SHAPE", "256x256").lower()
|
|
try:
|
|
return _FA4_BLOCK_SHAPES[requested]
|
|
except KeyError as exc:
|
|
choices = ", ".join(_FA4_BLOCK_SHAPES)
|
|
raise ValueError(
|
|
f"unsupported FASTVIDEO_VSA_FA4_BLOCK_SHAPE={requested!r}; choose one of {choices}"
|
|
) from exc
|
|
|
|
|
|
def _expand_mask_and_sizes_256_for_fa4(
|
|
logical_mask_256: torch.Tensor,
|
|
logical_kv_sizes_256: torch.Tensor,
|
|
q_block_size: int,
|
|
kv_block_size: int,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Expand a logical Q256/KV256 map to one supported physical FA4 shape.
|
|
|
|
Every child inherits its parent edge. KV valid-token counts are clamped
|
|
into each child's window, preserving partial and empty tail blocks.
|
|
"""
|
|
if _LOGICAL_BLOCK_SIZE % q_block_size or _LOGICAL_BLOCK_SIZE % kv_block_size:
|
|
raise ValueError(
|
|
f"FA4 QxKV blocks must divide {_LOGICAL_BLOCK_SIZE}, got "
|
|
f"{q_block_size}x{kv_block_size}"
|
|
)
|
|
q_factor = _LOGICAL_BLOCK_SIZE // q_block_size
|
|
kv_factor = _LOGICAL_BLOCK_SIZE // kv_block_size
|
|
expanded_mask = logical_mask_256.repeat_interleave(q_factor, dim=2)
|
|
expanded_mask = expanded_mask.repeat_interleave(kv_factor, dim=3)
|
|
|
|
sizes_i32 = logical_kv_sizes_256.to(torch.int32)
|
|
offsets = torch.arange(
|
|
kv_factor,
|
|
dtype=torch.int32,
|
|
device=sizes_i32.device,
|
|
) * kv_block_size
|
|
expanded_sizes = torch.clamp(
|
|
sizes_i32[:, None] - offsets[None, :],
|
|
min=0,
|
|
max=kv_block_size,
|
|
).reshape(-1)
|
|
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)
|
|
|
|
q_block_size, kv_block_size = _resolve_fa4_block_shape()
|
|
physical_mask, physical_sizes = _expand_mask_and_sizes_256_for_fa4(
|
|
logical_block_map_256,
|
|
logical_variable_block_sizes_256,
|
|
q_block_size,
|
|
kv_block_size,
|
|
)
|
|
from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
|
|
return block_sparse_attn_cute_fwd(q, k, v, physical_mask, physical_sizes)
|
|
|
|
|
|
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
|
|
|
|
q_block_size, kv_block_size = _resolve_fa4_block_shape()
|
|
physical_mask, physical_sizes = _expand_mask_and_sizes_256_for_fa4(
|
|
logical_block_map_256,
|
|
logical_variable_block_sizes_256,
|
|
q_block_size,
|
|
kv_block_size,
|
|
)
|
|
from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
|
|
return block_sparse_attn_cute_fwd_bshd(q, k, v, physical_mask, physical_sizes)
|