[perf] avoid full-sequence materialization in VSA coarse/sparse combine

[boxwrench]

Problem

The VSA compression branch computes a per-block coarse output, then expands it across the full sequence with repeat() before combining it with the sparse output:

out_c = out_c.view(batch, heads, q_num_blocks, 1, dim)
out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim)
...
return out_c * compress_attn_weight + out_s

The coarse value is constant within a block, so the expansion is avoidable. Gated, the combine then allocates two further [B, H, S, D] tensors for the product and the sum — three full-size temporaries where none is needed.

Change

Keep the coarse result at block resolution, [B, H, n_blocks, 1, D], and let it broadcast over the intra-block axis during the combine.

This mirrors the BSHD 128/256 combine, which already broadcasts out_c_blk.unsqueeze(2) instead of repeating. Under no_grad the combine accumulates into the sparse output with addcmul_, removing the remaining two full-sequence temporaries. With grad enabled it stays out-of-place, for the same reason the BSHD path documents in its comment — the sparse output is saved by FA4's autograd node for backward, so mutating it there would invalidate the graph.

The combine is shared by the 64/128/256 dispatch (the repeat() happened before the block-size branch), so all three paths benefit.

Correctness

New tests/test_vsa_combine.py, 17 cases, all CI-sized synthetic tensors with no model, checkpoint or pipeline dependency:

• Ungated: bit-exact against the previous implementation, both for the combine in isolation and end-to-end through video_sparse_attn.
• Gated: addcmul_ fuses the multiply-add instead of rounding the intermediate product to bf16, so the rounding differs from the old path and the two disagree on a fraction of elements. In these tests the fused result is the more accurate of the two when both are compared against an fp32 reference (mean error 1.52e-3 → 1.29e-3, max 2.29e-2 → 1.56e-2). The test asserts accuracy against fp32 rather than agreement with the old rounding.
• Autograd: the grad-enabled path does not mutate the sparse output, and backward still produces gradients.
• Routing and dispatch: top-k selection is unchanged, shapes and dtypes are unchanged, the 64-block path still calls block_sparse_attn, and the 128/256 kernels are still dispatched.
• Also covers a non-contiguous sparse output and asserts the combine no longer allocates a full-sequence temporary.

Performance

benchmarks/bench_vsa_combine.py. These numbers are an isolated microbenchmark of the combine alone — not an end-to-end workload. At an H3-like shape (bf16, 56 heads, head dim 128, 15488 tokens, block size 64):

	latency	peak allocated
gated, before	4.06 ms	848 MiB
gated, after	2.12 ms	212 MiB
ungated, before	2.72 ms	636 MiB
ungated, after	1.56 ms	212 MiB

The 636 MiB saved in the gated case is exactly the three full [B, H, S, D] bf16 tensors that are no longer materialized (3 × 211.8 MiB). The residual 212 MiB is the benchmark's own input clone.

For a full-workload data point, MiniMax H3 at 864×480/124f with topk 0.20 measured end-to-end:

	s per transformer forward	peak allocated
before	7.55 s	25.885 GiB
after	7.23 s	25.473 GiB

That end-to-end measurement is from an AMD Radeon AI PRO R9700 (gfx1201, ROCm 7.2.1, PyTorch 2.9.1, Triton 3.5.1) and is hardware-specific — it should not be extrapolated to other GPUs. The change itself is not platform-specific; the mechanism is allocation count and memory traffic.

Notes

• video_sparse_attn_bshd still uses .float().sum(dim=2) for its block means, which materializes an fp32 copy of q/k/v. That is a separate concern on a 128/256-only path and is left alone here to keep this diff minimal.
• Existing test suite: no new failures. On the machine used here the suite fails 99 tests both before and after this change, with identical failure sets (96 test_vmoba_correctness, 2 test_attn_qat_train, and one intermittent test_vsa_varlen.py::TestVSAVarlenBackward::test_backward_different_lengths that also fails on unmodified main). Those are unrelated to this change and appear to be ROCm-environment failures.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[github-actions]

Welcome to FastVideo! Thanks for your first pull request.

How our CI works:

PRs run a three-tier CI system:

1. Pre-commit — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
2. Fastcheck — six core GPU lanes run automatically via Buildkite (~10-15 min).
3. Merge gate — a reviewer adds ready; changed paths select only the relevant integration, training, golden, or SSIM coverage.

Before your PR is reviewed:

• pre-commit run --all-files passes locally
• You've added or updated tests for your changes
• The PR description explains what and why

If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and merge-gate results appear in the Checks section below.

Useful links:

• Contributing Guide
• Development Roadmap
• Slack

[SolitaryThinker]

Pushed three follow-up commits after review (all verified on a GB200):

1. Kernel (f9ff387f): one shared _combine_coarse_sparse for both layouts, so the 128/256 BSHD entry also stops allocating combine temporaries at inference. The grad branch uses out-of-place torch.addcmul (one temporary instead of two, bit-identical to the in-place path, so results no longer depend on the caller's grad mode). The in-place path is keyed on out_s.requires_grad rather than global grad mode and skipped on dtype promotion. Strict shape checks on out_s and the gate (same-numel wrong-layout inputs used to be accepted silently); the dead .contiguous() fallback is gone since splitting the sequence axis is always a view. torch.compile(fullgraph=True)-safe. Tests rewritten: fp32-oracle bounds in all four grad/no-grad modes, gradient parity with the old path, exact allocation pins (0 bytes in place, one full tensor out of place), end-to-end value/routing/gradient checks at 64/128/256 tiles in both layouts (77 cases on the Triton route, 44 on FA4 CuTe).
2. video_sparse_attn.py (28c68b9d): the gate is passed to the kernel as a transposed view instead of a contiguous copy, saving one full-sequence copy per layer with bit-identical output (1x39936x12x128: 639.8 -> 521.8 MiB peak, 11.64 -> 11.29 ms per call).
3. H3 backend (38412b28): same fused combine; gated values move by at most half a bf16 ulp, towards the fp32 value.

The benchmark now uses triton.testing.do_bench and measures peak memory outside the timed window (the previous numbers included an out_s.clone() inside both). At 56x15488x128 bf16, gated:

	latency us	peak alloc MiB
old combine	366	636
new, inference (in place)	304	0
new, training (addcmul)	304	212

End-to-end video_sparse_attn latency is unchanged within noise (the sparse kernel dominates); per-call peak drops by 2-3 full tensors on the 64-tile path and by about one on the CuTe 256-tile BSHD path. One note on the description: the MiniMax H3 backend does not call this combine (it has its own, now fused in commit 3), so the H3 end-to-end table cannot come from this diff; the beneficiaries are the Wan-family 64-tile path and the 128/256 paths.