[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