Skip to content

[perf] Add opt-in SM120 sparse attention forward kernel for MiniMax-H3 - #2048

Open
zhongwujie wants to merge 8 commits into
hao-ai-lab:mainfrom
zhongwujie:feat/h3-sm120-sparse-attention
Open

zhongwujie wants to merge 8 commits into
hao-ai-lab:mainfrom
zhongwujie:feat/h3-sm120-sparse-attention

Conversation

@zhongwujie

@zhongwujie zhongwujie commented Nov 1, 2026 •

Copy link
Copy Markdown

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

  • Kernel (fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_gluon.py): add a BF16 block-sparse attention forward kernel implemented in Triton Gluon, using mma_v2 for QK and PV matrix multiplications. The kernel handles partial tiles and empty selections, using FP32 accumulation.
  • Integration (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-off FASTVIDEO_H3_VSA_GLUON switch to replace the sparse attention while retaining the existing VSA selection and compression logic.
  • Memory (video_sparse_attn_h3.py): release temporary Q/K/V layout copies before allocating the output layout in both native and optimized routes.
  • Validation (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.

cd /root/code/FastVideo
export PYTHONPATH="$PWD/fastvideo-kernel/python${PYTHONPATH:+:$PYTHONPATH}"

pytest fastvideo-kernel/tests/test_vsa_gluon_forward.py -v
pytest fastvideo/tests/attention/test_vsa_h3_gluon.py -v
python fastvideo-kernel/benchmarks/bench_vsa_gluon.py

export MODEL_PATH=/path/to/FastH3-8-Step-V2
bash scripts/benchmarks/minimax_h3_sm120/run_1gpu.sh native
bash scripts/benchmarks/minimax_h3_sm120/run_1gpu.sh gluon

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.

Measurement Native Triton This PR Latency reduction
Sparse forward median 346.87 ms 247.04 ms 28.78%
DiT median 282.72 s 244.05 s 13.68%
E2E median 314.45 s 276.73 s 12.00%

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:

Quantity Maximum absolute error RMSE
Attention output 0.000244140625 1.3047e-7
Log2-LSE (softmax normalization statistic) 1.9073e-6 2.3576e-7

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.

Version 2 s 7 s 12 s
Native Triton Native at 2 s Native at 7 s Native at 12 s
This PR This PR at 2 s This PR at 7 s This PR at 12 s

Checklist

  • I ran pre-commit run --all-files and fixed all issues
  • I added or updated tests for my changes
  • I updated documentation if needed
  • I considered GPU memory impact of my changes

For model/pipeline changes, also check:

  • I verified SSIM regression tests pass

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.
@mergify

mergify Bot commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

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)\]

@mergify mergify Bot added type: perf Performance improvement scope: attention Attention backends (VSA, STA, Flash, etc.) scope: kernel CUDA kernels, fastvideo-kernel scope: docs Documentation scope: model Model architecture (DiTs, encoders, VAEs) labels Nov 1, 2026
@greptile-apps

greptile-apps Bot commented Nov 6, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The PR appears safe to merge; the latest changes introduce no actionable issues.

Summary

Adds a default-off Gluon sparse-attention forward for SM120 BF16 inference.

  • FastH3 can use a Gluon sparse-attention forward on SM120.
  • Sparse-attention routes free temporary layouts before making the output.

No new actionable issues were found.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[H3 tile selection] --> B{Gluon enabled?}
  B -->|No| C[Existing attention route]
  B -->|Yes| D[Check supported inputs]
  D --> E[Gluon sparse forward]
  C --> F[Optional gated compression]
  E --> F
  F --> G[Return attention output]
Loading

Reviews (4) · Last reviewed commit: "[test] Run H3 Gluon configuration checks..." · Reviewed by Greptile

@fv-superagent fv-superagent left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@mergify mergify Bot added the scope: infra CI, tests, Docker, build label Nov 7, 2026

@fv-superagent fv-superagent left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@mergify mergify Bot added the needs-rebase PR has merge conflicts label Nov 8, 2026
Rename the environment variable, internal flag, and benchmark mode.
Update tests, scripts, and documentation to use the Gluon names.
@fv-superagent

Copy link
Copy Markdown
Collaborator

Heads up — the new main makes this PR conflict in one file: docs/inference/fasth3-distilled.md.

Cause: #2004 landed and adds its ## Token reduction section at the same anchor this PR uses for ## Optional Gluon sparse forward on SM120 — right after "Model loading and decoding come on top of this." and before ## Apple Silicon. That is two insertions at one point.

The other overlapping file, fastvideo/models/dits/minimax_h3.py, merges cleanly: main's edits and this PR's two added lines sit in different hunks. I checked both with git apply --3way against main at ce970eb.

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 Silicon

Rebase on main and that is the only conflict to resolve.

— fv-superagent · two headings, one parking spot · run by aryan5v

@mergify mergify Bot removed the needs-rebase PR has merge conflicts label Nov 8, 2026
@zhongwujie

Copy link
Copy Markdown
Author

@fv-superagent The previous unit-test failure has been fixed, and the latest commit (27723cba) resolves the merge conflict with upstream main.

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!

@fv-superagent fv-superagent left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_SM120 reader, doc row or benchmark mode is left; fastvideo/envs.py:689 and the generated docs/contributing/env_vars.md row agree, and the old name was never on main, so nothing released changes.
  • Merge is clean. docs/inference/fasth3-distilled.md keeps both ## Optional Gluon sparse forward on SM120 and main's ## Token reduction; no conflict markers, and head differs from main in exactly the PR's 14 files.
  • On this head, Buildkite #2273 passed all six fastcheck lanes and pre-commit passed. 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:9 skips the whole module off SM120, so test_h3_conflicting_routes:55 never runs in CI even though it needs no GPU. It is the guard that would fail if the new _use_gluon checks 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 fv-superagent left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.py now scopes the SM120 skip to test_h3_dispatch_preserves_mask_and_compression, so test_h3_conflicting_routes runs on non-SM120 CI and exercises the new _use_gluon constructor guards.
  • docs/inference/fasth3-distilled.md restores "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

@zhongwujie

Copy link
Copy Markdown
Author

@fv-superagent Thanks for the latest 5/5 review!

The latest HEAD (442d75b3) has passed fastcheck, pre-commit, and build, and both your review and Greptile report no outstanding issues.

Could you clarify whether anything is still preventing an APPROVED review and /merge, or whether any further action is needed on my side?

Thanks!

@fv-superagent

Copy link
Copy Markdown
Collaborator

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.

@zhongwujie

Copy link
Copy Markdown
Author

@aryan5v Hi! The latest HEAD (442d75b3) has passed fastcheck, pre-commit, and build, with 5/5 reviews from both fv-superagent and Greptile.

fv-superagent confirmed that no further changes are needed, but noted that its earlier CHANGES_REQUESTED review needs to be dismissed by a maintainer before approval and merge can proceed.

Could you please take a look when you have a chance? Thanks!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

scope: attention Attention backends (VSA, STA, Flash, etc.) scope: docs Documentation scope: infra CI, tests, Docker, build scope: kernel CUDA kernels, fastvideo-kernel scope: model Model architecture (DiTs, encoders, VAEs) type: perf Performance improvement

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants