feat(attention): enable SM121 (GB10/DGX Spark) for NVFP4 attention - #3897
Conversation
📝 WalkthroughWalkthroughNVFP4 attention now accepts SM121 alongside SM120 across CUDA validation, AOT generation, JIT compilation flags, Python entry points, benchmark gating, and test gating. ChangesSM120/SM121 support expansion
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request extends support for the NVFP4 attention SM120 kernel to SM121 (compute capability 12.1) devices. This involves updating capability checks across benchmarks, CUDA bindings, JIT compilation, API decorators, and tests. The reviewer suggested using the existing is_sm12x_supported utility function in the test suite to simplify the check and make it more future-proof.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| from flashinfer.utils import is_sm120a_supported, is_sm121a_supported | ||
|
|
||
|
|
||
| def _require_sm120(): | ||
| if not torch.cuda.is_available(): | ||
| pytest.skip("CUDA is required") | ||
| if not is_sm120a_supported(torch.device("cuda")): | ||
| pytest.skip("SM120 GPU is required") | ||
| device = torch.device("cuda") | ||
| if not (is_sm120a_supported(device) or is_sm121a_supported(device)): | ||
| pytest.skip("SM120 or SM121 GPU is required") |
There was a problem hiding this comment.
Instead of importing and checking individual SM120 and SM121 support functions, we can use the existing is_sm12x_supported utility function. This simplifies the check and automatically supports any future SM12x variants without requiring code changes.
| from flashinfer.utils import is_sm120a_supported, is_sm121a_supported | |
| def _require_sm120(): | |
| if not torch.cuda.is_available(): | |
| pytest.skip("CUDA is required") | |
| if not is_sm120a_supported(torch.device("cuda")): | |
| pytest.skip("SM120 GPU is required") | |
| device = torch.device("cuda") | |
| if not (is_sm120a_supported(device) or is_sm121a_supported(device)): | |
| pytest.skip("SM120 or SM121 GPU is required") | |
| from flashinfer.utils import is_sm12x_supported | |
| def _require_sm120(): | |
| if not torch.cuda.is_available(): | |
| pytest.skip("CUDA is required") | |
| device = torch.device("cuda") | |
| if not is_sm12x_supported(device): | |
| pytest.skip("SM120 or SM121 GPU is required") |
There was a problem hiding this comment.
Current change is not too messy so this should be fine.
|
/bot run |
|
/bot run tests/attention |
saltyminty
left a comment
saltyminty
left a comment
There was a problem hiding this comment.
Non-blocking, but do we need to add sm121 to the JIT-cache wheels in .github/workflows/release.yml?
| ] | ||
| include_paths: list[str | Path] = [] | ||
| include_paths.extend(_nvfp4_attention_sm120_include_paths()) | ||
| nvcc_flags = current_compilation_context.get_nvcc_flags_list( |
There was a problem hiding this comment.
This change causes sm120 to no longer be supported on CUDA 12.8, as _normalize_cuda_arch enforces >= 12.9. Is this intended?
There was a problem hiding this comment.
Right thanks @saltyminty, this was unintended. Fixed in the latest commit with a gated compilation-context path
|
[FAILED] Pipeline #57442248: 7/20 passed |
|
Verified this head (a5922cd) on an RTX PRO 6000 (SM120, CUDA 13.0): all 7 tests pass and the JIT build.ninja now emits -gencode=arch=compute_120f,code=sm_120f, so SM120 is functionally unchanged on CUDA >= 12.9. On the 12.8 question: the gate is compilation_context._normalize_cuda_arch (raises for any SM12x below 12.9, compilation_context.py:56), and the same get_nvcc_flags_list(supported_major_versions=[12]) path already governs xqa, mla, gemm and fused_moe on main -- a CUDA 12.8 + SM120 install has already lost those modules, so this just aligns nvfp4_attention with the existing convention (the old sm120a flags were the last 12.8-compatible holdout besides fp4_quantization, which #3906 migrates). Since #3838 landed yesterday, a rebase would pick up its 5 extra tests for the SM121 leg too. |
It's a good point, and in fact there is #3907 that is WIP adding the wheel |
|
/bot run tests/attention |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/jit/nvfp4_attention_sm120.py`:
- Around line 93-99: Update the SM12 architecture handling in
_normalize_cuda_arch and the current_compilation_context flow so CUDA 12.8 with
SM120 can reach the sm120a_nvcc_flags fallback without raising. Preserve raw
SM12 capability for this path, while retaining normalization for CUDA 12.9+ and
existing behavior for other architectures.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8b8e367b-ac71-40e6-ba67-3bb172ddb339
📒 Files selected for processing (4)
csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cuflashinfer/jit/nvfp4_attention_sm120.pyflashinfer/nvfp4_attention_sm120.pytests/attention/test_nvfp4_attention_sm120.py
🚧 Files skipped from review as they are similar to previous changes (3)
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu
- tests/attention/test_nvfp4_attention_sm120.py
- flashinfer/nvfp4_attention_sm120.py
|
/bot run tests/attention |
4 participants
📌 Description
Enables the SM120 NVFP4 attention kernel (#3640) on SM121 (GB10 / DGX Spark). The kernel code needs no changes — it compiles and runs correctly for
sm_121a; this PR only widens the SM120-only gatespytest tests/attention/test_nvfp4_attention_sm120.py: 7/7 pass on this branch; 13/13 pass with fix(attention): correct SM120 NVFP4 qk_correction layout, row-sum reduction, and lse #3838's fixes and regression tests applied on top.Speedup is stable across head counts (2.2–2.6× for H=1…32 at D128); throughput saturates from B·H ≥ 16 at S4096. Speedup is around 1.5-1.6x for head dim 64.
Commands to reproduce the perf numbers
NVFP4 (shape lists zip together; one row printed per config):
The "NVFP4" column is the
attention_onlynumber (CUDA-graph replay, pure kernel time,quantization excluded);
end_to_endadditionally includesquantize_qkveach iteration.FA2 BF16 baseline (identical shapes; CUPTI kernel timing is on by default; uniform
full-length sequences — do not pass
--random_actual_seq_len):Environment: NVIDIA GB10 (SM121), CUDA 13.0, torch 2.11. FA2 median of 30 iters
(CUPTI); NVFP4 median of 10 iters (CUDA-graph). Comparing FA2 kernel time against
NVFP4
attention_onlyis apples-to-apples — both exclude host launch overhead.🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Bug Fixes