72795d0af8
Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
270 lines
12 KiB
Python
270 lines
12 KiB
Python
import math
|
|
import torch
|
|
from .block_sparse_attn import block_sparse_attn
|
|
from .block_sparse_attn_256 import (
|
|
block_sparse_attn_128,
|
|
block_sparse_attn_128_bshd,
|
|
block_sparse_attn_256,
|
|
block_sparse_attn_256_bshd,
|
|
)
|
|
from .triton_kernels.st_attn_triton import sliding_tile_attention_triton
|
|
from .triton_kernels.fused_compress_topk import fused_block_mean, fused_topk_mask
|
|
|
|
# Try to load the C++ extension
|
|
try:
|
|
from fastvideo_kernel._C import fastvideo_kernel_ops
|
|
sta_fwd = getattr(fastvideo_kernel_ops, "sta_fwd", None)
|
|
except ImportError:
|
|
sta_fwd = None
|
|
|
|
|
|
def sliding_tile_attention(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
window_size: list,
|
|
text_length: int,
|
|
has_text: bool = True,
|
|
seq_shape: str = "30x48x80",
|
|
) -> torch.Tensor:
|
|
# Check if the specific op is available
|
|
if sta_fwd is None:
|
|
return sliding_tile_attention_triton(q, k, v, window_size, text_length, has_text, seq_shape)
|
|
|
|
seq_length = q.shape[2]
|
|
shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3}
|
|
|
|
if has_text:
|
|
target_size = math.ceil(seq_length / 384) * 384
|
|
pad_size = target_size - seq_length
|
|
if pad_size > 0:
|
|
q = torch.cat([q, q[:, :, -pad_size:]], dim=2)
|
|
k = torch.cat([k, k[:, :, -pad_size:]], dim=2)
|
|
v = torch.cat([v, v[:, :, -pad_size:]], dim=2)
|
|
|
|
output = torch.empty_like(q)
|
|
flag = shape_map[seq_shape]
|
|
|
|
for head_idx, (t, h, w) in enumerate(window_size):
|
|
# Per-head slices are not contiguous in the batch dimension when batch>1
|
|
# (they keep the original head-stride). The TK kernel assumes contiguous
|
|
# [B, H, S, D] layout, so we materialize a contiguous [B,1,S,D] view.
|
|
q_h = q[:, head_idx:head_idx + 1].contiguous()
|
|
k_h = k[:, head_idx:head_idx + 1].contiguous()
|
|
v_h = v[:, head_idx:head_idx + 1].contiguous()
|
|
o_h = torch.empty_like(q_h)
|
|
sta_fwd(q_h, k_h, v_h, o_h, t, h, w, text_length, False, has_text, flag)
|
|
output[:, head_idx:head_idx + 1] = o_h
|
|
|
|
if has_text:
|
|
sta_fwd(q.contiguous(), k.contiguous(), v.contiguous(), output, 3, 3, 3, text_length, True, True, flag)
|
|
|
|
return output[:, :, :seq_length]
|
|
|
|
|
|
def video_sparse_attn(
|
|
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, 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``).
|
|
- 128 -> Triton fallback or CuTe FA4 block-sparse attention.
|
|
- 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 128/256-block paths.
|
|
"""
|
|
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 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 (fused Triton: bf16 read → fp32 accumulate → div → bf16 write)
|
|
q_c = fused_block_mean(q, q_variable_block_sizes, block_elements)
|
|
k_c = fused_block_mean(k, variable_block_sizes, block_elements)
|
|
v_c = fused_block_mean(v, variable_block_sizes, block_elements)
|
|
|
|
scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5)
|
|
attn = torch.softmax(scores, dim=-1)
|
|
# Kept at block resolution, [B, H, q_num_blocks, D]: the coarse result is
|
|
# constant within a block, so the combine broadcasts it over the intra-block
|
|
# axis instead of materializing it across the full sequence.
|
|
out_c = torch.matmul(attn, v_c)
|
|
|
|
# Sparse branch (fused Triton topk mask)
|
|
mask = fused_topk_mask(scores, topk)
|
|
|
|
if block_elements in (128, 256):
|
|
attention = block_sparse_attn_128 if block_elements == 128 else block_sparse_attn_256
|
|
out_s = attention(q, k, v, mask, variable_block_sizes)[0]
|
|
else:
|
|
out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
|
|
|
|
return _combine_coarse_sparse(out_c, out_s, compress_attn_weight, block_elements, seq_dim=2)
|
|
|
|
|
|
def _combine_coarse_sparse(
|
|
out_c: torch.Tensor,
|
|
out_s: torch.Tensor,
|
|
compress_attn_weight: torch.Tensor | None,
|
|
block_elements: int,
|
|
seq_dim: int,
|
|
) -> torch.Tensor:
|
|
"""Combine the block-resolution coarse output with the sparse output.
|
|
|
|
``out_s`` and the optional gate are full-sequence tensors whose sequence axis
|
|
is ``seq_dim`` (2 for [B, H, S, D], 1 for [B, S, H, D]); ``out_c`` has the
|
|
same layout with ``S // block_elements`` blocks on that axis. The coarse
|
|
result is constant within a block, so it broadcasts over the intra-block
|
|
axis of a view that splits the sequence axis into (blocks, block_elements)
|
|
and is never expanded to the full sequence. Splitting one dimension is
|
|
always expressible as a view, so no copy is made whatever the strides of
|
|
``out_s`` or the gate (the BHSD caller may pass a transposed BSHD gate
|
|
directly).
|
|
|
|
Numerics: ungated this is a plain broadcast add and bit-exact with the
|
|
previous ``out_c.repeat(...) + out_s``. Gated, both branches use
|
|
``addcmul``, which multiplies and accumulates in fp32 and rounds once, so
|
|
the result is identical whether or not it runs in place, and it is at least
|
|
as accurate as the old two-rounding ``out_c * w + out_s`` (see
|
|
``tests/test_vsa_combine.py``).
|
|
|
|
In-place contract: when ``out_s`` does not require grad no autograd node has
|
|
saved it (every sparse kernel's node saves its output for backward, so a
|
|
grad-tracking ``out_s`` is never mutated), and within ``video_sparse_attn``
|
|
nothing else aliases it. The combine then accumulates into ``out_s`` and
|
|
returns it, allocating nothing. The in-place path is skipped when the result
|
|
dtype would be promoted, since in place would silently downcast. Callers
|
|
must run the combine in the mode that produced ``out_s``: an inference
|
|
tensor combined outside ``inference_mode`` raises in the in-place update.
|
|
That is not guarded here because ``Tensor.is_inference`` is not traceable by
|
|
``torch.compile`` and would split the graph.
|
|
"""
|
|
full = tuple(out_s.shape)
|
|
q_num_blocks, remainder = divmod(full[seq_dim], block_elements)
|
|
coarse = full[:seq_dim] + (q_num_blocks, ) + full[seq_dim + 1:]
|
|
if remainder != 0 or tuple(out_c.shape) != coarse:
|
|
raise ValueError(f"expected out_c {list(coarse)} for out_s {list(full)} with block_elements="
|
|
f"{block_elements} on dim {seq_dim}, got out_c {list(out_c.shape)}")
|
|
if compress_attn_weight is not None and tuple(compress_attn_weight.shape) != full:
|
|
raise ValueError(f"compress_attn_weight must match out_s {list(full)}, got "
|
|
f"{list(compress_attn_weight.shape)}")
|
|
|
|
blocked = full[:seq_dim] + (q_num_blocks, block_elements) + full[seq_dim + 1:]
|
|
out_c = out_c.unsqueeze(seq_dim + 1)
|
|
out_s_b = out_s.view(*blocked)
|
|
gate_b = None if compress_attn_weight is None else compress_attn_weight.view(*blocked)
|
|
|
|
same_dtype = out_c.dtype == out_s.dtype and (gate_b is None or gate_b.dtype == out_s.dtype)
|
|
if not out_s.requires_grad and same_dtype:
|
|
if gate_b is not None:
|
|
out_s_b.addcmul_(out_c, gate_b)
|
|
else:
|
|
out_s_b.add_(out_c)
|
|
return out_s
|
|
|
|
if gate_b is not None:
|
|
combined = torch.addcmul(out_s_b, out_c, gate_b)
|
|
else:
|
|
combined = out_s_b + out_c
|
|
return combined.view(*full)
|
|
|
|
|
|
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 128/256-block paths; 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 not in (128, 256):
|
|
raise ValueError("video_sparse_attn_bshd is only defined for block_elements=128 or 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: match fused_block_mean's semantics).
|
|
# Padding values are expected to be zero; gradients are broadcast across
|
|
# the full padded block, just like the BHSD fused common path.
|
|
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().sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
|
|
k_c = (k_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
|
|
v_c = (v_c.float().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() # [B, q_num_blocks, H, D]
|
|
|
|
# Sparse branch (fused Triton topk mask + CuTe BSHD).
|
|
mask = fused_topk_mask(scores, topk)
|
|
attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd
|
|
out_s, _ = attention(q, k, v, mask, variable_block_sizes)
|
|
|
|
# Shared combine: out of place when ``out_s`` is saved by the kernel's
|
|
# autograd node (grad), in place otherwise.
|
|
return _combine_coarse_sparse(out_c_blk, out_s, compress_attn_weight, block_elements, seq_dim=1)
|