Integrate Rubin wgrad cutedsl kernel and fix an issue in rebased rubin branch - #456
Merged
Merged
Conversation
Signed-off-by: qiyuw <qiyuw@nvidia.com>
…tream Rubin, and added MXFP8 discrete-accumulation coverage
WanZzzzzz
changed the title
Contributor
📝 WalkthroughWalkthroughRubin-specific grouped GEMM Wgrad dispatch and validation were added, alongside an SM107 kernel specialization with TMEM planning and SF-window execution changes. Cache keys now include device type, and tests cover Rubin validation, dispatch, accumulated output, and SM107 configurations. ChangesRubin grouped GEMM Wgrad
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant API
participant RubinValidation
participant RubinKernel
participant TMEM
API->>RubinValidation: validate dtype, vector size, accumulation, and layout
RubinValidation->>RubinKernel: construct SM107 kernel
RubinKernel->>TMEM: create TMEM plan and staged layouts
RubinKernel->>RubinKernel: execute SF-window MMA accumulation
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (4)
python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py (1)
251-269: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRun Black with the repo's 160-column line length.
These parenthesized single-value assignments (and several call sites above) are wrapped at roughly 80 columns; Black at 160 collapses them, so the file will churn on the next format run.
As per coding guidelines: "Format Python code and notebooks with Black using a line length of 160."
🤖 Prompt for 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. In `@python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py` around lines 251 - 269, Run Black with the repository’s 160-column line length on the affected file, allowing parenthesized single-value assignments and nearby call sites to collapse to Black’s canonical formatting. Do not manually reformat beyond the formatter’s output.Source: Coding guidelines
test/python/fe_api/test_rubin_kernel_dispatch.py (1)
183-191: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
uint8cases to cover the packed-FP4 branch.
_is_supported_rubin_quantizationtreatstorch.uint8as FP4, but no parameter exercises it.💚 Suggested params
("float4_e2m1fn_x2", "float8_e4m3fn", 16, True), + ("uint8", "float8_e4m3fn", 16, True), + ("uint8", "float8_e8m0fnu", 32, True), + ("uint8", "float8_e4m3fn", 32, False), ("float4_e2m1fn_x2", "float8_e8m0fnu", 32, True),🤖 Prompt for 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. In `@test/python/fe_api/test_rubin_kernel_dispatch.py` around lines 183 - 191, Add uint8 parameter cases to the quantization test parameter list covering the packed-FP4 branch exercised by _is_supported_rubin_quantization, including the relevant scale dtypes, block sizes, and expected boolean outcomes. Keep the existing FP4 and FP8 cases unchanged.test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py (1)
176-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRubin FP4 K-padding rule is duplicated across two test files. The Rubin
FIX_PAD_SIZE = 256constraint is encoded twice with different spellings (compute_capability == 107vstorch.cuda.get_device_capability() == (10, 7)) and separate K lists, so a change to the padding granularity has to be found in both places.
test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py#L176-L180: extract a helper such asrubin_fp4_group_k_list(ab_dtype, base_list)(or a_requires_rubin_fp4_k_padding(ab_dtype)predicate) that owns the capability check and the 256 granularity, and use it here.test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py#L531-L534: replace the inline(10, 7)/ FP4 check with a call to that helper for the runtime K list.🤖 Prompt for 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. In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py` around lines 176 - 180, The Rubin FP4 K-padding rule is duplicated and uses inconsistent capability checks. In test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py:176-180, add a shared helper or predicate that owns the Rubin capability check, FP4 detection, and 256-token granularity, then use it to derive the grouped K list. In test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py:531-534, replace the inline capability/FP4 condition with that helper while preserving the existing runtime K-list behavior.python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py (1)
30-36: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer the shared FP4x2 predicate over a bare
uint8dtype check.
torch.uint8only means packed FP4 wheninterpret_uint8_as_fp4x2is set. It happens to hold here (line 80 hardcodesTrue), but the helper is also called directly fromapi.py, so the assumption is invisible at the call site. Passing anis_fp4flag (derived fromself._is_fp4x2(self.a_desc)) keeps this aligned withAPIBase._is_fp4x2.♻️ Suggested signature change
-def _is_supported_rubin_quantization(ab_dtype: torch.dtype, sf_dtype: torch.dtype, sf_vec_size: int) -> bool: - is_fp4 = ab_dtype in (torch.float4_e2m1fn_x2, torch.uint8) - if is_fp4: +def _is_supported_rubin_quantization(ab_dtype: torch.dtype, sf_dtype: torch.dtype, sf_vec_size: int, is_fp4: Optional[bool] = None) -> bool: + if is_fp4 is None: + is_fp4 = ab_dtype in (torch.float4_e2m1fn_x2, torch.uint8) + if is_fp4: return (sf_dtype == torch.float8_e4m3fn and sf_vec_size == 16) or ( sf_dtype == torch.float8_e8m0fnu and sf_vec_size == 32 )🤖 Prompt for 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. In `@python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py` around lines 30 - 36, Update _is_supported_rubin_quantization to accept an explicit is_fp4 flag instead of treating every torch.uint8 ab_dtype as packed FP4; derive the flag from APIBase._is_fp4x2(self.a_desc) at each relevant caller, including the API path and the hardcoded configuration path, and preserve the existing dtype/vector-size checks based on that flag.
🤖 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.
Nitpick comments:
In `@python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py`:
- Around line 30-36: Update _is_supported_rubin_quantization to accept an
explicit is_fp4 flag instead of treating every torch.uint8 ab_dtype as packed
FP4; derive the flag from APIBase._is_fp4x2(self.a_desc) at each relevant
caller, including the API path and the hardcoded configuration path, and
preserve the existing dtype/vector-size checks based on that flag.
In
`@python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py`:
- Around line 251-269: Run Black with the repository’s 160-column line length on
the affected file, allowing parenthesized single-value assignments and nearby
call sites to collapse to Black’s canonical formatting. Do not manually reformat
beyond the formatter’s output.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py`:
- Around line 176-180: The Rubin FP4 K-padding rule is duplicated and uses
inconsistent capability checks. In
test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py:176-180, add a
shared helper or predicate that owns the Rubin capability check, FP4 detection,
and 256-token granularity, then use it to derive the grouped K list. In
test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py:531-534, replace the
inline capability/FP4 condition with that helper while preserving the existing
runtime K-list behavior.
In `@test/python/fe_api/test_rubin_kernel_dispatch.py`:
- Around line 183-191: Add uint8 parameter cases to the quantization test
parameter list covering the packed-FP4 branch exercised by
_is_supported_rubin_quantization, including the relevant scale dtypes, block
sizes, and expected boolean outcomes. Keep the existing FP4 and FP8 cases
unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e1ed2b5a-7b3e-410a-8607-5f3551cf8eb4
📒 Files selected for processing (7)
python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/api.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.pytest/python/fe_api/test_rubin_kernel_dispatch.py
Merged
2 tasks
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Why
Related issues
API and compatibility impact
Testing
Summary by CodeRabbit
New Features
Bug Fixes
Tests