Fix/3170 dense blockscaled sm12x - #3180
Conversation
📝 WalkthroughWalkthroughThe launch-time architecture validation in the dense blockscaled GEMM kernel is expanded to accept both SM120 and SM121 GPU architectures instead of only SM120. The corresponding error message is updated to reflect the expanded supported architecture set. Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~2 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py (1)
1875-1875: Optional: consider plumbing the actual device SM version instead of hardcoding"sm_120".
dense_gemm()always passessm_version="sm_120"regardless of the running device, so on sm_121 the validation at line 1591 only succeeds because of the hardcoded literal — not because the actual device capability is propagated. Since SM120 and SM121 share the same 12.x spec this is functionally fine today (and the PR description explicitly calls this out as intentional), but it does mean the new"sm_121"branch in the validator is unreachable through this entry point. Ifdense_gemm()is ever called on a non-SM12x device, the hardcoded literal would also silently pass the check. Consider derivingsm_versionfroma_torch.device(e.g. viatorch.cuda.get_device_capability) for symmetry with the b12x decorator's gating; alternatively, leave a brief comment here noting the intentional hardcoding and SM120/SM121 equivalence.No action required for this PR — flagging as a potential follow-up.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` at line 1875, dense_gemm currently hardcodes sm_version="sm_120", which prevents propagating the actual device SM and makes the sm_121 validator branch unreachable; change dense_gemm to derive sm_version from the runtime device (e.g., read a_torch.device and call torch.cuda.get_device_capability to form "sm_{major}{minor}" or similar) so the validator sees the real SM, or if you intentionally want the hardcode leave a short comment next to sm_version="sm_120" explaining the deliberate choice and SM120/SM121 equivalence; reference dense_gemm, the sm_version literal, and the b12x decorator gating when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py`:
- Line 1875: dense_gemm currently hardcodes sm_version="sm_120", which prevents
propagating the actual device SM and makes the sm_121 validator branch
unreachable; change dense_gemm to derive sm_version from the runtime device
(e.g., read a_torch.device and call torch.cuda.get_device_capability to form
"sm_{major}{minor}" or similar) so the validator sees the real SM, or if you
intentionally want the hardcode leave a short comment next to
sm_version="sm_120" explaining the deliberate choice and SM120/SM121
equivalence; reference dense_gemm, the sm_version literal, and the b12x
decorator gating when making the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: da3b31b1-95e2-4da7-b04a-c322205ec4ae
📒 Files selected for processing (1)
flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.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 updates the DenseGemmKernel initialization to support both sm_120 and sm_121 architectures. Feedback suggests that while the check is updated, a hardcoded sm_version in the dense_gemm function call later in the file should be reviewed for consistency to ensure full compatibility with SM121.
|
/bot run |
|
[FAILED] Pipeline #49610128: 11/20 passed |
4 participants
📌 Description
After #3113 extended the b12x decorator to SM121, calling
mm_fp4(..., backend="b12x")on sm_121 still trips the dispatch-side check indense_blockscaled_gemm_sm120_b12x.py:1591: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
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.).No new tests; no SM121 hardware locally. Existing sm_120 path unchanged.
Reviewer Notes
Same audit cleanup batch as #3173 / #3174 / #3175.
Summary by CodeRabbit