Add CuTe DSL Rubin (sm107) grouped GEMM GLU/dGLU/quant kernels - #447
Conversation
* address comments Signed-off-by: qiyuw <qiyuw@nvidia.com> * further fix reduxkind issue related to internal cutlass Signed-off-by: qiyuw <qiyuw@nvidia.com> * unit test and skill update Signed-off-by: qiyuw <qiyuw@nvidia.com> * fix minor issue * rubin kernel support * disable epilog prefetch Signed-off-by: qiyuw <qiyuw@nvidia.com> * fix layout issue * war for rubin for api issue
…ntend!2237) * Support unit probability in grouped GEMM GLU kernels
…M kernels
develop grew a use_single_group_runtime_offsets kernel option after feature/rubin
forked, and the three *_rubin.py kernels were written against the older
constructor signature. Because those files are new on this branch, git applied
them with no conflict, so the mismatch only appeared at runtime:
TypeError: BlockScaledMoEGroupedGemmQuantKernel.__init__() got an
unexpected keyword argument 'use_single_group_runtime_offsets'
The call sites forward the option unconditionally, so this fires even when it is
False -- i.e. on every Rubin grouped GEMM compile, not just when the option is
requested. On an sm107 node this took test_grouped_gemm_quant.py from 138 passed
(develop baseline) to 131 failed.
Forward the option only to kernels that implement it, and raise NotImplementedError
when it is explicitly requested on Rubin rather than quietly running a different
schedule than the caller asked for -- the same capability-gating pattern this
branch already uses for geglu_alpha and the GLU clamp limits.
Comparing constructor signatures across all three develop/Rubin kernel pairs
confirms this was the only develop-side parameter the Rubin kernels reject.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdds SM107/Rubin device detection and internal kernel dispatch for grouped GEMM GLU, dGLU, and quantization APIs. New Rubin kernels support persistent MoE scheduling, block scaling, activation or backward fusion, SFD quantization, optional reductions, and architecture-specific compile/runtime signatures. ChangesRubin architecture dispatch
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Anerudhan
added
cat-feature
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-447-8a0e06e |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (12)
skills/cutedsl-kernel-integration/SKILL.md (1)
40-40: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInvoke the test from
test/pythonsopytest.ini/conftest.pyapply.As per coding guidelines, "Run Python tests from
test/pythonsopytest.iniandconftest.pyapply, and preserve the import-order and environment-variable requirements defined byconftest.py."📝 Suggested wording
-- For architecture-dispatch work, also run `pytest test/python/fe_api/test_rubin_kernel_dispatch.py`. On Rubin hardware, the existing FE API e2e tests for the affected operation should still pass without API changes. +- For architecture-dispatch work, also run `pytest fe_api/test_rubin_kernel_dispatch.py` from `test/python`. On Rubin hardware, the existing FE API e2e tests for the affected operation should still pass without API changes.🤖 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 `@skills/cutedsl-kernel-integration/SKILL.md` at line 40, Update the architecture-dispatch testing instruction in SKILL.md to run pytest from the test/python working directory, ensuring pytest.ini and conftest.py are applied and their import-order and environment-variable requirements are preserved.Source: Coding guidelines
skills/cutedsl-kernel-integration/references/integration-pattern.md (1)
97-101: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the examples list consistent.
The first bullet pairs default + Rubin modules; the next two list only the Rubin file, so the pairing convention the section is illustrating is lost.
📝 Suggested wording
- `moe_blockscaled_grouped_gemm_quant.py` + `moe_blockscaled_grouped_gemm_quant_rubin.py` -- `moe_blockscaled_grouped_gemm_glu_rubin.py` -- `moe_blockscaled_grouped_gemm_dglu_rubin.py` +- `moe_blockscaled_grouped_gemm_glu_bias.py` + `moe_blockscaled_grouped_gemm_glu_rubin.py` +- `moe_blockscaled_grouped_gemm_dglu_dbias.py` + `moe_blockscaled_grouped_gemm_dglu_rubin.py`🤖 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 `@skills/cutedsl-kernel-integration/references/integration-pattern.md` around lines 97 - 101, Update the Examples list in the integration-pattern documentation so each Rubin module is paired with its corresponding default module, matching the convention established by the first bullet. Add the default counterparts for the GLU and DGLU entries while preserving the existing module names and ordering.python/cudnn/api_base.py (1)
36-43: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider memoizing per device index.
get_device_type()is invoked on every wrapper call (cache-key construction ingrouped_gemm_glu/api.pyLine 617) and each call re-queriestorch.cuda.get_device_capability(). Anlru_cachekeyed ontorch.cuda.current_device()keeps it correct on mixed-arch systems while removing the repeated query.♻️ Suggested memoization
+@functools.lru_cache(maxsize=None) +def _device_type_for(device_index: int) -> str: + return "rubin" if torch.cuda.get_device_capability(device_index) == (10, 7) else "blackwell" + + def is_sm107_device() -> bool: """Return True when the current CUDA device is Rubin (SM107).""" - return torch.cuda.is_available() and torch.cuda.get_device_capability(torch.cuda.current_device()) == (10, 7) + return get_device_type() == "rubin" def get_device_type() -> str: """Return the architecture family used by SM100 grouped GEMM wrappers.""" - return "rubin" if is_sm107_device() else "blackwell" + if not torch.cuda.is_available(): + return "blackwell" + return _device_type_for(torch.cuda.current_device())Note: the tests in
test/python/fe_api/test_rubin_kernel_dispatch.pypatchtorch.cuda.get_device_capabilityand would need cache clearing if this is adopted.🤖 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/api_base.py` around lines 36 - 43, Memoize device-type detection per CUDA device index to avoid repeated capability queries while preserving correctness on mixed-architecture systems. Update is_sm107_device or the nearest device-type helper to use an lru_cache keyed by torch.cuda.current_device(), and ensure tests that patch torch.cuda.get_device_capability clear the cache between cases.python/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.py (2)
1080-1089: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe Rubin/non-Rubin tuning-arg branch is now repeated four times.
Dense compile (Line 885), dense
tensor_api(Line 944), discrete compile (Line 1080), and discretetensor_api(Line 1150) all encode the same rule: appendgeglu_alpha/glu_clamp_max/glu_clamp_minonly for non-Rubin. A single helper returning the extra positional tuple would keep the four call sites from drifting apart.♻️ Sketch
def _glu_tune_args(is_rubin_kernel, geglu_alpha, glu_clamp_max, glu_clamp_min): if is_rubin_kernel: return () return ( cutlass.Float32(geglu_alpha), cutlass.Float32(glu_clamp_max), cutlass.Float32(glu_clamp_min), )- if self._is_rubin_kernel: - _compiled_kernel(*kernel_args) - else: - _compiled_kernel( - *kernel_args, - cutlass.Float32(geglu_alpha), - cutlass.Float32(glu_clamp_max), - cutlass.Float32(glu_clamp_min), - ) + _compiled_kernel( + *kernel_args, + *_glu_tune_args(self._is_rubin_kernel, geglu_alpha, glu_clamp_max, glu_clamp_min), + )Also applies to: 1150-1158
🤖 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_glu/_blockscaled_api.py` around lines 1080 - 1089, Extract the repeated Rubin/non-Rubin tuning-argument logic into a shared helper, such as _glu_tune_args, returning an empty tuple for Rubin kernels and the three Float32-wrapped values otherwise. Replace the branching at dense compile, dense tensor_api, discrete compile, and discrete tensor_api call sites with this helper, preserving each site’s existing argument order and values.
63-85: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHoist the shared Rubin GLU guard
python/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.pyandpython/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.pyboth define the same Rubin tuning defaults and rejection helper; moving them intopython/cudnn/grouped_gemm/grouped_gemm_utils.pywould keep the contract in one place._GEGGLU_ALPHA_DEFAULTalso looks like a typo for_GEGLU_ALPHA_DEFAULT.🤖 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_glu/_blockscaled_api.py` around lines 63 - 85, Move the shared Rubin GLU tuning defaults and _reject_unsupported_rubin_glu_tune_params helper from both blockscaled API modules into grouped_gemm_utils.py, then import and reuse them in each caller. Rename _GEGGLU_ALPHA_DEFAULT to _GEGLU_ALPHA_DEFAULT consistently while preserving the existing default values and rejection behavior.python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py (2)
2828-2843: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRun Black over this file.
Lines 2832-2843 chain statements with semicolons (Ruff E702) and the
for _m_halfbody at Lines 2929-2941 uses 2-space indentation while the rest of the file uses 4; there is also trailing whitespace at Lines 193 and 1510. As per coding guidelines, "Format Python code and notebooks with Black using a line length of 160."The OpenGrep "credit card number" hits at Lines 1508 and 1588 are false positives on the
LOG2_Eliteral.Also applies to: 2929-2941
🤖 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_glu/moe_blockscaled_grouped_gemm_glu_rubin.py` around lines 2828 - 2843, Run Black on the entire Python file using a line length of 160, ensuring chained semicolon statements near the grouped-GEMM partition setup are split into formatted statements, the for _m_half body uses consistent four-space indentation, and trailing whitespace is removed. Do not alter the LOG2_E literals identified as false-positive OpenGrep matches.Sources: Coding guidelines, Linters/SAST tools
1505-1513: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueName the baked-in GeGLU constants and note the API coupling.
1.702,7.0, and-7.0are hard-coded here and mirrored by_GEGGLU_ALPHA_DEFAULT/_GLU_CLAMP_*_DEFAULTingrouped_gemm_glu/_blockscaled_api.py, which rejects any other value on Rubin. A module-level named constant plus a short comment pointing at that guard makes the coupling explicit for whoever changes either side. Note the unpacked path (Line 1578) also usessilu_f32_geglu_scaled, which presumably bakes the same alpha.Also applies to: 3075-3081
🤖 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_glu/moe_blockscaled_grouped_gemm_glu_rubin.py` around lines 1505 - 1513, Define module-level named constants for the GeGLU alpha and clamp bounds currently hard-coded as 1.702, 7.0, and -7.0, then use them in geglu_act and the corresponding logic around the second affected location. Add a brief comment referencing _GEGGLU_ALPHA_DEFAULT and _GLU_CLAMP_*_DEFAULT in _blockscaled_api.py to document the Rubin API coupling, and verify the unpacked silu_f32_geglu_scaled path remains aligned with the same alpha.test/python/fe_api/test_rubin_kernel_dispatch.py (1)
128-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThese two tests don't exercise the production dispatch path.
test_kernel_selection_uses_rubin_on_sm107re-implements the_get_rubin_kernel() if is_sm107_device() else defaultexpression inside the test body, and patches_get_rubin_kernelto return the object it already fetched — so the assertions hold regardless of what_blockscaled_api.pyLine 268 actually does.test_grouped_gemm_quant_has_rubin_compile_branchesgrepsapi.pyfor literal source strings, which breaks on any rename without indicating a real regression.Patching
cudnn.api_base.get_device_typeand asserting on a constructed API object's_kernel/_is_rubin_kernelwould cover the real contract:with mock.patch("cudnn.api_base.get_device_type", return_value="rubin"): api = GroupedGemmGluBlockScaledAPI(...) # minimal sample descriptors assert api._kernel is api_mod._get_rubin_kernel()🤖 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 128 - 167, Replace both tests with behavioral coverage of the production dispatch path: construct a minimal GroupedGemmGluBlockScaledAPI instance while patching cudnn.api_base.get_device_type to return "rubin", then assert its _kernel is api_mod._get_rubin_kernel() and its _is_rubin_kernel state is correct. Remove the source-text assertions and the test-side reimplementation of the device-selection conditional so the tests validate actual constructor behavior.python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py (2)
3285-3286: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReword the inline comment.
"Don't ask why, AST is shit tracking the constexpr values to loop args" is not appropriate for shipped source. Describe the actual constraint (the tiled copy must be rebuilt inside the loop because the constexpr shapes are not propagated to the loop-carried values) so the next reader knows whether it can be hoisted.
🤖 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_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py` around lines 3285 - 3286, Replace the informal inline comment above epilog_tmem_copy_and_partition with a professional explanation that the tiled copy must be rebuilt inside the loop because constexpr shapes are not propagated to loop-carried values; clarify that this prevents hoisting the operation.
220-220: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRun Black over this file.
Ruff reports several non-Black constructs in the new file:
E701at Line 220 (if ...: return Falseon one line),E702at Line 2086 (y1_0 = 0.0; y1_1 = 0.0; ...), plus trailing whitespace at Lines 2517 and 3472. Line 873 also tripsE712— preferif cutlass.const_expr(not self.generate_sfd):. As per coding guidelines, "Format Python code and notebooks with Black using a line length of 160."Also applies to: 873-873, 2086-2086, 2517-2517, 3472-3472
🤖 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_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py` at line 220, Run Black on moe_blockscaled_grouped_gemm_dglu_rubin.py with a 160-character line length, then manually resolve remaining Ruff issues: expand the one-line if/return at the validation logic, split the semicolon-separated assignments, replace the E712 comparison near generate_sfd with the recommended const_expr form, and remove trailing whitespace.Sources: Coding guidelines, Linters/SAST tools
python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py (1)
999-1024: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHoist the
cutlass.Float32(self.linear_offset)conversion out of the per-launch closure.The surrounding code already caches launch-invariant values (
cached_workspace_ptr,cached_n,cached_k,cached_b_stride);linear_offsetis equally invariant, so building the wrapper on everyexecute()is inconsistent and needlessly re-allocates. Same applies to the discrete path at Line 1229.♻️ Proposed refactor (dense path)
cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + cached_linear_offset = cutlass.Float32(self.linear_offset) if self._is_rubin_kernel else Noneif self._is_rubin_kernel: - _compiled_kernel(*kernel_args, cutlass.Float32(self.linear_offset), dbias_tensor, stream) + _compiled_kernel(*kernel_args, cached_linear_offset, dbias_tensor, stream) else: _compiled_kernel(*kernel_args, dbias_tensor, stream)🤖 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_dglu/_blockscaled_api.py` around lines 999 - 1024, Hoist the launch-invariant cutlass.Float32(self.linear_offset) conversion out of the per-launch closure in the dense path around _compiled_kernel and reuse the cached value when invoking the Rubin kernel. Apply the same change to the discrete path around its corresponding _compiled_kernel invocation near the second launch site, while preserving the existing non-Rubin argument handling.python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py (1)
591-594: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRecord this as a temporary workaround with a version/tracking reference.
Forcing
epilogue_prefetch_more = Falsedisables the ping-pong prefetch that previously ran for the fp8-in/fp8-out path, so this is a real throughput regression on that config, and it also makes theepilogue_prefetch_morebranches at Lines 2993-3001 and 3011-3014 permanently dead. Please note the affected CuTe DSL version (or a tracking issue) in the comment so the prefetch path can be re-enabled once the lowering bug is fixed. Want me to open an issue to track re-enabling it?🤖 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_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py` around lines 591 - 594, Document the assignment to self.epilogue_prefetch_more in the nearby comment as a temporary CuTe DSL lowering workaround, including the affected version or an existing tracking issue reference. Note that it disables fp8-in/fp8-out ping-pong prefetching and should be re-enabled when the lowering bug is fixed.
🤖 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_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py`:
- Around line 3630-3641: Define thr_copy_t2r in the store_d_directly branch
before its partition_D call, using tiled_copy_t2r.get_slice(epi_tidx), so the
branch traces without a NameError. Keep the existing gD_sub_loop and
cute.filter_zeros flow unchanged.
- Around line 2053-2072: Update Rubin’s dgeglu implementation to match the
reference backward: add and propagate square_alpha through dgeglu and its
caller, applying square_alpha to the vector accumulator and scalar fc2 gradient.
In the vector and scalar clamp-mask logic, make in-range x1_filter and x2_filter
values 1.0 rather than the clamped activations, preserving clamped values only
for the activation computation.
In
`@python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py`:
- Around line 811-813: Set self.generate_sfd from sfd_row_tensor and
norm_const_tensor before calling _setup_attributes(), so stage and sD_col sizing
use the final flag; remove the later recomputation. In the conditional that
disables self.discrete_col_sfd, replace the equality comparison with the direct
negation not self.generate_sfd.
- Around line 3167-3173: Replace the fixed subtile_idx == 6 gate in the SFD
row/column write block with a check for the loop’s final iteration, subtile_cnt
- 2. Preserve the existing bounds checks and scale-factor stores so supported
N=192 tiles execute the SFD writes.
In `@python/cudnn/grouped_gemm/grouped_gemm_quant/api.py`:
- Around line 1180-1185: Update the Rubin row-scale rejection in execute() to
raise NotImplementedError instead of ValueError, matching the behavior of
check_support() and the wrapper for the same unsupported condition. Preserve the
existing validation message and all other branches.
In
`@python/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.py`:
- Around line 783-786: Resolve generate_sfd and discrete_col_sfd before
_setup_attributes(), using not self.generate_sfd, so SMEM sizing and all
consumers use the settled tensor-based state; update the anchor site
accordingly. At the sibling tma_atom_d_col construction, gate the column-TMA
path on generate_sfd and fall back to d when disabled, avoiding dereferencing
d_col=None.
- Around line 736-752: Update the dtype initialization in the constructor around
self.a_dtype and self.b_dtype so b_dtype is derived from b.element_type,
including for DENSE mode, rather than from a.element_type. Keep the existing A/B
dtype equality guard unchanged so genuinely mixed inputs are rejected, and
ensure downstream b_dtype-based logic such as mma_inst_k receives the actual B
dtype.
- Around line 1238-1239: Update the column quantization path near acc_scale_col
to call fmin with nan=True, matching the quant_sfd_row behavior and ensuring NaN
accumulators produce consistent scale factors for both SFD outputs.
- Around line 567-578: Update the SMEM accounting around `amax_bytes` and
`num_ab_stage`: reserve the unconditional `SharedStorage.sAmax` allocation for
every output dtype, using its actual `num_epilog_warps`-by-Float32 footprint
rather than gating it on BF16, then validate the computed `num_ab_stage` before
constructing the pipeline so zero or negative capacity fails clearly instead of
producing a degenerate configuration.
- Around line 2440-2456: Update the SFD flush condition in the non-breuse path
to derive the trigger from real_subtile_idx, flushing when its four-scale-factor
group is complete rather than checking raw subtile_idx. Replace the hard-coded
3/7 checks with a condition that also flushes the final group for the actual
subtile count, including the mma_tiler[1] != 256 case, while preserving the
existing bounds checks and copy operations.
In `@test/python/fe_api/test_rubin_kernel_dispatch.py`:
- Around line 51-55: Add the pytest.mark.L0 decorator to each of the four new
fast, import-only tests, including test_rubin_kernel_module_is_present and the
tests at the referenced locations. Keep the existing parameterization and test
logic unchanged.
---
Nitpick comments:
In `@python/cudnn/api_base.py`:
- Around line 36-43: Memoize device-type detection per CUDA device index to
avoid repeated capability queries while preserving correctness on
mixed-architecture systems. Update is_sm107_device or the nearest device-type
helper to use an lru_cache keyed by torch.cuda.current_device(), and ensure
tests that patch torch.cuda.get_device_capability clear the cache between cases.
In `@python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py`:
- Around line 999-1024: Hoist the launch-invariant
cutlass.Float32(self.linear_offset) conversion out of the per-launch closure in
the dense path around _compiled_kernel and reuse the cached value when invoking
the Rubin kernel. Apply the same change to the discrete path around its
corresponding _compiled_kernel invocation near the second launch site, while
preserving the existing non-Rubin argument handling.
In
`@python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py`:
- Around line 591-594: Document the assignment to self.epilogue_prefetch_more in
the nearby comment as a temporary CuTe DSL lowering workaround, including the
affected version or an existing tracking issue reference. Note that it disables
fp8-in/fp8-out ping-pong prefetching and should be re-enabled when the lowering
bug is fixed.
In
`@python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py`:
- Around line 3285-3286: Replace the informal inline comment above
epilog_tmem_copy_and_partition with a professional explanation that the tiled
copy must be rebuilt inside the loop because constexpr shapes are not propagated
to loop-carried values; clarify that this prevents hoisting the operation.
- Line 220: Run Black on moe_blockscaled_grouped_gemm_dglu_rubin.py with a
160-character line length, then manually resolve remaining Ruff issues: expand
the one-line if/return at the validation logic, split the semicolon-separated
assignments, replace the E712 comparison near generate_sfd with the recommended
const_expr form, and remove trailing whitespace.
In `@python/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.py`:
- Around line 1080-1089: Extract the repeated Rubin/non-Rubin tuning-argument
logic into a shared helper, such as _glu_tune_args, returning an empty tuple for
Rubin kernels and the three Float32-wrapped values otherwise. Replace the
branching at dense compile, dense tensor_api, discrete compile, and discrete
tensor_api call sites with this helper, preserving each site’s existing argument
order and values.
- Around line 63-85: Move the shared Rubin GLU tuning defaults and
_reject_unsupported_rubin_glu_tune_params helper from both blockscaled API
modules into grouped_gemm_utils.py, then import and reuse them in each caller.
Rename _GEGGLU_ALPHA_DEFAULT to _GEGLU_ALPHA_DEFAULT consistently while
preserving the existing default values and rejection behavior.
In
`@python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py`:
- Around line 2828-2843: Run Black on the entire Python file using a line length
of 160, ensuring chained semicolon statements near the grouped-GEMM partition
setup are split into formatted statements, the for _m_half body uses consistent
four-space indentation, and trailing whitespace is removed. Do not alter the
LOG2_E literals identified as false-positive OpenGrep matches.
- Around line 1505-1513: Define module-level named constants for the GeGLU alpha
and clamp bounds currently hard-coded as 1.702, 7.0, and -7.0, then use them in
geglu_act and the corresponding logic around the second affected location. Add a
brief comment referencing _GEGGLU_ALPHA_DEFAULT and _GLU_CLAMP_*_DEFAULT in
_blockscaled_api.py to document the Rubin API coupling, and verify the unpacked
silu_f32_geglu_scaled path remains aligned with the same alpha.
In `@skills/cutedsl-kernel-integration/references/integration-pattern.md`:
- Around line 97-101: Update the Examples list in the integration-pattern
documentation so each Rubin module is paired with its corresponding default
module, matching the convention established by the first bullet. Add the default
counterparts for the GLU and DGLU entries while preserving the existing module
names and ordering.
In `@skills/cutedsl-kernel-integration/SKILL.md`:
- Line 40: Update the architecture-dispatch testing instruction in SKILL.md to
run pytest from the test/python working directory, ensuring pytest.ini and
conftest.py are applied and their import-order and environment-variable
requirements are preserved.
In `@test/python/fe_api/test_rubin_kernel_dispatch.py`:
- Around line 128-167: Replace both tests with behavioral coverage of the
production dispatch path: construct a minimal GroupedGemmGluBlockScaledAPI
instance while patching cudnn.api_base.get_device_type to return "rubin", then
assert its _kernel is api_mod._get_rubin_kernel() and its _is_rubin_kernel state
is correct. Remove the source-text assertions and the test-side reimplementation
of the device-selection conditional so the tests validate actual constructor
behavior.
🪄 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: 6048a52f-8594-4e69-9db6-ba396c4e775d
📒 Files selected for processing (16)
python/cudnn/api_base.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/api.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.pypython/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.pypython/cudnn/grouped_gemm/grouped_gemm_glu/api.pypython/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.pypython/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.pypython/cudnn/grouped_gemm/grouped_gemm_quant/api.pypython/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.pypython/cudnn/grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.pypython/cudnn/grouped_gemm/grouped_gemm_utils.pyskills/cutedsl-kernel-integration/SKILL.mdskills/cutedsl-kernel-integration/references/integration-pattern.mdtest/python/fe_api/test_rubin_kernel_dispatch.py
| @cute.jit | ||
| def dgeglu(self, | ||
| acc_vec: cute.Tensor, | ||
| x1_vec_load: cute.Tensor, | ||
| x2_vec_load: cute.Tensor, | ||
| mProb: cute.Tensor, | ||
| linear_offset: Float32, | ||
| dprob_swiglu: Optional[cute.Tensor] = None | ||
| ): | ||
| LOG2_E = cutlass.Float32(1.4426950408889634) | ||
| x_dtype = x1_vec_load.element_type | ||
| geglu_max_value = x_dtype(7.0) | ||
| geglu_min_value = x_dtype(-7.0) | ||
| zero_x_dtype = x_dtype(0.0) | ||
| fmul2 = partial(cute.arch.mul_packed_f32x2, rnd='rn', ftz=False) | ||
| fadd2 = partial(cute.arch.add_packed_f32x2, rnd='rn', ftz=False) | ||
| scale_1702 = (1.702, 1.702) | ||
| ones2 = (1.0, 1.0) | ||
| mprob2 = (mProb, mProb) | ||
| linear_offset2 = (linear_offset, linear_offset) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
Rubin dgeglu drops square_alpha and multiplies the gradients by the clamped activations instead of a 0/1 mask.
Two divergences from BlockScaledMoEGroupedGemmDgluDbiasKernel.dgeglu in moe_blockscaled_grouped_gemm_dglu_dbias.py, both of which change the numerical result rather than just the schedule:
-
Missing
alpha²scaling. The SM100 kernel starts withacc = fmul2(acc, (square_alpha, square_alpha))and the scalar path usesfc2_dgrad = acc_vec[i] * square_alpha. The Rubin version takes nosquare_alphaargument at all (Line 2054) and the caller at Line 3471 does not pass one, so the per-expert alpha scaling is silently dropped fordgeglu.dswigluon the same kernel does apply it (Line 3469), so any expert withalpha != 1produces wrong dGeGLU gradients. -
Clamp mask uses the value, not 1.0. Lines 2150-2151 / 2159-2162 set
x1_filter = y1_0(andx2_filter = y2_0) inside the in-range case, then multiplydy1/dy2by it. The SM100 kernel uses1.0there — the filter is meant to be a pass-through mask, so multiplying by the clamped activation scales the gradient byyinstead of gating it. Same issue in the scalar path at Lines 2200-2202.
Please confirm against the reference dGeGLU backward before merge.
🐛 Sketch of the expected form (vectorized path)
def dgeglu(self,
acc_vec: cute.Tensor,
x1_vec_load: cute.Tensor,
x2_vec_load: cute.Tensor,
mProb: cute.Tensor,
+ square_alpha: Float32,
linear_offset: Float32,
dprob_swiglu: Optional[cute.Tensor] = None
): acc = (acc_vec[i], acc_vec[i + 1])
+ acc = fmul2(acc, (square_alpha, square_alpha))- x1_filter_0 = y1_0 if x1_0 <= geglu_max_value else cutlass.Float32(0.0)
- x1_filter_1 = y1_1 if x1_1 <= geglu_max_value else cutlass.Float32(0.0)
+ x1_filter_0 = cutlass.Float32(1.0) if x1_0 <= geglu_max_value else cutlass.Float32(0.0)
+ x1_filter_1 = cutlass.Float32(1.0) if x1_1 <= geglu_max_value else cutlass.Float32(0.0)- x2_filter_0 = x2_0 if x2_0 <= geglu_max_value else x_dtype(0.0)
- x2_filter_1 = x2_1 if x2_1 <= geglu_max_value else x_dtype(0.0)
- x2_filter_0 = y2_0 if x2_filter_0 >= geglu_min_value else cutlass.Float32(0.0)
- x2_filter_1 = y2_1 if x2_filter_1 >= geglu_min_value else cutlass.Float32(0.0)
+ x2_filter_0 = cutlass.Float32(1.0) if (x2_0 >= geglu_min_value and x2_0 <= geglu_max_value) else cutlass.Float32(0.0)
+ x2_filter_1 = cutlass.Float32(1.0) if (x2_1 >= geglu_min_value and x2_1 <= geglu_max_value) else cutlass.Float32(0.0)Also applies to: 2140-2163, 2186-2205
🧰 Tools
🪛 OpenGrep (1.26.0)
[ERROR] 2062-2062: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🤖 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_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py`
around lines 2053 - 2072, Update Rubin’s dgeglu implementation to match the
reference backward: add and propagate square_alpha through dgeglu and its
caller, applying square_alpha to the vector accumulator and scalar fc2 gradient.
In the vector and scalar clamp-mask logic, make in-range x1_filter and x2_filter
values 1.0 rather than the clamped activations, preserving clamped values only
for the activation computation.
| if cutlass.const_expr(self.store_d_directly): | ||
| self.epilog_sync_barrier.arrive_and_wait() | ||
| d_idx_mn = (mma_tile_coord_mnl[0], mma_tile_coord_mnl[1]) | ||
| d_epilogue_subtile = ( | ||
| cute.make_layout(128), | ||
| cute.make_layout(self.mma_tiler[1] * 2), | ||
| ) | ||
| gD_sub_loop = cute.local_tile( | ||
| real_d, d_epilogue_subtile, (None, None, None) | ||
| ) | ||
| tCgD_mnl_loop = thr_copy_t2r.partition_D(gD_sub_loop) | ||
| tCgD_mnl_loop = cute.filter_zeros(tCgD_mnl_loop) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
thr_copy_t2r is undefined in the store_d_directly branch.
Ruff flags F821 at Line 3640. The branch is currently unreachable because self.store_d_directly is hard-coded to False (Line 530), but it will NameError at trace time the moment anyone flips that flag. Either bind thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) in this scope or drop the dead branch.
🧰 Tools
🪛 Ruff (0.16.0)
[error] 3640-3640: Undefined name thr_copy_t2r
(F821)
🤖 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_dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py`
around lines 3630 - 3641, Define thr_copy_t2r in the store_d_directly branch
before its partition_D call, using tiled_copy_t2r.get_slice(epi_tidx), so the
branch traces without a NameError. Keep the existing gD_sub_loop and
cute.filter_zeros flow unchanged.
Source: Linters/SAST tools
| self.generate_sfd = sfd_row_tensor is not None and norm_const_tensor is not None | ||
| if cutlass.const_expr(self.generate_sfd == False): | ||
| self.discrete_col_sfd = False |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Compare the ordering with the SM100 sibling kernel
fd -t f 'moe_blockscaled_grouped_gemm_glu_bias.py' python/cudnn --exec rg -nP -C4 'generate_sfd\s*=|_setup_attributes\(\)' {}Repository: NVIDIA/cudnn-frontend
Length of output: 1614
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py'
echo '--- around _setup_attributes call and tensor setup ---'
sed -n '760,860p' "$file"
echo
echo '--- _setup_attributes implementation ---'
rg -n 'def _setup_attributes|compute_stages|sD_col_size|generate_sfd|discrete_col_sfd' "$file" -n -C3Repository: NVIDIA/cudnn-frontend
Length of output: 12128
Move self.generate_sfd before _setup_attributes(). _setup_attributes() uses the constructor value to size stages, but the flag is recomputed later from the tensors, so a mismatch can size SMEM and sD_col from different assumptions. Also replace self.generate_sfd == False with not self.generate_sfd.
🧰 Tools
🪛 Ruff (0.16.0)
[error] 812-812: Avoid equality comparisons to False; use not self.generate_sfd: for false checks
Replace with not self.generate_sfd
(E712)
🤖 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_glu/moe_blockscaled_grouped_gemm_glu_rubin.py`
around lines 811 - 813, Set self.generate_sfd from sfd_row_tensor and
norm_const_tensor before calling _setup_attributes(), so stage and sD_col sizing
use the final flag; remove the later recomputation. In the conditional that
disables self.discrete_col_sfd, replace the equality comparison with the direct
negation not self.generate_sfd.
Source: Linters/SAST tools
| if subtile_idx == 6: | ||
| if sfd_row_idx_mn[1] * 32 * regPerSubtile < cute.size(cute.shape(mSFDRow_mnl.layout, mode=[1])): | ||
| tCrSFDRow.store(tCrSFDRow_pvscale.load().to(self.sf_dtype)) | ||
| cute.autovec_copy(tCrSFDRow, tCgSFDRow) | ||
| if sfd_col_idx_mn[1] * 32 * regPerSubtile < cute.size(cute.shape(mSFDCol_mnl.layout, mode=[1])): | ||
| tCrSFDCol.store(tCrSFDCol_pvscale.load().to(self.sf_dtype)) | ||
| cute.autovec_copy(tCrSFDCol, tCgSFDCol) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how the SM100 sibling gates the SFD store, and what subtile counts are reachable
fd -t f 'moe_blockscaled_grouped_gemm_glu_bias.py' python/cudnn --exec rg -nP -C6 'subtile_idx ==|subtile_cnt' {}Repository: NVIDIA/cudnn-frontend
Length of output: 3530
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the epilogue loop, subtile count derivation, and any matching sibling implementation.
sed -n '2920,3195p' python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py
printf '\n--- sibling search ---\n'
rg -n -C4 'subtile_idx == 6|subtile_cnt =|iter_acc_early_release_in_epilogue|can_implement|mma_tiler_mn|epi_tile' python/cudnn/grouped_gemm -g '*.py'Repository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect only the relevant regions in the target file and the sibling for exact loop shape and validation.
echo '--- target: gate region ---'
sed -n '3135,3195p' python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py
echo
echo '--- target: setup/validation region ---'
rg -n -C3 'epi_tile|subtile_cnt|subtile_idx == 6|can_implement|mma_tiler_mn' python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py
echo
echo '--- sibling: matching gate region ---'
sed -n '1758,1805p' python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_grouped_gemm_dglu_dbias.pyRepository: NVIDIA/cudnn-frontend
Length of output: 23671
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '2928,3176p' python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.py
printf '\n--- can_implement region ---\n'
sed -n '143,205p' python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.pyRepository: NVIDIA/cudnn-frontend
Length of output: 15805
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect how the D CTA tile and epilogue tile determine the subtile count.
rg -n -C4 'cta_tile_shape_mnk_d|cta_tile_shape_mnk|epi_tile_n_required|num_d_stage|enable_breuse|subtile_cnt = cute.size\(tTR_tAcc_h.shape, mode=\[3\]\)' \
python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.pyRepository: NVIDIA/cudnn-frontend
Length of output: 17941
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C5 'mma_tiler_d\s*=|cta_tile_shape_mnk_d\s*=|epi_tile_cnt\s*=|epi_tile_n_required\s*=' \
python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_rubin.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1853
Hard-coded subtile_idx == 6 skips SFD writes for supported N=192 tiles.
This loop advances by 2, so the last iteration is subtile_cnt - 2, not a fixed 6. can_implement allows mma_tiler_mn[1] == 192, where this gate never fires and the row/col scale factors are never written.
🤖 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_glu/moe_blockscaled_grouped_gemm_glu_rubin.py`
around lines 3167 - 3173, Replace the fixed subtile_idx == 6 gate in the SFD
row/column write block with a check for the loop’s final iteration, subtile_cnt
- 2. Preserve the existing bounds checks and scale-factor stores so supported
N=192 tiles execute the SFD writes.
| if self._is_rubin_kernel: | ||
| self._value_error_if( | ||
| row_scale_tensor is not None, | ||
| "row_scale_tensor is not supported on Rubin (sm107)", | ||
| ) | ||
| elif self.row_scale_desc is None: |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Use NotImplementedError for the Rubin row-scale rejection to match the other two gates.
check_support() (Line 302) and the wrapper (Line 1461) both raise NotImplementedError for this exact condition, but execute() raises ValueError, so callers cannot catch it uniformly.
🐛 Proposed fix
if self._is_rubin_kernel:
- self._value_error_if(
+ self._not_implemented_error_if(
row_scale_tensor is not None,
"row_scale_tensor is not supported on Rubin (sm107)",
)📝 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.
| if self._is_rubin_kernel: | |
| self._value_error_if( | |
| row_scale_tensor is not None, | |
| "row_scale_tensor is not supported on Rubin (sm107)", | |
| ) | |
| elif self.row_scale_desc is None: | |
| if self._is_rubin_kernel: | |
| self._not_implemented_error_if( | |
| row_scale_tensor is not None, | |
| "row_scale_tensor is not supported on Rubin (sm107)", | |
| ) | |
| elif self.row_scale_desc is None: |
🤖 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 1180 -
1185, Update the Rubin row-scale rejection in execute() to raise
NotImplementedError instead of ValueError, matching the behavior of
check_support() and the wrapper for the same unsupported condition. Preserve the
existing validation message and all other branches.
| self.a_dtype: Type[cutlass.Numeric] = a.element_type | ||
| self.b_dtype: Type[cutlass.Numeric] = a.element_type | ||
| self.c_dtype: Type[cutlass.Numeric] = c.element_type | ||
| self.d_dtype: Type[cutlass.Numeric] = d.element_type | ||
| self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type | ||
| self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() | ||
| self.c_layout = utils.LayoutEnum.from_tensor(c) | ||
| self.d_layout = utils.LayoutEnum.from_tensor(d) | ||
| self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 | ||
|
|
||
| if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): | ||
| self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() | ||
| else: | ||
| self.b_major_mode = b_major_mode | ||
|
|
||
| if cutlass.const_expr(self.a_dtype != self.b_dtype): | ||
| raise TypeError(f"A/B dtype must match: {self.a_dtype} != {self.b_dtype}") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
b_dtype is derived from a, making the A/B dtype guard dead code.
Line 737 assigns a.element_type to self.b_dtype, so the check at Lines 751-752 can never fire. For DENSE mode b is a real tensor whose element type should be read directly; the mma_inst_k selection (Line 334) explicitly tests self.b_dtype.width independently of A, so a genuinely mixed A/B configuration would be silently mis-typed here rather than rejected.
🐛 Proposed fix
self.a_dtype: Type[cutlass.Numeric] = a.element_type
- self.b_dtype: Type[cutlass.Numeric] = a.element_type
+ if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE):
+ self.b_dtype: Type[cutlass.Numeric] = b.element_type
+ else:
+ # Discrete mode: `b` is a raw pointer array, so B shares A's element type.
+ self.b_dtype: Type[cutlass.Numeric] = a.element_type📝 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.
| self.a_dtype: Type[cutlass.Numeric] = a.element_type | |
| self.b_dtype: Type[cutlass.Numeric] = a.element_type | |
| self.c_dtype: Type[cutlass.Numeric] = c.element_type | |
| self.d_dtype: Type[cutlass.Numeric] = d.element_type | |
| self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type | |
| self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() | |
| self.c_layout = utils.LayoutEnum.from_tensor(c) | |
| self.d_layout = utils.LayoutEnum.from_tensor(d) | |
| self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 | |
| if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): | |
| self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() | |
| else: | |
| self.b_major_mode = b_major_mode | |
| if cutlass.const_expr(self.a_dtype != self.b_dtype): | |
| raise TypeError(f"A/B dtype must match: {self.a_dtype} != {self.b_dtype}") | |
| self.a_dtype: Type[cutlass.Numeric] = a.element_type | |
| if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): | |
| self.b_dtype: Type[cutlass.Numeric] = b.element_type | |
| else: | |
| # Discrete mode: `b` is a raw pointer array, so B shares A's element type. | |
| self.b_dtype: Type[cutlass.Numeric] = a.element_type | |
| self.c_dtype: Type[cutlass.Numeric] = c.element_type | |
| self.d_dtype: Type[cutlass.Numeric] = d.element_type | |
| self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type | |
| self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() | |
| self.c_layout = utils.LayoutEnum.from_tensor(c) | |
| self.d_layout = utils.LayoutEnum.from_tensor(d) | |
| self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 | |
| if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): | |
| self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() | |
| else: | |
| self.b_major_mode = b_major_mode | |
| if cutlass.const_expr(self.a_dtype != self.b_dtype): | |
| raise TypeError(f"A/B dtype must match: {self.a_dtype} != {self.b_dtype}") |
🤖 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/moe_blockscaled_grouped_gemm_quant_rubin.py`
around lines 736 - 752, Update the dtype initialization in the constructor
around self.a_dtype and self.b_dtype so b_dtype is derived from b.element_type,
including for DENSE mode, rather than from a.element_type. Keep the existing A/B
dtype equality guard unchanged so genuinely mixed inputs are rejected, and
ensure downstream b_dtype-based logic such as mma_inst_k receives the actual B
dtype.
| # ---- SFD setup ---- | ||
| self.generate_sfd = sfd_row_tensor is not None and norm_const_tensor is not None | ||
| if cutlass.const_expr(self.generate_sfd == False): | ||
| self.discrete_col_sfd = False |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Column-SFD state is resolved too late in __call__. self.generate_sfd is only settled from tensor presence at Line 784, after _setup_attributes() has already sized the SMEM stages with the constructor value, and the d_col TMA atom is built with no gate at all. Hoist the SFD resolution above _setup_attributes() (or validate the constructor flag against the tensors) so every consumer sees one settled value.
python/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.py#L783-L786: move thegenerate_sfd/discrete_col_sfdresolution before the_setup_attributes()call at Line 754, and replace== Falsewithnot ....python/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.py#L937-L942: gate thetma_atom_d_colconstruction on the settledgenerate_sfd, falling back todwhen the column path is disabled sod_col=Nonecannot be dereferenced.
🧰 Tools
🪛 Ruff (0.16.0)
[error] 785-785: Avoid equality comparisons to False; use not self.generate_sfd: for false checks
Replace with not self.generate_sfd
(E712)
📍 Affects 1 file
python/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.py#L783-L786(this comment)python/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.py#L937-L942
🤖 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/moe_blockscaled_grouped_gemm_quant_rubin.py`
around lines 783 - 786, Resolve generate_sfd and discrete_col_sfd before
_setup_attributes(), using not self.generate_sfd, so SMEM sizing and all
consumers use the settled tensor-based state; update the anchor site
accordingly. At the sibling tma_atom_d_col construction, gate the column-TMA
path on generate_sfd and fall back to d when disabled, avoiding dereferencing
d_col=None.
| fp32_max = cutlass.Float32(3.40282346638528859812e38) | ||
| acc_scale_col = fmin(acc_scale_col, fp32_max) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
NaN handling diverges between the row and column quantization paths.
quant_sfd_row clamps with fmin(acc_scale, fp32_max, nan=True) (Line 1189) while the column path omits nan, so a NaN accumulator yields different scale factors for the two SFD outputs from identical input data. Make the flag consistent unless the asymmetry is deliberate.
♻️ Proposed fix
- acc_scale_col = fmin(acc_scale_col, fp32_max)
+ acc_scale_col = fmin(acc_scale_col, fp32_max, nan=True)📝 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.
| fp32_max = cutlass.Float32(3.40282346638528859812e38) | |
| acc_scale_col = fmin(acc_scale_col, fp32_max) | |
| fp32_max = cutlass.Float32(3.40282346638528859812e38) | |
| acc_scale_col = fmin(acc_scale_col, fp32_max, nan=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
`@python/cudnn/grouped_gemm/grouped_gemm_quant/moe_blockscaled_grouped_gemm_quant_rubin.py`
around lines 1238 - 1239, Update the column quantization path near acc_scale_col
to call fmin with nan=True, matching the quant_sfd_row behavior and ensuring NaN
accumulators produce consistent scale factors for both SFD outputs.
| if cutlass.const_expr(self.mma_tiler[1] == 256): | ||
| sfd_n = epi_work_tile_info.tile_n_idx * 2 + (real_subtile_idx >> 2) | ||
| else: | ||
| sfd_n = epi_work_tile_info.tile_n_idx | ||
| sfd_row_idx_mn = (global_sfd_m, sfd_n) | ||
| sfd_col_idx_mn = sfd_row_idx_mn | ||
| if cutlass.const_expr(self.discrete_col_sfd): | ||
| sfd_col_idx_mn = (epi_work_tile_info.tile_m_idx, sfd_n) | ||
| tCgSFDRow = tCgSFDRow_mn[(None, None, None, *sfd_row_idx_mn)] | ||
| tCgSFDCol = tCgSFDCol_mn[(None, None, None, *sfd_col_idx_mn)] | ||
| if subtile_idx == 3 or subtile_idx == 7: | ||
| if sfd_row_idx_mn[1] * 32 * regPerSubtile < cute.size(cute.shape(mSFDRow_mnl.layout, mode=[1])): | ||
| tCrSFDRow.store(tCrSFDRow_pvscale.load().to(self.sf_dtype)) | ||
| cute.autovec_copy(tCrSFDRow, tCgSFDRow) | ||
| if sfd_col_idx_mn[1] * 32 * regPerSubtile < cute.size(cute.shape(mSFDCol_mnl.layout, mode=[1])): | ||
| tCrSFDCol.store(tCrSFDCol_pvscale.load().to(self.sf_dtype)) | ||
| cute.autovec_copy(tCrSFDCol, tCgSFDCol) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
SFD flush trigger uses subtile_idx while the pvscale slot and destination use real_subtile_idx.
Two problems on this non-breuse path:
- The pvscale group is filled at slot
real_subtile_idx % 4(Lines 2421, 2430) and addressed viareal_subtile_idx >> 2(Line 2441), but the flush fires on rawsubtile_idx == 3 or 7. Underoverlapping_accumwithreverse_subtile(Lines 2237-2239)real_subtile_idx = 7 - subtile_idx, so the flush happens whenreal_subtile_idx % 4 == 0— after only the first of the four scale factors has been written, storing three stale lanes. The breuse path gets this right by deriving the trigger fromreal_subtile_idx(Line 2397,sfd_write = n_sub % 4 == 3). - The literals
3/7hard-code an 8-subtile tile (N=256 withepi_tile[1] == 32). Line 2443 explicitly handlesmma_tiler[1] != 256, where the subtile count is 6 and the final group is never flushed.
🐛 Proposed fix
- if subtile_idx == 3 or subtile_idx == 7:
+ if real_subtile_idx % 4 == 3:🤖 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/moe_blockscaled_grouped_gemm_quant_rubin.py`
around lines 2440 - 2456, Update the SFD flush condition in the non-breuse path
to derive the trigger from real_subtile_idx, flushing when its four-scale-factor
group is complete rather than checking raw subtile_idx. Replace the hard-coded
3/7 checks with a condition that also flushes the final group for the actual
subtile count, including the mma_tiler[1] != 256 case, while preserving the
existing bounds checks and copy operations.
| @pytest.mark.parametrize( | ||
| "api_module_path,default_module_path,default_kernel_name,rubin_filename", | ||
| RUBIN_DISPATCH_CASES, | ||
| ) | ||
| def test_rubin_kernel_module_is_present( |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add a level marker to each new test.
Only test_grouped_gemm_quant_has_rubin_compile_branches carries @pytest.mark.L0. As per coding guidelines, "Mark every new Python test with a level from L0 through L4; keep L0 tests fast and place large parameter sweeps at higher levels." These four are fast, import-only checks, so L0 fits.
💚 Example
+@pytest.mark.L0
`@pytest.mark.parametrize`(
"api_module_path,default_module_path,default_kernel_name,rubin_filename",
RUBIN_DISPATCH_CASES,
)
def test_rubin_kernel_module_is_present(Also applies to: 66-70, 103-107, 124-128
🤖 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 51 - 55, Add
the pytest.mark.L0 decorator to each of the four new fast, import-only tests,
including test_rubin_kernel_module_is_present and the tests at the referenced
locations. Keep the existing parameterization and test logic unchanged.
Source: Coding guidelines
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-447-8a0e06e |
5 participants
What's here
moe_blockscaled_grouped_gemm_{glu,dglu,quant}_rubin.py_get_rubin_kernel()selection in the_blockscaled_api.pyfacades for glu/dglu,device_typeprepended to the block-scaled cache keys so Rubin and non-Rubin compiles don't collidegeglu_alpha, the GLU clamps, oruse_single_group_runtime_offsets. These are passed conditionally, and explicitly requesting an unsupported one raisesNotImplementedError(which the test harness turns into a skip with a reason).test/python/fe_api/test_rubin_kernel_dispatch.pyValidation
Validated on Rubin (sm107) silicon, A/B against a
developbaseline, overtest_grouped_gemm_{glu,dglu,quant}.py(317 cases, identical subset both sides):developbaselineNo regressions. The delta is exactly 6 tests, all moving
passed → skippedwith an explicit reason, none failing:Rubin grouped GEMM quant does not support row_scale fusionThe Rubin grouped GEMM kernels do not support use_single_group_runtime_offsetsSummary by CodeRabbit
New Features
Bug Fixes