Skip to content

Commit 38412b2

Browse files
[perf]: fuse the MiniMax H3 VSA gate combine, accumulate in place at inference
torch.addcmul keeps one full-sequence temporary under grad instead of two, and when the sparse output is not tracked by autograd the gate branch accumulates into it directly, allocating nothing. The product is no longer rounded to bf16 before the add, so gated values move by at most half a bf16 ulp (towards the fp32 value). test_vsa_h3_backward.py passes on the Triton and FA4 CuTe backends on a GB200. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NjFimeedTQWxgqSzP5xo4j
1 parent 28c68b9 commit 38412b2

1 file changed

Lines changed: 12 additions & 5 deletions

File tree

fastvideo/attention/backends/video_sparse_attn_h3.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -763,11 +763,18 @@ def forward( # type: ignore[override]
763763
out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled) # [B, H, n_tiles, D]
764764
out_c = out_c.permute(0, 2, 1, 3).to(out.dtype) # [B, n_tiles, H, D]
765765
batch, seq_len, heads, dim = out.shape
766-
# Out-of-place: on the CuTe backend ``out`` is the tensor FA4's
767-
# autograd node saved for its backward, so an in-place add here
768-
# bumps its version counter and backward dies with "one of the
769-
# variables needed for gradient computation has been modified".
770766
out_tiled = out.view(batch, n_tiles, tile_elems, heads, dim)
771767
gate_tiled = logical_gate.view(batch, n_tiles, tile_elems, heads, dim)
772-
out = (out_tiled + out_c.unsqueeze(2) * gate_tiled).view(batch, seq_len, heads, dim)
768+
out_c = out_c.unsqueeze(2)
769+
if out.requires_grad or gate_tiled.dtype != out.dtype:
770+
# Out-of-place: ``out`` is the tensor the attention kernel's
771+
# autograd node saved for its backward, so an in-place add here
772+
# bumps its version counter and backward dies with "one of the
773+
# variables needed for gradient computation has been modified".
774+
# Fused addcmul keeps one full-sequence temporary instead of two.
775+
out = torch.addcmul(out_tiled, out_c, gate_tiled).view(batch, seq_len, heads, dim)
776+
else:
777+
# Inference: ``out`` is a fresh kernel output nobody else holds,
778+
# so accumulate into it and allocate nothing.
779+
out_tiled.addcmul_(out_c, gate_tiled)
773780
return out

0 commit comments

Comments
 (0)