267 lines
8.9 KiB
Python
267 lines
8.9 KiB
Python
"""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
|
|
|
|
|
|
# Q-side tile size; 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),
|
|
)
|
|
|
|
# _flash_attn_fwd returns (out, lse, p, row_max); keep the first two.
|
|
out, lse = _flash_attn_fwd(
|
|
q_bshd,
|
|
k_bshd,
|
|
v_bshd,
|
|
tile_mn=(_M_BLOCK_SIZE_DEFAULT, 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,
|
|
)[:2]
|
|
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
|