nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf] Add FA4 CuTe backward support for VSA-256 (#1639)

Browse Source 
Co-authored-by: Hyunsung Lee <hyunsungl@sizigistudios.com>
Co-authored-by: alexzms <3036648523@qq.com>
This commit is contained in:
Hyunsung Lee
2026-08-19 11:46:49 -07:00
committed by GitHub
co-authored by Hyunsung Lee alexzms
parent 8537dcd6de
commit 00338aa9ca
13 changed files with 929 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo-kernel/python/fastvideo_kernel/ops.py
+24 -22
View File
@@ -2,6 +2,8 @@ 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,
)
@@ -74,12 +76,13 @@ def video_sparse_attn(
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 256-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)
@@ -119,8 +122,9 @@ def video_sparse_attn(
# Sparse branch (fused Triton topk mask)
mask = fused_topk_mask(scores, topk)
if block_elements == 256:
out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0]
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]
@@ -142,14 +146,14 @@ def video_sparse_attn_bshd(
"""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
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 != 256:
raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 "
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
@@ -171,19 +175,15 @@ def video_sparse_attn_bshd(
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)
# 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() * 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_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()
@@ -195,13 +195,15 @@ def video_sparse_attn_bshd(
# Sparse branch (fused Triton topk mask + CuTe BSHD).
mask = fused_topk_mask(scores, topk)
out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes)
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)
out = out_s
out_view = out.view(batch, q_num_blocks, block_elements, heads, dim)
# Out-of-place: ``out_s`` is the tensor FA4's autograd node saved for its
# backward, so mutating it in place invalidates the graph.
out_view = out_s.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)
out = out_view + out_c_blk.unsqueeze(2) * gate_view
else:
out_view.add_(out_c_blk.unsqueeze(2))
return out
out = out_view + out_c_blk.unsqueeze(2)
return out.view(batch, q_seq_len, heads, dim)