Allow grouped GEMM quant without probability tensor - #458
Conversation
📝 WalkthroughWalkthroughGrouped GEMM quant APIs and kernels now accept omitted probability tensors. Probability loading and scaling are conditionally compiled and executed. Tests cover equivalence with unit probabilities and Blackwell/Rubin kernel branching. ChangesOptional probability execution
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant grouped_gemm_quant_api
participant grouped_gemm_quant_kernel
participant d_tensor
Caller->>grouped_gemm_quant_api: execute with prob_tensor=None
grouped_gemm_quant_api->>grouped_gemm_quant_kernel: compile optional probability path
grouped_gemm_quant_kernel->>d_tensor: compute and store unscaled output
d_tensor-->>Caller: return output
🚥 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}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@python/cudnn/grouped_gemm/grouped_gemm_quant/api.py`:
- Around line 1112-1113: Update execute() to validate that runtime probability
presence matches self.prob_desc: reject a provided probability tensor when the
kernel was compiled without one, and reject None when compiled with one. Raise a
clear error before invocation for either mismatch, while preserving execution
for matching forms.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py`:
- Around line 523-524: Add an L0–L4 pytest level marker to
test_grouped_gemm_quant_wrapper_without_prob_tensor, choosing the level that
matches its GPU compilation cost and existing test-suite conventions.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3fd83076-5819-4f74-a9b5-5431949807d8
📒 Files selected for processing (5)
python/cudnn/grouped_gemm/grouped_gemm_quant/api.pypython/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.pypython/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_quant.pytest/python/fe_api/test_rubin_kernel_dispatch.py
| :param prob_tensor: Optional probability tensor for per-row gating. When | ||
| omitted, the kernel compiles out the probability load and multiply. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep the execute-time probability ABI consistent with compilation.
prob_desc selects a compile-time kernel variant, but execute() now accepts either runtime form. Direct API users can compile with sample_prob=None then pass a tensor, or compile with a sample then pass None, causing a mismatched compiled invocation rather than a clear error. Require runtime presence to match self.prob_desc.
Proposed fix
+ if self.prob_desc is None:
+ self._value_error_if(
+ prob_tensor is not None,
+ "prob_tensor must be omitted when the API was compiled without sample_prob",
+ )
+ else:
+ self._value_error_if(
+ prob_tensor is None,
+ "prob_tensor must be provided when the API was compiled with sample_prob",
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| :param prob_tensor: Optional probability tensor for per-row gating. When | |
| omitted, the kernel compiles out the probability load and multiply. | |
| if self.prob_desc is None: | |
| self._value_error_if( | |
| prob_tensor is not None, | |
| "prob_tensor must be omitted when the API was compiled without sample_prob", | |
| ) | |
| else: | |
| self._value_error_if( | |
| prob_tensor is None, | |
| "prob_tensor must be provided when the API was compiled with sample_prob", | |
| ) |
🤖 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_quant/api.py` around lines 1112 -
1113, Update execute() to validate that runtime probability presence matches
self.prob_desc: reject a provided probability tensor when the kernel was
compiled without one, and reject None when compiled with one. Raise a clear
error before invocation for either mismatch, while preserving execution for
matching forms.
| def test_grouped_gemm_quant_wrapper_without_prob_tensor(request): | ||
| """The block-scaled wrapper supports a compile-time unit probability.""" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add a test-level marker.
This test has no L0–L4 marker, so level-based test selection will omit its classification. Add the level appropriate for its GPU compilation cost.
As per coding guidelines, “Mark every new Python test with a level from L0 through L4.”
🤖 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_quant.py` around lines 523
- 524, Add an L0–L4 pytest level marker to
test_grouped_gemm_quant_wrapper_without_prob_tensor, choosing the level that
matches its GPU compilation cost and existing test-suite conventions.
Source: Coding guidelines
saltyminty
left a comment
saltyminty
left a comment
There was a problem hiding this comment.
non-blocking but grouped_gemm_quant docs should also be updated to mark prob_tensor as optional
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-458-78f3aa3 |
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)
python/cudnn/gemm/cutedsl/grouped/quant/api.py (1)
269-270: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAdd dtype and stride validation for optional probability descriptors.
When
self.prob_descis present, this path checks only shape. Add conditional_check_dtypeand_check_tensor_stridechecks forself.prob_descbefore compilation. Keep the checks disabled whenself.prob_desc is None.As per coding guidelines:
python/cudnn/**/api.pyrequirescheck_support()to validate dtype, shape, stride, architecture, and configuration using the provided_check_tensor_*and_value_error_ifhelpers.🤖 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/gemm/cutedsl/grouped/quant/api.py` around lines 269 - 270, Update the validation path around the probability descriptor to conditionally call _check_dtype and _check_tensor_stride for self.prob_desc when it is present, using the existing descriptor and validation conventions. Keep these checks disabled when self.prob_desc is None, and ensure check_support() validates the optional descriptor’s dtype, shape, and stride before compilation.Source: Coding guidelines
♻️ Duplicate comments (1)
python/cudnn/gemm/cutedsl/grouped/quant/api.py (1)
1112-1113:⚠️ Potential issue | 🟠 MajorRe-raise the unresolved execute-time probability ABI check.
prob_tensorselects a compile-time variant throughself.prob_desc, but the removed guard allows either runtime form. A tensor is ignored when the kernel was compiled without probability.Nonereaches a variant compiled to load probability data when probability was provided at compilation. Reject both mismatches before kernel invocation. This is the same unresolved issue from the previous review.Also applies to: 1136-1137
🤖 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/gemm/cutedsl/grouped/quant/api.py` around lines 1112 - 1113, Restore the execute-time ABI validation in the API method handling prob_tensor: compare whether prob_tensor is provided with the compile-time probability variant represented by self.prob_desc, and reject both mismatches before kernel invocation. Preserve valid tensor/probability-enabled and None/probability-disabled combinations.
🤖 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.
Outside diff comments:
In `@python/cudnn/gemm/cutedsl/grouped/quant/api.py`:
- Around line 269-270: Update the validation path around the probability
descriptor to conditionally call _check_dtype and _check_tensor_stride for
self.prob_desc when it is present, using the existing descriptor and validation
conventions. Keep these checks disabled when self.prob_desc is None, and ensure
check_support() validates the optional descriptor’s dtype, shape, and stride
before compilation.
---
Duplicate comments:
In `@python/cudnn/gemm/cutedsl/grouped/quant/api.py`:
- Around line 1112-1113: Restore the execute-time ABI validation in the API
method handling prob_tensor: compare whether prob_tensor is provided with the
compile-time probability variant represented by self.prob_desc, and reject both
mismatches before kernel invocation. Preserve valid tensor/probability-enabled
and None/probability-disabled combinations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ce65e7fd-641e-408b-b8bc-205998cca07f
📒 Files selected for processing (5)
python/cudnn/gemm/cutedsl/grouped/quant/api.pypython/cudnn/gemm/cutedsl/grouped/quant/grouped_gemm_quant.pypython/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_quant.pytest/python/fe_api/test_rubin_kernel_dispatch.py
🚧 Files skipped from review as they are similar to previous changes (2)
- test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py
- test/python/fe_api/test_rubin_kernel_dispatch.py
Anerudhan
added
mod-frontend
4 participants
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
Tests