Fix: Extend b12x FP4 GEMM support to SM121 (GB10/DGX Spark) - #3113
Conversation
SM121 (GB10, used in DGX Spark and RTX Pro 6000) supports the same b12x CuTe DSL warp-level MMA FP4 GEMM kernels as SM120, but was excluded from _b12x_gemm_fp4_requirement's supported_compute_capability list. On SM121, calling mm_fp4(..., backend="b12x") raised: BackendSupportedError: mm_fp4 does not support backend 'b12x' with capability 121 Add 121 to the supported compute capability list so that mm_fp4(..., backend="b12x") works correctly on SM121. Signed-off-by: Meenakshi Venkataraman <meenakshiv@nvidia.com>
📝 WalkthroughWalkthroughUpdated FP4 Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
meena-at-work
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/gemm/gemm_base.py (1)
4541-4569:⚠️ Potential issue | 🟡 MinorExtend the auto-selection heuristic to SM121 as well.
Line 4541 enables explicit
backend="b12x"on SM121, but_heuristic_func_mm_fp4still only prefersb12xwhenmajor == 12 and minor == 0. On SM121,backend="auto"will continue falling through tocudnn/cutlasseven thoughb12xis now marked suitable.Suggested follow-up
- is_sm120 = major == 12 and minor == 0 + is_sm12x = major == 12 and minor in (0, 1) - # SM120 + CUDA 13: prefer b12x (warp-level MMA, underfill tile selection) - if is_sm120 and use_nvfp4 and cuda_major >= 13: + # SM120/121 + CUDA 13: prefer b12x (warp-level MMA, underfill tile selection) + if is_sm12x and use_nvfp4 and cuda_major >= 13: return [c for c in ("b12x", "cutlass", "cudnn") if c in suitable_backends]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/gemm/gemm_base.py` around lines 4541 - 4569, The auto-selection heuristic in _heuristic_func_mm_fp4 must be updated to treat SM121 the same as SM120 so backend="auto" prefers b12x when _b12x_gemm_fp4_requirement is valid; locate the conditional that currently checks for major == 12 and minor == 0 and expand it to include minor == 1 (or check minor in {0,1} / minor <= 1) so SM121 is considered eligible for the b12x backend selection.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@flashinfer/gemm/gemm_base.py`:
- Around line 4541-4569: The auto-selection heuristic in _heuristic_func_mm_fp4
must be updated to treat SM121 the same as SM120 so backend="auto" prefers b12x
when _b12x_gemm_fp4_requirement is valid; locate the conditional that currently
checks for major == 12 and minor == 0 and expand it to include minor == 1 (or
check minor in {0,1} / minor <= 1) so SM121 is considered eligible for the b12x
backend selection.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 25b11123-e78f-40f7-9583-db1c35fd71b3
📒 Files selected for processing (1)
flashinfer/gemm/gemm_base.py
![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 expands FP4 GEMM support to include compute capability 121 (Blackwell) within the _b12x_gemm_fp4_requirement function. A review comment identifies that the heuristic selection logic in _heuristic_func_mm_fp4 needs a corresponding update to ensure SM 121 devices are correctly handled when the backend is set to "auto".
|
|
||
|
|
||
| @supported_compute_capability([120]) | ||
| @supported_compute_capability([120, 121]) |
There was a problem hiding this comment.
While adding 121 to the supported compute capability list enables the b12x backend when explicitly requested, the heuristic function _heuristic_func_mm_fp4 (around line 5070) also needs to be updated to include SM121. Currently, it only checks for is_sm120 (major 12, minor 0), which means backend="auto" will not select b12x on SM121 devices even if CUDA 13+ is present. Consider updating the heuristic logic to use major == 12 or _match_sm_version(a.device, ["120", "121"]) to ensure consistent behavior across Blackwell variants.
There was a problem hiding this comment.
Auto not selecting b12x on SM121 devices is intentional at this point.
|
/bot run |
|
[FAILED] Pipeline #48833682: 13/20 passed |
## 📌 Description After #3113 extended the b12x decorator to SM121, calling `mm_fp4(..., backend="b12x")` on sm_121 still trips the dispatch-side check in `dense_blockscaled_gemm_sm120_b12x.py:1591`: ``` ValueError: dense_gemm launch only supports sm_120, got sm_121 ``` SM120 and SM121 share the same 12.x spec (MMA atoms, SMEM), so the same kernel applies. Allow sm_121 through the check. The error message is updated to match. Other call sites in this file (`sm_version="sm_120"` at line 1875, `get_smem_capacity_in_bytes("sm_120")` at lines 122/1461) are left alone — they work for both arches. ## 🔍 Related Issues Refs #3170 (Action Item 4). Follows up on #3113. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). No new tests; no SM121 hardware locally. Existing sm_120 path unchanged. ## Reviewer Notes Same audit cleanup batch as #3173 / #3174 / #3175. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Expanded GPU architecture support to include additional SM12x devices (sm_120 and sm_121). <!-- end of auto-generated comment: release notes by coderabbit.ai -->
3 participants
SM121 (GB10, used in DGX Spark and RTX Pro 6000) supports the same b12x CuTe DSL warp-level MMA FP4 GEMM kernels as SM120, but was excluded from _b12x_gemm_fp4_requirement's supported_compute_capability list. On SM121, calling mm_fp4(..., backend="b12x") raised:
BackendSupportedError: mm_fp4 does not support backend 'b12x' with capability 121
Add 121 to the supported compute capability list so that mm_fp4(..., backend="b12x") works correctly on SM121.
📌 Description
🔍 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
Release Notes