[feat] VSA-256 fastpath on Blackwell via FA4 CuTe block-sparse attention - #1354
Conversation
Adds FoundationResearch/flash-attention on the vsapad branch under fastvideo-kernel/include/flash-attention. This carries the CuTe-DSL block-sparse attention forward kernel (flash_attn.cute.block_sparsity) required by the upcoming VSA-256 fastpath on Blackwell. The submodule is Python-only (CuTe DSL JIT); no C++ build step needed.
New fastvideo_kernel.block_sparse_attn_cute_fwd. Provides a thin Python wrapper around flash_attn.cute.interface._flash_attn_fwd that adapts VSA's (block_map, variable_block_sizes) inputs into FA4's BlockSparseTensorsTorch + per-KV-block validity mask. Exposes BHSD and BSHD entrypoints (block_sparse_attn_cute_fwd and block_sparse_attn_cute_fwd_bshd respectively); the BSHD variant is intended for the upcoming VSA-256 fastpath which keeps inputs in BSHD layout to avoid hot-path transposes. Forward only; no autograd registration yet. Requires the FoundationResearch flash-attention submodule plus nvidia-cutlass-dsl and quack-kernels.
Adds VSA's 256-token block path on top of the existing 64-token path.
The new path is intended for FA4's CuTe block-sparse attention forward
kernel on Blackwell, where 256-token KV tiles give ~1.4 PFLOPs forward
throughput on GB200.
`video_sparse_attn(block_size=...)` now auto-dispatches on
`block_elements = prod(block_size)`:
- 64 -> existing index-native TK/Triton path (unchanged behavior).
- 256 -> new `block_sparse_attn_256` wrapper, which expands the
logical 256-block map into the FA4 kernel's 128-token KV
layout and calls into the CuTe BSA forward.
`video_sparse_attn_bshd` is added for callers that already have BSHD
tensors and want to skip the BHSD<->BSHD round-trip on the CuTe hot
path; it is defined only for block_elements=256.
Backend selection honors a small set of opt-in env vars (all unset by
default):
- FASTVIDEO_VSA_TRITON=1 forces Triton in either path. The 256 path
uses a route-A 256->64 expansion.
- FASTVIDEO_VSA_TK=1 prefers the sm_90 TK kernel in the 64 path
(no-op if the extension isn't available).
- FASTVIDEO_VSA_CUTEDSL=1 prefers CuTe in the 256 path (default).
- FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1 is kept as a backward-compat
alias for FASTVIDEO_VSA_TRITON.
mergify
Bot
added
type: feat
Merge ProtectionsYour pull request matches the following merge protections and will not be merged until they are valid. 🔴 PR merge requirementsWaiting for
This rule is failing.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a 256-block sparse attention path optimized for Blackwell GPUs, utilizing the CuTe FA4 block-sparse attention kernel. It introduces new wrappers for 256-block attention, adds a BSHD-native fastpath to reduce memory transposes, and integrates these into the existing VSA backend with configurable tile sizes. Review feedback identifies high-severity issues including a potential JIT compilation error in the CuTe DSL mask logic and shape inconsistencies in the Log-Sum-Exp (LSE) tensor across different layouts. Additionally, the reviewer suggested simplifying redundant backend selection logic and optimizing tensor expansion operations.
| 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) |
There was a problem hiding this comment.
In the CuTe DSL, kv_blk is a scalar SSA value derived from n_idx // block_size_ssa. Indexing it with [0] is likely incorrect and may cause a JIT compilation error. It should be used directly as the index for kv_sizes.
| valid = utils.scalar_to_ssa(kv_sizes[kv_blk[0]], cutlass.Int32) | |
| valid = utils.scalar_to_ssa(kv_sizes[kv_blk], cutlass.Int32) |
| 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() |
There was a problem hiding this comment.
There is a shape inconsistency for the Log-Sum-Exp (LSE) tensor. Flash Attention typically returns LSE in [B, H, S] layout. Since this function transposes the output to BHSD, the LSE should remain in [B, H, S] layout to be consistent. Transposing it to [B, S, H] on line 210 conflicts with the fallback shape on line 205.
| 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() | |
| if lse_bshd is None: | |
| lse = torch.zeros( | |
| (q.shape[0], q.shape[1], q.shape[2]), | |
| dtype=torch.float32, | |
| device=q.device, | |
| ) | |
| else: | |
| lse = lse_bshd.contiguous() |
| 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() |
There was a problem hiding this comment.
In the BSHD variant, the fallback LSE shape on line 238 is [B, H, S] (since q.shape[2] is H and q.shape[1] is S), but the transposed return value on line 243 is [B, S, H]. The fallback should match the expected BSHD-consistent layout [B, S, H].
| 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() | |
| if lse_bshd is None: | |
| lse = torch.zeros( | |
| (q.shape[0], q.shape[1], q.shape[2]), | |
| dtype=torch.float32, | |
| device=q.device, | |
| ) | |
| else: | |
| lse = lse_bshd.transpose(1, 2).contiguous() |
| elif _force_tk(): | ||
| use_sm90 = sm90_available | ||
| else: | ||
| use_sm90 = sm90_available |
There was a problem hiding this comment.
The elif _force_tk(): block is redundant because it performs the exact same assignment as the else: block. This can be simplified to improve maintainability.
else:
use_sm90 = sm90_available| if os.environ.get("FASTVIDEO_VSA_CUTEDSL", "0") == "1": | ||
| return "cutedsl" |
There was a problem hiding this comment.
This explicit check for FASTVIDEO_VSA_CUTEDSL is redundant as it returns the same value as the default case on line 48. Removing it simplifies the logic.
return "cutedsl"| 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 |
There was a problem hiding this comment.
The construction of expanded_sizes can be performed more efficiently and concisely using torch.stack and flatten.
expanded_sizes = torch.stack([child0, child1], dim=1).flatten()forward() now derives block_elements = math.prod(VSA_TILE_SIZE) and, when it is 256 and the CuTe entrypoint is importable, calls video_sparse_attn_bshd directly (inputs already arrive in [B,S,H,D], so the BHSD transpose round-trip is skipped). The default (4,4,4) tile keeps the existing 64-element TK/Triton path byte-for-byte. VSA_TILE_SIZE stays the single module-level constant it already was on main -- no env var, no per-pipeline plumbing, no metadata/config threading. Switching to the Blackwell fastpath is a one-line change to that constant; everything else (build/tile/construct_variable_block_sizes) reads it unchanged.
Four GPU correctness tests covering the new VSA-256 path:
- test_vsa256_forward.py: CuTe forward vs dense torch reference
(qk_equal and qk_diff shapes).
- test_vsa256_forward_vbs.py: CuTe forward with variable KV block sizes
(<256) vs token-masked torch reference.
- test_vsa256_triton.py: Route-A Triton fwd + bwd vs token-masked
torch reference (forces FASTVIDEO_VSA_TRITON).
- test_vsa256_forward_cross.py: Three-way parity: torch ref vs CuTe vs
Triton on the same inputs.
All tests skip cleanly when CUDA is unavailable. They exercise the public
fastvideo_kernel.video_sparse_attn entrypoint with block_size=(4, 8, 8)
and rely on its built-in block_elements=256 dispatch.
There was a problem hiding this comment.
Could you check if this fork is needed?
There was a problem hiding this comment.
I think I will need to merge the kernel level change of FA4 into our trusted repo https://github.com/hao-ai-lab/flash-attention-fp4
There was a problem hiding this comment.
@Davids048 Good catch — confirmed the fork is not needed, and it's been removed.
I traced it through: the submodule pointed at our fork's vsapad branch, which turned out to be Dao-AILab/flash-attention @ c19cd20e plus a single line adding torchvision to cute/pyproject.toml — i.e. zero kernel changes. The flash_attn.cute we build against (interface.py, block_sparsity.py) is byte-identical to that upstream revision, and the VSA-256 parity tests pass against it. (The only real kernel work in the fork is an unrelated dev/blocksize64 experiment that this PR doesn't use — it ships the 256→128 expansion path, i.e. native block=128.)
So this supersedes my earlier note about merging into flash-attention-fp4 — there's nothing to merge. The FA4 CuTe backend is now an optional, lazily-imported dependency (VSA falls back to Triton when it's absent), documented in fastvideo-kernel/README.md (e9320eb) with a pinned upstream install:
pip install "git+https://github.com/Dao-AILab/flash-attention.git@c19cd20e#subdirectory=flash_attn/cute"
Pinned to c19cd20e because that revision's _flash_attn_fwd takes m_block_size/n_block_size (what the wrapper calls); later upstream reshaped it into a tile_mn tuple.
…module Review feedback: the VSA-256 fastpath vendored FA4 CuTe via a git submodule pointing at a personal-org SSH fork (git@github.com:FoundationResearch/flash-attention.git @ vsapad), making it a hard build/clone dependency of the public repo. - Remove the flash-attention submodule (.gitmodules + gitlink). - Lazily import flash_attn.cute (_load_fa4_cute) with a clear actionable error; it is no longer pulled at module load. - Default the VSA-256 backend to Triton; the FA4 CuTe fastpath is opt-in via FASTVIDEO_VSA_CUTEDSL=1 (same optional-dependency model as the NVFP4 FA4 path already on main). - VSA-256 CuTe tests skip cleanly when the optional FA4 CuTe build is absent and explicitly opt into CuTe when present. Follow-up: upstream the block-sparsity delta into hao-ai-lab/flash-attention-fp4 so the CuTe fastpath can depend on it the same optional-pip way as NVFP4.
The VSA-256 fastpath's optional flash_attn.cute dependency is provided upstream by Dao-AILab/flash-attention @ c19cd20e: the installed cute (interface.py + block_sparsity.py) is byte-identical to that revision, so no FoundationResearch fork or kernel patch is required. Document the pinned install in the kernel README. Pin rationale: that revision's _flash_attn_fwd uses m_block_size/n_block_size; later upstream moved to a tile_mn tuple and is not drop-in compatible.
|
Hi @alexzms — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off. TL;DRSubstantial, well-scoped feature PR: clean Verdict: approve-with-followup
Findings[S2-important] CuTe-256 forward has no autograd; training silently broken if 256-tile + CuTe is enabledWhat. Why it matters. PR carries Suggested fix (pick one):
Evidence: [S2-important]
|
3 participants
Summary
Adds a 256-element VSA tile path that dispatches to the FA4 CuTe block-sparse
attention forward kernel on Blackwell, on top of the existing 64-element
TK/Triton path which remains the default and is unchanged.
Measured ~1.4 PFLOPs forward kernel on a single GB200 (bf16, head_dim=128,
seq_len=40960, 8 heads) for the sparse kernel itself; the 64-tile path is
byte-identical to current behavior at the default settings.
Design
VSA tile shape flows through the same channel that already carries
VSA_sparsity:FastVideoArgs.VSA_tile_size(CLI:--VSA-tile-size T H W)→
attn_metadata_builder.build()→VideoSparseAttentionMetadata.VSA_tile_size→ backend
forward().(4, 4, 4)keeps the existing 64-tile path untouched.(4, 8, 8)(256-element blocks) triggers the new CuTe BSHDfastpath inside
forward()automatically (dispatch byblock_elements = prod(tile_size)).Kernel-level backend overrides for debugging / perf comparison (all unset by
default, no behavior change):
FASTVIDEO_VSA_TRITON=1forces Triton (route-A 256→64 expansion for the256 path).
FASTVIDEO_VSA_TK=1prefers the sm_90 TK kernel (64 path only).FASTVIDEO_VSA_CUTEDSL=1prefers CuTe (256 path default).FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1kept as a backward-compat alias forFASTVIDEO_VSA_TRITON.New dependencies (256-tile path only)
fastvideo-kernel/include/flash-attention→FoundationResearch/flash-attention@vsapad. Python-only (CuTe DSLJIT); no C++ build step.
nvidia-cutlass-dsl>=4.3andquack-kernels.The default 64-tile path adds no new dependencies and is unaffected if the
submodule / CuTe DSL is absent.
Test plan
pytest fastvideo-kernel/tests/test_vsa256_forward.py— CuTe forward vs dense torch reference (2/2)pytest fastvideo-kernel/tests/test_vsa256_forward_vbs.py— CuTe variable KV block size (1/1)pytest fastvideo-kernel/tests/test_vsa256_triton.py— Triton fwd+bwd vs torch reference (1/1)pytest fastvideo-kernel/tests/test_vsa256_forward_cross.py— three-way torch/CuTe/Triton parity (1/1)(4,4,4)default path is unchanged vsmain.path is untouched but worth confirming).
Hardware tested
NVIDIA GB200 (sm_100), CUDA 12.9, torch 2.9.1+cu128, aarch64 Linux.
Draft: opening for early review of the dispatch design and the submodule /
dependency story before backward (autograd) support for the CuTe path lands.