Repository navigation
[perf] Add opt-in SM120 sparse attention forward kernel for MiniMax-H3 - #2048
zhongwujie wants to merge 8 commits into
Conversation
This kernel overlaps asynchronous K/V copies with attention computation and includes numerical tests, benchmarks and documentation. On RTX PRO 6000, it reduces sparse-forward latency by 28.78%, DiT latency by 13.68% and E2E latency by 12.00%. It is disabled by default and is not bit-exact. Numerical tests passed, and the five-prompt mean MS-SSIM was 0.6384.
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
[Medium impact] The PR appears safe to merge; the latest changes introduce no actionable issues.
|
fv-superagent
left a comment
There was a problem hiding this comment.
fv-superagent score: 3/5
Unit lane is red on this head, and it is this PR's change.
Build #2092 (commit 9dec65453) ends the unit lane at 2 failed, 3481 passed, and both failures are the same test:
FAILED fastvideo/tests/attention/test_minimax_h3_fine_override.py::test_compile_warns_that_the_fine_override_is_ignored[None]
FAILED fastvideo/tests/attention/test_minimax_h3_fine_override.py::test_compile_warns_that_the_fine_override_is_ignored[override1]
E AttributeError: 'types.SimpleNamespace' object has no attribute '_sm120'
fastvideo/attention/backends/video_sparse_attn_h3.py:731
Cause: prepare_for_compile now reads self._sm120 (fastvideo/attention/backends/video_sparse_attn_h3.py:731), but the test calls the unbound method on a stub built as SimpleNamespace(fine_attention_override=override, layer_idx=3) (fastvideo/tests/attention/test_minimax_h3_fine_override.py:72), so the new attribute is missing. The other red jobs in that build are broken only because the unit lane failed first.
Fix: add _sm120=False to that stub, and to the other SimpleNamespace impls in the same file if they reach this path; or read the flag as getattr(self, "_sm120", False).
Greptile's comment-rule findings (fastvideo-kernel/tests/test_vsa_sm120.py:68 and the del out_bhsd comment) also still stand; those are P2s and I have not touched them.
— fv-superagent, automated maintainer run by aryan5v
Move the kernel into triton_kernels behind a lazy public interface. Separate kernel and H3 tests, rename the benchmark, and update docs.
fv-superagent
left a comment
There was a problem hiding this comment.
fv-superagent score: 4/5 (was 3/5)
Re-reviewed head a6cb034. The blocker is fixed exactly as suggested: the stub is now SimpleNamespace(fine_attention_override=override, layer_idx=3, _sm120=False) (fastvideo/tests/attention/test_minimax_h3_fine_override.py:72), and build #2265 is green across all six fastcheck lanes plus pre-commit. I re-ran Greptile on this head.
The only item still open from the earlier review is Greptile's comment-rule P2s (fastvideo-kernel/tests/test_vsa_sm120.py:68 and the del out_bhsd comment in video_sparse_attn_h3.py); make sure its new pass is clean or those are removed before merge.
— fv-superagent · the stub finally learned what SM120 means · run by aryan5v
Rename the environment variable, internal flag, and benchmark mode. Update tests, scripts, and documentation to use the Gluon names.
|
Heads up — the new Cause: #2004 landed and adds its The other overlapping file, Suggested resolution — keep both sections, either order: GPU; the output is identical. Model loading and decoding come on top of this.
+## Token reduction
+...main's section...
+
+## Optional Gluon sparse forward on SM120
+...this PR's section...
+
## Apple SiliconRebase on — fv-superagent · two headings, one parking spot · run by aryan5v |
|
@fv-superagent The previous unit-test failure has been fixed, and the latest commit ( Fastcheck and pre-commit are both green on the latest HEAD, and Greptile reports no remaining actionable issues. Could you please re-review the latest HEAD when possible? The PR is currently mergeable, but further upstream changes may introduce new conflicts. Thanks! |
There was a problem hiding this comment.
fv-superagent score: 4/5 (was 4/5)
Re-reviewed head 27723cba; the rename and the merge are both correct, and the earlier unit-test blocker stays fixed. Nothing here blocks a merge.
- Rename is complete. No
VSA_SM120reader, doc row or benchmark mode is left;fastvideo/envs.py:689and the generateddocs/contributing/env_vars.mdrow agree, and the old name was never onmain, so nothing released changes. - Merge is clean.
docs/inference/fasth3-distilled.mdkeeps both## Optional Gluon sparse forward on SM120and main's## Token reduction; no conflict markers, and head differs frommainin exactly the PR's 14 files. - On this head, Buildkite #2273 passed all six fastcheck lanes and
pre-commitpassed. Greptile's fresh pass on this head reports no new actionable issue; the earlier comment-rule P2 is gone from the code, and no P1 is open.
Two minor, non-blocking items:
fastvideo/tests/attention/test_vsa_h3_gluon.py:9skips the whole module off SM120, sotest_h3_conflicting_routes:55never runs in CI even though it needs no GPU. It is the guard that would fail if the new_use_gluonchecks disappeared; scoping the skip to the CUDA tests would give it real coverage.- The merge rewrote the Gluon doc section and dropped the explicit note that backward, DiT compilation and conflicting fine-attention overrides are rejected; one sentence brings the code-enforced limits back.
— fv-superagent · renamed before it shipped, so nobody noticed · run by aryan5v
Limit the SM120 skip to GPU tests and clarify inference restrictions.
fv-superagent
left a comment
There was a problem hiding this comment.
fv-superagent score: 5/5 (was 4/5)
Re-reviewed head 442d75b3; both items from the last review are fixed exactly.
fastvideo/tests/attention/test_vsa_h3_gluon.pynow scopes the SM120 skip totest_h3_dispatch_preserves_mask_and_compression, sotest_h3_conflicting_routesruns on non-SM120 CI and exercises the new_use_gluonconstructor guards.docs/inference/fasth3-distilled.mdrestores "Backward, DiT compilation and conflicting fine-attention overrides are rejected."
On this head, Buildkite #2278 passed all six fastcheck lanes, and pre-commit and build are green; Greptile reports 5/5 with no actionable issue. Nothing outstanding.
— fv-superagent · one less test that never ran · run by aryan5v
|
@fv-superagent Thanks for the latest 5/5 review! The latest HEAD ( Could you clarify whether anything is still preventing an Thanks! |
|
Nothing needed on your side. On 442d75b everything is green - all six fastcheck lanes, pre-commit, build, Greptile 5/5 and our review at 5/5. Two repo-side gates remain: an approving review, and full-suite-passed, which Mergify runs from that approval. One wrinkle worth knowing: GitHub still shows "Changes requested" from my earlier review of an older head, and later comment reviews do not clear that flag, so a maintainer needs to dismiss it alongside approving. |
|
@aryan5v Hi! The latest HEAD ( fv-superagent confirmed that no further changes are needed, but noted that its earlier Could you please take a look when you have a chance? Thanks! |
Purpose
Add an opt-in BF16 sparse-attention forward kernel for FastH3 on RTX PRO 6000 Blackwell (SM120). Compared with native Triton, it reduces synthetic sparse-forward latency by 28.78%. For FastH3-8-Step-V2, median DiT latency decreases by 13.68% and E2E latency by 12.00%. The kernel uses a fixed launch configuration to avoid runtime autotuning of the sparse forward. The reported latency reductions are measured after warmup and do not include savings from avoiding autotuning.
Changes
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_gluon.py): add a BF16 block-sparse attention forward kernel implemented in Triton Gluon, usingmma_v2for QK and PV matrix multiplications. The kernel handles partial tiles and empty selections, using FP32 accumulation.fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py,fastvideo/attention/backends/video_sparse_attn_h3.py,fastvideo/envs.py,fastvideo/models/dits/minimax_h3.py): expose a lazily imported forward interface and add the default-offFASTVIDEO_H3_VSA_GLUONswitch to replace the sparse attention while retaining the existing VSA selection and compression logic.video_sparse_attn_h3.py): release temporary Q/K/V layout copies before allocating the output layout in both native and optimized routes.fastvideo-kernel/tests/test_vsa_gluon_forward.py,fastvideo/tests/attention/test_vsa_h3_gluon.py,fastvideo-kernel/benchmarks/bench_vsa_gluon.py,scripts/benchmarks/minimax_h3_sm120/): add separate kernel and H3 integration tests, a CUDA-event microbenchmark and a paired single-GPU generation script.The kernel supports eager inference only and is not bit-exact with native Triton.
Test Plan
Tests were run on one NVIDIA RTX PRO 6000 Blackwell Server Edition with PyTorch 2.12.0+cu130 and Triton 3.7.0.
Test Results
All tests and benchmarks ran on one RTX PRO 6000 GPU (SM120, 96 GB) with PyTorch 2.12.0+cu130 and Triton 3.7.0. FastH3-8-Step-V2 used BF16, 1344×768 resolution, 362 output frames at 24 FPS and eight DiT forwards. Kernel-level numerical tests, edge-case checks and H3 integration tests passed.
The kernel-only microbenchmark uses synthetic BHSD inputs
[1, 56, 115328, 128]at sparsity 0.8, with 20 measurements after warmup. Model timings are medians of three generations using one fixed prompt and seed 42, with identical AdaLN cache and CPU-offload settings. E2E includes transfers, decoding and saving, but excludes initial generator construction and warmup. Both routes peak at approximately 65.98 GiB allocated during measured denoising.Microbenchmark errors against native Triton:
The low RMSE values indicate close numerical agreement with native Triton on this input. The results are not bit-exact, and video consistency is evaluated separately below.
SSIM test
Five fixed prompts were tested at seed 42, with one native and one optimized video per prompt. The average of their per-prompt mean MS-SSIM scores was 0.638419, measured on aligned frames without resizing or truncation.
Example frames
Prompt: "A red fox walks slowly through a snowy forest, leaving footprints in fresh snow. A stationary wide shot shows the fox's full body and detailed fur. Soft wind and crisp footsteps are audible."
Resolution: 1344×768. Duration: approximately 15.08 seconds (362 frames at 24 FPS). Seed: 42. Paired frames are sampled at the same timestamps.
Checklist
pre-commit run --all-filesand fixed all issuesFor model/pipeline changes, also check: