Reorganize Gemm fusion - #459
Conversation
All GEMM fusion packages were scattered at the top of `python/cudnn/` with
redundant, inconsistent prefixes (`gemm_swiglu/`, `grouped_gemm/grouped_gemm_glu/`,
`discrete_grouped_gemm/discrete_grouped_gemm_swiglu/`). This collects them into a
single family tree grouped by operand layout, matching the internal layout:
python/cudnn/gemm/
├── cutedsl/
│ ├── dense/{amax,dsrelu,proj_rope_mxfp8,srelu,swiglu}/
│ ├── grouped/{dglu,dsrelu,dswiglu,glu,glu_hadamard,quant,srelu,
│ │ swiglu,unfused,wgrad}/
│ └── discrete_grouped/{dswiglu,swiglu}/
├── ops/ # backend-independent torch custom-op contracts
└── reference/ # pure-PyTorch MATMUL/POINTWISE correctness engine
The per-fusion directories drop the now-redundant prefix, so
`grouped_gemm/grouped_gemm_glu/` becomes `gemm/cutedsl/grouped/glu/`.
No public API changes. Every symbol keeps its name and stays reachable as
`cudnn.<symbol>`; `_LAZY_OPTIONAL_IMPORTS` in `python/cudnn/__init__.py` is
repointed at the new module paths, so the supported top-level entry point is
unchanged. `cudnn.grouped_gemm` and `cudnn.discrete_grouped_gemm` still resolve
as attributes via the same table.
Two adjacent modules move with the family so the tree is complete:
- `engines/reference_matmul_engine.py` -> `gemm/reference/`, re-exported from
`cudnn.engines` so `cudnn.engines.ReferenceMatmulEngine` keeps working.
- `experimental/ops/moe_grouped_matmul.py` -> `gemm/ops/`, aliased into
`sys.modules` from `cudnn.experimental.ops` so both the attribute and the
`import cudnn.experimental.ops.moe_grouped_matmul` form keep resolving.
GitHub-only APIs added since the internal reorg are carried over in place:
the unfused BF16 grouped GEMM (`GroupedGemmSm100`), the split BF16/MXFP8-input
proj-RoPE kernels, the `_bf16_api`/`_blockscaled_api` split, and the Rubin
(sm107) kernel variants.
`grouped_gemm/grouped_gemm_utils.py` is renamed to
`gemm/cutedsl/grouped/backend_utils.py` (backend enum + stream context), keeping
the kernel-side `utils.py` distinct.
Also fixes a latent import in the moved reference engine: its relative imports
(`.base`, `.engine_ids`, `..graph_types`) are now absolute `cudnn.engines.*` /
`cudnn.graph_types`, and the TYPE_CHECKING-only `..pygraph` now points at the
real `cudnn._pygraph`.
Docs, README, ACKNOWLEDGEMENTS, THIRD_PARTY_LICENSES and the
cutedsl-kernel-integration skill are updated to the new paths.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe pull request consolidates GEMM APIs, kernels, utilities, exports, documentation, and tests under ChangesGEMM CuTeDSL namespace migration
Estimated code review effort: 5 (Critical) | ~120 minutes 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 |
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-459-c84a177 |
test_grouped_gemm_bf16_wrapper reached into the kernel package for grouped_gemm_wrapper_sm100, but that symbol is exported at the top level. Import it from `cudnn` so the test is decoupled from the module layout. The remaining deep imports in the GEMM tests cannot go through the public surface: they target module-private wrapper caches (`_cache_of_*Objects`), monkeypatch targets that must be the module object, private API classes, and test-only reference oracles - none of which are exported by design. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-459-b6fd4b0 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 13
🧹 Nitpick comments (32)
python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py (2)
1346-1348: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead call:
m_idx/n_idxare computed and never used.
create_and_partition_new_SFDColderives its offsets fromcompute_expert_token_rangeon line 1348, so thetile_info_to_mn_idxcall is pure overhead (and Ruff RUF059 flags both unpacked names).♻️ Proposed fix
- m_idx, n_idx = self.tile_info_to_mn_idx(tile_info) expert_idx = tile_info[0] cumsum_tokens, tokens_this_group = compute_expert_token_range(padded_offsets, expert_idx)🤖 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/discrete_grouped/swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py` around lines 1346 - 1348, Remove the unused tile_info_to_mn_idx call and m_idx/n_idx unpacking from create_and_partition_new_SFDCol, retaining expert_idx and the compute_expert_token_range flow unchanged.Source: Linters/SAST tools
2430-2433: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueRemove the unused
d_col_pipelineconstruction. Thed_colTMA store is covered byd_pipeline.producer_tail(); no separate tail is required, so this is redundant rather than a store-drain race.🤖 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/discrete_grouped/swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py` around lines 2430 - 2433, Remove the unused d_col_pipeline construction by deleting the pipeline.PipelineTmaStore.create call near the d_pipeline setup. Keep d_pipeline.producer_tail() as the sole mechanism covering the d_col TMA store, and leave the surrounding pipeline configuration unchanged.python/cudnn/AGENTS.md (1)
22-31: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUntagged fenced code blocks in the new
gemm/layout trees. Both docs gained an identical package-layout tree fenced without a language, which markdownlint flags as MD040.
python/cudnn/AGENTS.md#L22-L31: change the opening fence on line 22 to```text.python/cudnn/README.md#L32-L40: change the opening fence on line 32 to```text.🤖 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/AGENTS.md` around lines 22 - 31, Change the opening fence for the package-layout tree to a text-labeled fence in both python/cudnn/AGENTS.md lines 22-31 and python/cudnn/README.md lines 32-40; leave the tree contents unchanged.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py (1)
664-668: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove
import os/import loggingto the module header.
osis referenced at line 156 (__init__) and line 372 (compile), but imported here — below the class. It works today only because module-level statements all execute before any method call; any future split of this file, or a top-level use ofosadded above, breaks immediately. Same forlogging/_logger.♻️ Proposed fix
Add to the header near line 13:
+import logging +import os from cuda.bindings import driver as cuda import torchThen trim the block below the class:
-import logging -import os - _logger = logging.getLogger(__name__) _cache_of_GroupedGemmDswigluSm100Objects = {}🤖 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/dswiglu/api.py` around lines 664 - 668, Move the os and logging imports from the block below the class to the module header, alongside the existing top-level imports. Keep the _logger initialization at module scope near the header, and remove the duplicated lower block without changing the cache definition or other class behavior.python/cudnn/__init__.py (1)
332-361: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winTarget grouped kernels by fusion package in
_LAZY_OPTIONAL_IMPORTS.The family initializers re-export all registered symbols, but importing one grouped or discrete-grouped symbol initializes every API and its kernel module. Use per-fusion paths while retaining family paths for the module aliases.
🤖 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/__init__.py` around lines 332 - 361, Update the symbol entries in _LAZY_OPTIONAL_IMPORTS so each grouped and discrete-grouped kernel class or wrapper resolves through its specific fusion package instead of the shared family module path. Keep the grouped_gemm and discrete_grouped_gemm module-alias entries on their existing family paths, and preserve each symbol’s exported attribute name.Source: Coding guidelines
python/cudnn/gemm/cutedsl/grouped/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py (1)
72-72: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicate
amax_reduction_per_thread: imported helper is shadowed by the method.Line 72 imports
amax_reduction_per_threadfrom..moe_kernel_helpers, but the method defined at 910-918 is the one actually called (self.amax_reduction_per_thread). Drop one of the two to avoid divergence.Also applies to: 910-918
🤖 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/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py` at line 72, Remove the unused imported amax_reduction_per_thread from the import list, keeping the class method amax_reduction_per_thread used through self.amax_reduction_per_thread as the single implementation.python/cudnn/gemm/cutedsl/dense/swiglu/dense_gemm_persistent_swiglu.py (1)
1439-1444: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePromised shared-memory assertion is missing; the three totals are dead code.
The comment says the total usage is checked, but
total_ab_smem/total_output_smem/total_smem_usedare computed and discarded. Alsonum_ab_stage(line 1428) can go<= 0for large tiles/dtypes and would fail later inside layout construction with an opaque error. Since this is a relocation, feel free to defer, but an explicit check here is cheap.♻️ Suggested guard
total_ab_smem = occupancy * ab_bytes_per_stage * num_ab_stage total_output_smem = occupancy * (ab12_bytes_per_stage * num_ab12_stage + c_bytes_per_stage * num_c_stage) total_smem_used = total_ab_smem + total_output_smem + occupancy * mbar_helpers_bytes + assert num_ab_stage > 0, f"insufficient smem for A/B staging: computed num_ab_stage={num_ab_stage}" + assert total_smem_used <= smem_capacity, f"smem usage {total_smem_used}B exceeds capacity {smem_capacity}B" return num_acc_stage, num_ab_stage, num_ab12_stage, num_c_stage🤖 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/dense/swiglu/dense_gemm_persistent_swiglu.py` around lines 1439 - 1444, In the stage-count validation before returning from the surrounding configuration function, add explicit checks that reject non-positive num_ab_stage and any total_smem_used exceeding the available shared-memory capacity. Use the existing total_ab_smem, total_output_smem, and mbar_helpers_bytes calculations, preserve the valid return tuple, and raise a clear error before layout construction when either guard fails.python/cudnn/gemm/ops/moe_grouped_matmul.py (1)
97-125: 📐 Maintainability & Code Quality | 🔵 Trivial
_graph_cacheis unbounded and keyed on the full token shape.MoE workloads typically see a different
total_tokenson every step, so each call adds a new compiled graph that is never evicted — steady memory growth plus a rebuild on every step. Consider a bounded LRU, or enabling dynamic-shape graphs so the M dimension is excluded from the key.Also applies to: 308-321
🤖 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/ops/moe_grouped_matmul.py` around lines 97 - 125, Update _graph_cache and the _make_cache_key call path so varying total_tokens does not create an unbounded compiled-graph entry per step. Prefer a bounded LRU cache with eviction, or enable dynamic-shape graph compilation and omit the token-count-dependent M dimension from the cache key while preserving keys for other shape and configuration differences.python/cudnn/gemm/cutedsl/dense/amax/api.py (1)
335-337: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead cleanup loop: no
sample_*attributes are ever stored.
__init__only keeps*_descdescriptors, so this loop never matches anything and gives a false sense that sample tensors are released after compile.♻️ Suggested removal
- for attr_name in tuple(vars(self)): - if attr_name.startswith("sample_"): - setattr(self, attr_name, 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/gemm/cutedsl/dense/amax/api.py` around lines 335 - 337, Remove the dead sample_* cleanup loop from the surrounding compile/reset method, since the object stores only *_desc descriptors and no sample_* attributes. Keep the existing descriptor cleanup behavior unchanged.python/cudnn/gemm/cutedsl/grouped/dswiglu/grouped_gemm_dswiglu_quant.py (2)
3536-3541: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead
total_bytesand commented-out debug block.
total_bytesis only consumed by the commentedcute.printf, so it is computed and discarded on every compile. Either drop both or keep it behind a real debug flag.🤖 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/dswiglu/grouped_gemm_dswiglu_quant.py` around lines 3536 - 3541, Remove the unused total_bytes calculation and the adjacent commented-out cute.printf debug block in the stage-count computation. Keep the return of num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, and num_tile_stage unchanged.
2826-2828: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueList comprehension used only for side effects.
Building and discarding a list to mutate
resin place obscures intent; a plain loop is clearer and allocation-free.♻️ Suggested refactor
- # let every res[?] be cute.arch.rcp_approx(res[?]) - [res.__setitem__(i, cute.arch.rcp_approx(res[i])) for i in range(cute.size(res.shape))] + # let every res[?] be cute.arch.rcp_approx(res[?]) + for i in cutlass.range_constexpr(cute.size(res.shape)): + res[i] = cute.arch.rcp_approx(res[i])🤖 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/dswiglu/grouped_gemm_dswiglu_quant.py` around lines 2826 - 2828, Replace the side-effect-only list comprehension in the reciprocal-update block with a plain loop over range(cute.size(res.shape)), assigning each res element through res.__setitem__. Preserve the subsequent res.load() behavior and the existing rcp_approx transformation.python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py (1)
21-35: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
get_divisibilityhelper. It has no calls, imports, or repository-wide references.🤖 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/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py` around lines 21 - 35, Remove the unused get_divisibility helper, including its divisibility mapping and unsupported-type error handling, without changing surrounding GEMM functionality.python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py (1)
3591-3722: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_compute_stagesreturn annotation is wrong andtotal_bytesis dead.The signature declares
Tuple[int, int, int]but five values are returned (line 3722), andtotal_bytes(line 3720) is computed then discarded. Neither affects behavior, but the annotation actively misleads callers of a 130-line helper.♻️ Fix annotation, drop dead computation
- ) -> Tuple[int, int, int]: + ) -> Tuple[int, int, int, int, int]:- total_bytes = occupancy * (ab_bytes_per_stage * num_ab_stage + epi_bytes + sinfo_bytes + mbar_helpers_bytes) - return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage🤖 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/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py` around lines 3591 - 3722, Update _compute_stages to annotate its five-element return tuple, matching the returned values num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, and num_tile_stage. Remove the unused total_bytes calculation while leaving the stage computation and return behavior unchanged.python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py (1)
19-19: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer a relative import for consistency with the migrated tree.
Sibling module
_blockscaled_api.py(line 22) usesfrom ..backend_utils import .... Using an absolutecudnn.…path inside the package is also more exposed to import-ordering surprises given the lazy-optional-dependency setup inpython/cudnn/__init__.py. Separately,_require_pointer_tensoris a private helper being pulled across thegrouped/→discrete_grouped/boundary; consider promoting it into a shared helpers module (e.g.grouped/backend_utils.py) if both families need it.♻️ Relative import
-from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from ...discrete_grouped.discrete_kernel_utils import _require_pointer_tensor🤖 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/dglu/_bf16_api.py` at line 19, Replace the absolute import of _require_pointer_tensor in the grouped dglu API module with the appropriate relative import, matching the migrated tree’s sibling modules and avoiding package-root resolution. Keep the helper behavior unchanged; only relocate it to a shared grouped backend utility if required to make the relative import valid for both grouped families.python/cudnn/gemm/cutedsl/grouped/moe_kernel_helpers.py (1)
461-477: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
if True:guards an unconditional debug dump — remove or gate it.
compare_and_report_mismatchesprints the first 8 elements plus a header on every invocation, including successful comparisons, before it even evaluates tolerances. That is a leftover debug artifact and makes passing test runs noisy.♻️ Gate behind an opt-in flag
-def compare_and_report_mismatches( - gpu_tensor, - ref_tensor, - name="Tensor", - atol=1e-05, - rtol=1e-05, - max_mismatches=8, -): +def compare_and_report_mismatches( + gpu_tensor, + ref_tensor, + name="Tensor", + atol=1e-05, + rtol=1e-05, + max_mismatches=8, + dump_head=False, +):- if True: + if dump_head: print(f"\n{name} - First 8 elements:")🤖 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/moe_kernel_helpers.py` around lines 461 - 477, Remove the unconditional `if True:` debug block in `compare_and_report_mismatches`, or gate the header and first-eight-elements dump behind an explicit opt-in debug flag. Ensure successful comparisons remain silent by default while preserving the existing diagnostic output when debugging is enabled.python/cudnn/gemm/cutedsl/dense/srelu/api.py (2)
85-86: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the unused C/D shape unpacks.
c_m/c_n/c_landd_m/d_n/d_lare never read —_check_tensor_shapeon lines 89-90 already enforces those shapes. Ruff flags all six (RUF059).🧹 Suggested cleanup
- c_m, c_n, c_l = self._tensor_shape(self.c_desc, name="sample_c") - d_m, d_n, d_l = self._tensor_shape(self.d_desc, name="sample_d")🤖 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/dense/srelu/api.py` around lines 85 - 86, Remove the unused c_m/c_n/c_l and d_m/d_n/d_l assignments from the shape-validation flow in the surrounding method, while preserving the existing _check_tensor_shape calls that enforce C and D tensor shapes.Source: Linters/SAST tools
563-564: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace
assertwith an explicit raise in the public wrapper.
assertis stripped underpython -O, so an unsupported configuration would fall through tocompile()instead of failing fast. The other relocated wrappers (python/cudnn/gemm/cutedsl/grouped/wgrad/api.py,python/cudnn/gemm/cutedsl/grouped/dglu/api.py) raiseRuntimeError; the"Unsupported testcase"wording also reads as test-only for a public entry point.♻️ Suggested change
- assert op.check_support(), "Unsupported testcase" - op.compile() + if not op.check_support(): + raise RuntimeError("Unsupported configuration for gemm_srelu_wrapper_sm100") + op.compile()🤖 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/dense/srelu/api.py` around lines 563 - 564, Replace the assert guarding op.check_support() in the public wrapper with an explicit RuntimeError using production-appropriate unsupported-configuration wording, and keep op.compile() unreachable when support validation fails. Match the established error-handling pattern in the relocated grouped wrappers.python/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.py (5)
772-776: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGuard
num_ab_stageand drop the unusedtotal_bytes.If
mbar_helpers_bytes + epi_bytes + sinfo_bytesexceeds the per-occupancy SMEM budget,num_ab_stagebecomes<= 0and is returned as a valid pipeline depth, pushing the failure deep into kernel construction.total_byteson line 774 is computed but never used or returned.🛡️ Suggested fix
num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage - - total_bytes = occupancy * (ab_bytes_per_stage * num_ab_stage + epi_bytes + sinfo_bytes + mbar_helpers_bytes) + if num_ab_stage < 1: + raise ValueError( + f"Insufficient shared memory for A/B pipeline: capacity={num_smem_capacity}, " + f"occupancy={occupancy}, epilogue+mbar+sinfo={mbar_helpers_bytes + epi_bytes + sinfo_bytes}" + ) return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage🤖 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/discrete_grouped/discrete_kernel_utils.py` around lines 772 - 776, Update the pipeline-depth calculation in the surrounding utility function to validate num_ab_stage after computing it, rejecting or otherwise handling values less than or equal to zero before returning the stage counts. Remove the unused total_bytes calculation, while preserving the existing return tuple for valid configurations.
46-56: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_require_pointer_tensoris underscore-private but consumed cross-package.
grouped/dglu/api.py,grouped/dsrelu/api.py, andgrouped/quant/api.pyall import it from thisdiscrete_groupedmodule via function-local imports. Since the relocation already establishes a shared-helper layer, consider hosting it next to the other shared validators (e.g.grouped/backend_utils.py) under a non-underscore name so the dependency direction and visibility match the actual usage.🤖 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/discrete_grouped/discrete_kernel_utils.py` around lines 46 - 56, Move _require_pointer_tensor into the shared validator module alongside the other backend utilities, rename it to a non-underscore public helper, and update the function-local imports and call sites in grouped/dglu/api.py, grouped/dsrelu/api.py, and grouped/quant/api.py to use the new symbol. Remove the obsolete definition or imports from discrete_grouped while preserving all existing validation behavior.
125-127: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional lint cleanup in the relocated PTX helpers.
Ruff flags the constant-only f-strings (F541 at 125/127/133/143/145/151/186/188/233/235) and the
is_power_of_2lambda assignment (E731 at 570). Since these lines are being touched by the move anyway, dropping the strayfprefixes and converting the lambda to adefkeeps the new package lint-clean.Also applies to: 143-145, 186-188, 233-235, 570-570
🤖 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/discrete_grouped/discrete_kernel_utils.py` around lines 125 - 127, Remove the unnecessary f-string prefixes from the constant PTX instruction assignments in the relocated helpers, including the branches around min/max and related instructions. Replace the is_power_of_2 lambda assignment with a named def in the same utility scope, preserving its existing behavior and call sites.Source: Linters/SAST tools
405-421: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the
if True:debug block.This dead conditional unconditionally dumps the first 8 elements to stdout on every call, including when validation passes. Ruff also flags line 409 (
f"\n"with no placeholders).🧹 Suggested cleanup
- if True: - print(f"\n{name} - First 8 elements:") - print(f"{'Index':<6} {'Coordinate':<30} {'GPU Data':<20} {'CPU Data':<20} {'Abs Error':<20}") - print("-" * 100) - print(f"\n") - - flat_gpu = gpu_data.flatten() - flat_ref = ref_data.flatten() - num_elements = min(8, flat_gpu.numel()) - - for i in range(num_elements): - idx_tuple = torch.unravel_index(torch.tensor(i), gpu_data.shape) - coord = tuple(idx.item() for idx in idx_tuple) - gpu_val = gpu_data[coord].item() - ref_val = ref_data[coord].item() - abs_error = abs(gpu_val - ref_val) - print(f"{i + 1:<6} {str(coord):<30} {gpu_val:<20.6f} {ref_val:<20.6f} {abs_error:<20.6f}") - diff = torch.abs(gpu_data - ref_data)🤖 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/discrete_grouped/discrete_kernel_utils.py` around lines 405 - 421, Remove the unconditional debug-print block beginning with if True, including the first-eight-elements formatting and iteration, so the validation helper no longer writes to stdout on every call. This also eliminates the unnecessary f"\n" expression flagged by Ruff.Source: Linters/SAST tools
268-286: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueRemove or reject the unused
positive_onlymode.The signed
i32atomicMAXis valid only for non-negative FP32 values. Current callers use the default with non-negativeamaxreductions, butpositive_only=Falsewould silently mishandle negative inputs.🤖 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/discrete_grouped/discrete_kernel_utils.py` around lines 268 - 286, Remove the unused positive_only parameter from atomic_max_float32 and update all call sites to stop passing it. Do not retain a mode that permits negative inputs with signed i32 atomic MAX; if the API must remain configurable, explicitly reject positive_only=False before the atomic operation.python/cudnn/gemm/cutedsl/grouped/backend_utils.py (1)
36-54: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd annotations and report every forbidden scale control.
The rest of this package is fully annotated; this selector is the one untyped public helper. It also surfaces only
forbidden[0], so a caller passing several BF16-incompatible scale controls has to fix them one round-trip at a time.♻️ Suggested change
def select_grouped_gemm_backend( *, - operation, - a_dtype, - b_dtype, - scale_controls, - block_scaled_dtype_pairs, -): + operation: str, + a_dtype: torch.dtype, + b_dtype: torch.dtype, + scale_controls: Iterable[tuple[str, object]], + block_scaled_dtype_pairs: Container[tuple[torch.dtype, torch.dtype]], +) -> GroupedGemmBackend: bf16_operands = (a_dtype == torch.bfloat16, b_dtype == torch.bfloat16) if any(bf16_operands): if not all(bf16_operands): raise ValueError(f"{operation}: mixed dtype families: a_dtype={a_dtype}, " f"b_dtype={b_dtype}") forbidden = [name for name, value in scale_controls if value is not None] if forbidden: - raise ValueError(f"{operation}: BF16 forbids scale control {forbidden[0]}") + raise ValueError(f"{operation}: BF16 forbids scale controls {forbidden}") return GroupedGemmBackend.BF16Requires widening the typing import:
-from typing import Iterator, Optional +from typing import Container, Iterable, Iterator, Optional🤖 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/backend_utils.py` around lines 36 - 54, Add complete type annotations to the public select_grouped_gemm_backend function, including its parameters and return type, using the package’s existing typing conventions and symbols. In the BF16 branch, report all non-None forbidden scale-control names rather than only forbidden[0], while preserving the existing validation and backend-selection behavior.python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py (1)
3574-3705: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStale return annotation and dead local in
_compute_stages.The signature/docstring advertise a 3-tuple but the function returns 5 values (
num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage), matching the unpack at lines 467-473.total_bytes(line 3703) is computed and discarded.♻️ Proposed cleanup
- ) -> Tuple[int, int, int]: + ) -> Tuple[int, int, int, int, int]:num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage - total_bytes = occupancy * (ab_bytes_per_stage * num_ab_stage + epi_bytes + sinfo_bytes + mbar_helpers_bytes) - return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage🤖 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/discrete_grouped/dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py` around lines 3574 - 3705, Update _compute_stages to annotate and document its actual five-value return tuple, including ACC, A/B, C, D, and tile stages, consistent with its callers. Remove the unused total_bytes calculation since it is computed and discarded.python/cudnn/gemm/cutedsl/grouped/moe_persistent_scheduler.py (2)
920-925: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring doesn't match the body.
"Gather a full warp of arrives per CTA" describes a fan-in that isn't here — the body issues a single
arriveon the empty barrier.📝 Proposed docstring fix
- """Gather a full warp of arrives per CTA to leader's empty barrier.""" + """Signal the leader's empty barrier that this CTA consumed the broadcast slot."""🤖 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/moe_persistent_scheduler.py` around lines 920 - 925, Update the docstring for _cluster_consumer_release to describe its actual behavior: issuing a single arrive on the consumer empty barrier using the current consumer state index. Remove the inaccurate wording about gathering a full warp of arrivals per CTA.
244-261: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the forced
use_dynamic_schedfor2Dx2D.A caller passing
use_dynamic_sched=Falsewithscenario="2Dx2D"silently gets dynamic scheduling. It's required by the CLC path ininternal_init, but the class docstring doesn't mention it — worth a line so callers don't debug why their static request was ignored.🤖 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/moe_persistent_scheduler.py` around lines 244 - 261, The constructor’s forced dynamic-scheduling behavior for scenario "2Dx2D" is undocumented. Update the class docstring associated with the scheduler to state that "2Dx2D" requires and forces use_dynamic_sched=True, while leaving the existing assignment logic in __init__ unchanged.python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py (1)
132-134: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the class-level padding constant consistently. The N=192 branch reads the separately imported
FIX_PAD_SIZE, while the constructor and fallback useBlockScaledMoEGroupedGemmQuantKernel.FIX_PAD_SIZE. Use the class attribute here to prevent future divergence.🤖 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/moe_blockscaled_grouped_gemm_quant_rubin.py` around lines 132 - 134, Update the N=192 branch in the relevant kernel validation logic to compare m_aligned against BlockScaledMoEGroupedGemmQuantKernel.FIX_PAD_SIZE instead of the separately imported FIX_PAD_SIZE, matching the constructor and fallback behavior.python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py (2)
897-898: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCommented-out debug
printfblocks carried over.Two dead debug blocks (scheduler warp and TMA warp). Good opportunity to drop them during the move.
Also applies to: 1004-1013
🤖 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/wgrad/moe_blockscaled_grouped_gemm_wgrad.py` around lines 897 - 898, Remove the commented-out debug printf blocks near the scheduler warp and TMA warp, including the block containing work_tile_info and the additional block around lines 1004-1013. Leave the surrounding execution logic unchanged.
20-23: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStale extension names in the module docstring.
moe_sched_extension.pydefinesWgradScaledGemmSchedExtension, notWgradDense / WgradDiscrete.🤖 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/wgrad/moe_blockscaled_grouped_gemm_wgrad.py` around lines 20 - 23, Update the module docstring near the scheduler and extension references to replace the stale “WgradDense / WgradDiscrete” names with the actual WgradScaledGemmSchedExtension symbol, leaving the surrounding documentation unchanged.python/cudnn/gemm/cutedsl/grouped/moe_sched_extension.py (1)
10-22: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring lists only two of the six extensions in this module.
WgradScaledGemmSchedExtension,DiscreteWeightGroupedGemmSchedExtension,ContiguousGroupedGemmSchedExtension, andWgradGemmSchedExtensionare also defined here. Worth extending the "Two concrete extensions are provided" section while the file is being relocated.🤖 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/moe_sched_extension.py` around lines 10 - 22, Update the module-level docstring’s extension overview to list all six defined extensions, including WgradScaledGemmSchedExtension, DiscreteWeightGroupedGemmSchedExtension, ContiguousGroupedGemmSchedExtension, and WgradGemmSchedExtension, with an accurate brief description for each while preserving the existing entries.python/cudnn/gemm/cutedsl/grouped/wgrad/api.py (1)
15-17: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRelocated modules reach back through the top-level
cudnnnamespace instead of using relative imports. After the move, several intra-package imports are absolute (from cudnn....), which mixes styles within the same package and re-enters the parent package during import — fragile for the lazy optional-import boundary.
python/cudnn/gemm/cutedsl/grouped/wgrad/api.py#L15-L17: switch line 15 tofrom ....discrete_grouped.discrete_kernel_utils import _require_pointer_tensorso it matches the relative..backend_utilsimport on line 17.python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py#L16-L16: apply the same relative form for_require_pointer_tensor.python/cudnn/gemm/reference/reference_matmul_engine.py#L22-L27: usefrom ...engines.base import BaseEngine,from ...engines.engine_ids import PYTHON_ENGINE_ID_BASE, andfrom ...graph_types import NodeType.🤖 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/wgrad/api.py` around lines 15 - 17, Replace the absolute intra-package imports with relative imports to preserve the lazy optional-import boundary: in python/cudnn/gemm/cutedsl/grouped/wgrad/api.py lines 15-17 and python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py line 16, update the _require_pointer_tensor import to use the relative discrete_grouped path; in python/cudnn/gemm/reference/reference_matmul_engine.py lines 22-27, update BaseEngine, PYTHON_ENGINE_ID_BASE, and NodeType to their specified relative engine and graph_types paths.python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py (1)
65-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueParenthesize the mixed
*///expression.
max(tile_n // 128, 1) * tile_k // sf_vec_sizeevaluates as(max(...) * tile_k) // sf_vec_size. It is equivalent here only becausetile_k % sf_vec_size == 0is validated above; explicit parentheses make the intent robust to future edits.🤖 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/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py` at line 65, Parenthesize the scaling-factor column calculation assigned to sfb_columns so the multiplication and integer division order is explicit, preserving the current result while making the intended grouping robust to future changes.
🤖 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/gemm/cutedsl/dense/amax/api.py`:
- Around line 105-108: Remove the unused f-string prefix from the error message
passed in the _value_error_if call, leaving the message text unchanged.
In
`@python/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py`:
- Around line 1383-1390: Reject FP8 D configurations requiring SFD before the
epilogue store in GemmDsreluSm100’s
is_valid_dtypes_and_scale_factor_vec_size/can_implement or check_support path,
rather than printing and storing uninitialized tRS_rD; apply equivalent
can_implement gating in
python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py
(lines 1374-1381), while preserving supported non-SFD behavior.
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Line 539: Update the cache-key construction in the affected API paths to
include cos.dtype and sin.dtype, plus x_scale.dtype and w_scale.dtype for the
MXFP8 path, alongside the existing shape, output, and device fields. Ensure
calls with identical shapes but different rotary-table or scale dtypes cannot
reuse an incompatible compiled kernel.
- Line 525: Replace the public-input `assert` validations in the affected
wrapper checks, including the dtype check and the guards at the referenced
scale/output validation sites, with explicit `ValueError` raises that preserve
the current messages and conditions. Ensure mismatched dtypes or scales still
fail before `from_dlpack` or kernel launch, including when Python runs with
optimizations enabled.
In
`@python/cudnn/gemm/cutedsl/dense/swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py`:
- Around line 2080-2103: Update the local fmin helper to honor its nan parameter
by selecting the appropriate PTX min.f32 or NaN-propagating variant, matching
discrete_grouped/discrete_kernel_utils.fmin; preferably reuse that shared helper
instead of maintaining a divergent copy. Remove the unnecessary f-string
prefixes from the placeholder-free PTX and constraint strings.
In
`@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py`:
- Around line 3166-3195: In the store_d_directly branch, replace the undefined
thr_copy_t2r reference used to partition gD_sub_loop with a thread slice derived
from the existing tiled_copy_t2r binding. Preserve the current tCgD1/tCgD2
indexing and global-store behavior, ensuring the branch traces successfully when
store_d_directly is enabled.
In
`@python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py`:
- Around line 3187-3196: Bind the per-thread copy slice before partitioning D in
both direct-store epilogues: in
python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py
lines 3187-3196 and
python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py
lines 3604-3615, create thr_copy_t2r from tiled_copy_t2r.get_slice(epi_tidx)
before calling partition_D; both sites require the same change.
In
`@python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py`:
- Around line 2027-2046: Update the Rubin dgeglu flow around the dgeglu method
to accept the per-expert square_alpha and beta_val values, then apply them to
the accumulator and C operands before computing gradients. Thread both values
from their existing alpha_val/beta_val loads at the call site, matching sibling
implementations, while preserving the current behavior when scales are 1.0.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py`:
- Around line 810-837: Update the cache-miss initialization in the grouped GEMM
wrapper so a caller-provided amax_tensor_buf is assigned to amax_tensor
regardless of d_dtype, including fp8 paths. Preserve allocation of
cached_amax_tensor only for bfloat16/float16 when no caller buffer is supplied,
and ensure the first invocation matches cache-hit behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/grouped_gemm_dswiglu_quant.py`:
- Around line 601-602: Update the condition in the generate_sfd handling block
to use the idiomatic negation of self.generate_sfd instead of comparing it
explicitly to False, preserving the existing cutlass.const_expr behavior.
In
`@python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py`:
- Around line 141-154: Update the `use_2cta_instrs` and `mma_tiler_mn[0] == 512`
validation branch in `can_implement` to reject `m_aligned` values that are not
divisible by the 512-row tile height. Add the alignment check before returning
success, while preserving the existing dtype, layout, tensor-alignment, tile-N,
and cluster-M validations.
In
`@python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py`:
- Around line 1485-1490: Remove the duplicated current_scale_b_iter assignment
in the block handling global scales. Update the guard around
current_scale_a_iter/current_scale_b_iter and alpha calculation to require both
global_scale_a and global_scale_b to be present before dereferencing either
iterator, preserving safe behavior for an a-only call path.
In `@python/cudnn/gemm/ops/moe_grouped_matmul.py`:
- Around line 83-89: Update _get_handle and its callers to avoid sharing a
mutable cuDNN handle across concurrent streams: use a handle scoped per thread
or per CUDA stream, and ensure handle creation is synchronized as needed.
Preserve the existing current-stream selection and pass the isolated handle
through the execution path around graph.execute so concurrent submissions cannot
interleave set_stream with execution.
---
Nitpick comments:
In `@python/cudnn/__init__.py`:
- Around line 332-361: Update the symbol entries in _LAZY_OPTIONAL_IMPORTS so
each grouped and discrete-grouped kernel class or wrapper resolves through its
specific fusion package instead of the shared family module path. Keep the
grouped_gemm and discrete_grouped_gemm module-alias entries on their existing
family paths, and preserve each symbol’s exported attribute name.
In `@python/cudnn/AGENTS.md`:
- Around line 22-31: Change the opening fence for the package-layout tree to a
text-labeled fence in both python/cudnn/AGENTS.md lines 22-31 and
python/cudnn/README.md lines 32-40; leave the tree contents unchanged.
In `@python/cudnn/gemm/cutedsl/dense/amax/api.py`:
- Around line 335-337: Remove the dead sample_* cleanup loop from the
surrounding compile/reset method, since the object stores only *_desc
descriptors and no sample_* attributes. Keep the existing descriptor cleanup
behavior unchanged.
In `@python/cudnn/gemm/cutedsl/dense/srelu/api.py`:
- Around line 85-86: Remove the unused c_m/c_n/c_l and d_m/d_n/d_l assignments
from the shape-validation flow in the surrounding method, while preserving the
existing _check_tensor_shape calls that enforce C and D tensor shapes.
- Around line 563-564: Replace the assert guarding op.check_support() in the
public wrapper with an explicit RuntimeError using production-appropriate
unsupported-configuration wording, and keep op.compile() unreachable when
support validation fails. Match the established error-handling pattern in the
relocated grouped wrappers.
In
`@python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py`:
- Around line 21-35: Remove the unused get_divisibility helper, including its
divisibility mapping and unsupported-type error handling, without changing
surrounding GEMM functionality.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/dense_gemm_persistent_swiglu.py`:
- Around line 1439-1444: In the stage-count validation before returning from the
surrounding configuration function, add explicit checks that reject non-positive
num_ab_stage and any total_smem_used exceeding the available shared-memory
capacity. Use the existing total_ab_smem, total_output_smem, and
mbar_helpers_bytes calculations, preserve the valid return tuple, and raise a
clear error before layout construction when either guard fails.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.py`:
- Around line 772-776: Update the pipeline-depth calculation in the surrounding
utility function to validate num_ab_stage after computing it, rejecting or
otherwise handling values less than or equal to zero before returning the stage
counts. Remove the unused total_bytes calculation, while preserving the existing
return tuple for valid configurations.
- Around line 46-56: Move _require_pointer_tensor into the shared validator
module alongside the other backend utilities, rename it to a non-underscore
public helper, and update the function-local imports and call sites in
grouped/dglu/api.py, grouped/dsrelu/api.py, and grouped/quant/api.py to use the
new symbol. Remove the obsolete definition or imports from discrete_grouped
while preserving all existing validation behavior.
- Around line 125-127: Remove the unnecessary f-string prefixes from the
constant PTX instruction assignments in the relocated helpers, including the
branches around min/max and related instructions. Replace the is_power_of_2
lambda assignment with a named def in the same utility scope, preserving its
existing behavior and call sites.
- Around line 405-421: Remove the unconditional debug-print block beginning with
if True, including the first-eight-elements formatting and iteration, so the
validation helper no longer writes to stdout on every call. This also eliminates
the unnecessary f"\n" expression flagged by Ruff.
- Around line 268-286: Remove the unused positive_only parameter from
atomic_max_float32 and update all call sites to stop passing it. Do not retain a
mode that permits negative inputs with signed i32 atomic MAX; if the API must
remain configurable, explicitly reject positive_only=False before the atomic
operation.
In
`@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py`:
- Around line 3574-3705: Update _compute_stages to annotate and document its
actual five-value return tuple, including ACC, A/B, C, D, and tile stages,
consistent with its callers. Remove the unused total_bytes calculation since it
is computed and discarded.
In
`@python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py`:
- Around line 1346-1348: Remove the unused tile_info_to_mn_idx call and
m_idx/n_idx unpacking from create_and_partition_new_SFDCol, retaining expert_idx
and the compute_expert_token_range flow unchanged.
- Around line 2430-2433: Remove the unused d_col_pipeline construction by
deleting the pipeline.PipelineTmaStore.create call near the d_pipeline setup.
Keep d_pipeline.producer_tail() as the sole mechanism covering the d_col TMA
store, and leave the surrounding pipeline configuration unchanged.
In `@python/cudnn/gemm/cutedsl/grouped/backend_utils.py`:
- Around line 36-54: Add complete type annotations to the public
select_grouped_gemm_backend function, including its parameters and return type,
using the package’s existing typing conventions and symbols. In the BF16 branch,
report all non-None forbidden scale-control names rather than only forbidden[0],
while preserving the existing validation and backend-selection behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Line 19: Replace the absolute import of _require_pointer_tensor in the grouped
dglu API module with the appropriate relative import, matching the migrated
tree’s sibling modules and avoiding package-root resolution. Keep the helper
behavior unchanged; only relocate it to a shared grouped backend utility if
required to make the relative import valid for both grouped families.
In
`@python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py`:
- Around line 3591-3722: Update _compute_stages to annotate its five-element
return tuple, matching the returned values num_acc_stage, num_ab_stage,
num_c_stage, num_d_stage, and num_tile_stage. Remove the unused total_bytes
calculation while leaving the stage computation and return behavior unchanged.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py`:
- Around line 664-668: Move the os and logging imports from the block below the
class to the module header, alongside the existing top-level imports. Keep the
_logger initialization at module scope near the header, and remove the
duplicated lower block without changing the cache definition or other class
behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/grouped_gemm_dswiglu_quant.py`:
- Around line 3536-3541: Remove the unused total_bytes calculation and the
adjacent commented-out cute.printf debug block in the stage-count computation.
Keep the return of num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, and
num_tile_stage unchanged.
- Around line 2826-2828: Replace the side-effect-only list comprehension in the
reciprocal-update block with a plain loop over range(cute.size(res.shape)),
assigning each res element through res.__setitem__. Preserve the subsequent
res.load() behavior and the existing rcp_approx transformation.
In
`@python/cudnn/gemm/cutedsl/grouped/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py`:
- Line 72: Remove the unused imported amax_reduction_per_thread from the import
list, keeping the class method amax_reduction_per_thread used through
self.amax_reduction_per_thread as the single implementation.
In `@python/cudnn/gemm/cutedsl/grouped/moe_kernel_helpers.py`:
- Around line 461-477: Remove the unconditional `if True:` debug block in
`compare_and_report_mismatches`, or gate the header and first-eight-elements
dump behind an explicit opt-in debug flag. Ensure successful comparisons remain
silent by default while preserving the existing diagnostic output when debugging
is enabled.
In `@python/cudnn/gemm/cutedsl/grouped/moe_persistent_scheduler.py`:
- Around line 920-925: Update the docstring for _cluster_consumer_release to
describe its actual behavior: issuing a single arrive on the consumer empty
barrier using the current consumer state index. Remove the inaccurate wording
about gathering a full warp of arrivals per CTA.
- Around line 244-261: The constructor’s forced dynamic-scheduling behavior for
scenario "2Dx2D" is undocumented. Update the class docstring associated with the
scheduler to state that "2Dx2D" requires and forces use_dynamic_sched=True,
while leaving the existing assignment logic in __init__ unchanged.
In `@python/cudnn/gemm/cutedsl/grouped/moe_sched_extension.py`:
- Around line 10-22: Update the module-level docstring’s extension overview to
list all six defined extensions, including WgradScaledGemmSchedExtension,
DiscreteWeightGroupedGemmSchedExtension, ContiguousGroupedGemmSchedExtension,
and WgradGemmSchedExtension, with an accurate brief description for each while
preserving the existing entries.
In
`@python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py`:
- Around line 132-134: Update the N=192 branch in the relevant kernel validation
logic to compare m_aligned against
BlockScaledMoEGroupedGemmQuantKernel.FIX_PAD_SIZE instead of the separately
imported FIX_PAD_SIZE, matching the constructor and fallback behavior.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/api.py`:
- Around line 15-17: Replace the absolute intra-package imports with relative
imports to preserve the lazy optional-import boundary: in
python/cudnn/gemm/cutedsl/grouped/wgrad/api.py lines 15-17 and
python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py line 16, update the
_require_pointer_tensor import to use the relative discrete_grouped path; in
python/cudnn/gemm/reference/reference_matmul_engine.py lines 22-27, update
BaseEngine, PYTHON_ENGINE_ID_BASE, and NodeType to their specified relative
engine and graph_types paths.
In
`@python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py`:
- Line 65: Parenthesize the scaling-factor column calculation assigned to
sfb_columns so the multiplication and integer division order is explicit,
preserving the current result while making the intended grouping robust to
future changes.
In
`@python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py`:
- Around line 897-898: Remove the commented-out debug printf blocks near the
scheduler warp and TMA warp, including the block containing work_tile_info and
the additional block around lines 1004-1013. Leave the surrounding execution
logic unchanged.
- Around line 20-23: Update the module docstring near the scheduler and
extension references to replace the stale “WgradDense / WgradDiscrete” names
with the actual WgradScaledGemmSchedExtension symbol, leaving the surrounding
documentation unchanged.
In `@python/cudnn/gemm/ops/moe_grouped_matmul.py`:
- Around line 97-125: Update _graph_cache and the _make_cache_key call path so
varying total_tokens does not create an unbounded compiled-graph entry per step.
Prefer a bounded LRU cache with eviction, or enable dynamic-shape graph
compilation and omit the token-count-dependent M dimension from the cache key
while preserving keys for other shape and configuration differences.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
🛑 Comments failed to post (13)
python/cudnn/gemm/cutedsl/dense/amax/api.py (1)
105-108: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Drop the unused
fprefix.Ruff flags F541 here; the string has no placeholders.
🧹 Proposed fix
- f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported", + "Unsupported ab_dtype and sf_vec_size combination: {float8_e5m2, float8_e4m3fn} and 16 is not supported",📝 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._value_error_if( ab_dtype in {torch.float8_e5m2, torch.float8_e4m3fn} and self.sf_vec_size == 16, "Unsupported ab_dtype and sf_vec_size combination: {float8_e5m2, float8_e4m3fn} and 16 is not supported", )🧰 Tools
🪛 Ruff (0.16.0)
[error] 107-107: f-string without any placeholders
Remove extraneous
fprefix(F541)
🤖 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/dense/amax/api.py` around lines 105 - 108, Remove the unused f-string prefix from the error message passed in the _value_error_if call, leaving the message text unchanged.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py (1)
1383-1390: 🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
Unimplemented SFD epilogue in both relocated dense kernels writes uninitialized
D. Both kernels take thegenerate_sfdbranch, print"SFD not implemented", and fall through to the shared-memory/TMA store without ever populatingtRS_rD— the store then commits stale register contents. The frontend requiressfd/norm_constwheneverDis FP8, so this path is reachable from a supported configuration rather than dead code.
python/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py#L1383-L1390: reject FP8Dinis_valid_dtypes_and_scale_factor_vec_size/can_implement(or raise fromGemmDsreluSm100.check_support) until the quantized epilogue lands, instead of printing and storing garbage.python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py#L1374-L1381: apply the same gating in this kernel'scan_implementso the sReLU API surfaces an explicit error too.📍 Affects 2 files
python/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py#L1383-L1390(this comment)python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py#L1374-L1381🤖 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/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py` around lines 1383 - 1390, Reject FP8 D configurations requiring SFD before the epilogue store in GemmDsreluSm100’s is_valid_dtypes_and_scale_factor_vec_size/can_implement or check_support path, rather than printing and storing uninitialized tRS_rD; apply equivalent can_implement gating in python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py (lines 1374-1381), while preserving supported non-SFD behavior.python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py (2)
525-525: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Use explicit exceptions instead of
assertfor public-input validation.These guards validate user-supplied dtypes/scales at the public entry point, but
assertis elided underpython -O, at which point mismatched inputs fall through intofrom_dlpack/kernel launch. The sibling wrappers in this PR (gemm_amax_wrapper_sm100,_allocate_dense_output) raiseValueErrorfor the same class of check.🛡️ Proposed fix
- assert x.dtype == w.dtype, f"x and w must share a dtype (both bfloat16 or both float8_e4m3fn); got x {x.dtype}, w {w.dtype}" + if x.dtype != w.dtype: + raise ValueError(f"x and w must share a dtype (both bfloat16 or both float8_e4m3fn); got x {x.dtype}, w {w.dtype}")- assert x_scale is None and w_scale is None, "bf16 inputs must not be given MXFP8 scales (x_scale/w_scale); those are for the float8_e4m3fn path" + if x_scale is not None or w_scale is not None: + raise ValueError("bf16 inputs must not be given MXFP8 scales (x_scale/w_scale); those are for the float8_e4m3fn path")- assert x_scale is not None and w_scale is not None, "MXFP8 (float8_e4m3fn) inputs require x_scale and w_scale (E8M0 rowwise block scales)" + if x_scale is None or w_scale is None: + raise ValueError("MXFP8 (float8_e4m3fn) inputs require x_scale and w_scale (E8M0 rowwise block scales)")- raise AssertionError(f"unsupported input dtype {x.dtype}; expected bfloat16 (BF16 GEMM) or float8_e4m3fn (MXFP8 GEMM)") + raise ValueError(f"unsupported input dtype {x.dtype}; expected bfloat16 (BF16 GEMM) or float8_e4m3fn (MXFP8 GEMM)")Also applies to: 537-538, 558-559, 597-598
🤖 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/dense/proj_rope_mxfp8/api.py` at line 525, Replace the public-input `assert` validations in the affected wrapper checks, including the dtype check and the guards at the referenced scale/output validation sites, with explicit `ValueError` raises that preserve the current messages and conditions. Ensure mismatched dtypes or scales still fail before `from_dlpack` or kernel launch, including when Python runs with optimizations enabled.
539-539: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Cache key omits
cos/sin(and scale) dtypes.
check_supportvalidatescos/sindtype only when the object is first constructed; later calls with the samex/wshapes but e.g. fp32 rotary tables reuse the compiled kernel and reinterpret memory as bf16. Shapes are derivable fromtokens, but dtypes are not — consider addingcos.dtype/sin.dtype(andx_scale.dtype/w_scale.dtypeon the MXFP8 path) to the key.Also applies to: 565-565
🤖 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/dense/proj_rope_mxfp8/api.py` at line 539, Update the cache-key construction in the affected API paths to include cos.dtype and sin.dtype, plus x_scale.dtype and w_scale.dtype for the MXFP8 path, alongside the existing shape, output, and device fields. Ensure calls with identical shapes but different rotary-table or scale dtypes cannot reuse an incompatible compiled kernel.python/cudnn/gemm/cutedsl/dense/swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py (1)
2080-2103: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
fminsilently ignores itsnanargument.Callers at lines 1596, 1597, and 1609 pass
nan=True, but the template is unconditionallymin.f32— nomin.NaN.f32variant.discrete_grouped/discrete_kernel_utils.fminimplements the same helper correctly; consider importing/sharing it instead of a local divergent copy. (Thefprefixes on the two placeholder-free strings are also extraneous per Ruff F541.)🐛 Proposed fix
- ptx_instr = f"min.f32 $0, $1, $2;" + ptx_instr = "min.NaN.f32 $0, $1, $2;" if nan else "min.f32 $0, $1, $2;" return cutlass.Float32( cutlass._mlir.dialects.llvm.inline_asm( cutlass.cutlass_dsl.T.f32(), [ cutlass.Float32(a).ir_value(loc=loc, ip=ip), cutlass.Float32(b).ir_value(loc=loc, ip=ip), ], - f"{ptx_instr}", - f"=f,f,f", + ptx_instr, + "=f,f,f",📝 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.# TODO `@mingyangw` def fmin( a: Union[float, cutlass.Float32], b: Union[float, cutlass.Float32], *, loc=None, ip=None, nan=True, ) -> cutlass.Float32: ptx_instr = "min.NaN.f32 $0, $1, $2;" if nan else "min.f32 $0, $1, $2;" return cutlass.Float32( cutlass._mlir.dialects.llvm.inline_asm( cutlass.cutlass_dsl.T.f32(), [ cutlass.Float32(a).ir_value(loc=loc, ip=ip), cutlass.Float32(b).ir_value(loc=loc, ip=ip), ], ptx_instr, "=f,f,f", has_side_effects=True, is_align_stack=False, asm_dialect=cutlass._mlir.dialects.llvm.AsmDialect.AD_ATT, ) )🧰 Tools
🪛 Ruff (0.16.0)
[error] 2089-2089: f-string without any placeholders
Remove extraneous
fprefix(F541)
[error] 2098-2098: f-string without any placeholders
Remove extraneous
fprefix(F541)
🤖 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/dense/swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py` around lines 2080 - 2103, Update the local fmin helper to honor its nan parameter by selecting the appropriate PTX min.f32 or NaN-propagating variant, matching discrete_grouped/discrete_kernel_utils.fmin; preferably reuse that shared helper instead of maintaining a divergent copy. Remove the unnecessary f-string prefixes from the placeholder-free PTX and constraint strings.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py (1)
3166-3195: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Undefined
thr_copy_t2rin thestore_d_directlypath.Line 3174 uses
thr_copy_t2r, which is never bound in the epilogue scope — onlytiled_copy_t2r(line 2733) andthr_copy_t2r_local(line 2862) exist. The branch is currently unreachable becausestore_d_directlyis hardcodedFalseat line 464, but re-enabling the commented-out condition there will fail at trace time. Bind the thread slice fromtiled_copy_t2rrather than relying on the branch staying dead.🐛 Proposed fix
gD_sub_loop = cute.local_tile(real_d, d_epilogue_subtile, (None, None, None)) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) tCgD_mnl_loop = thr_copy_t2r.partition_D(gD_sub_loop)📝 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 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)) thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) tCgD_mnl_loop = thr_copy_t2r.partition_D(gD_sub_loop) tCgD_mnl_loop = cute.filter_zeros(tCgD_mnl_loop) tCgD1 = tCgD_mnl_loop[ ( None, 0, # T2R_M 2 * real_subtile_idx + 0, # T2R_N *d_idx_mn, # RestM/N 0, # RestL ) ] tCgD2 = tCgD_mnl_loop[ ( None, 0, # T2R_M 2 * real_subtile_idx + 1, # T2R_N *d_idx_mn, # RestM/N 0, # RestL ) ] self.store_global_memory_256b(tCgD1, tRS_rD1) self.store_global_memory_256b(tCgD2, tRS_rD2)🧰 Tools
🪛 Ruff (0.16.0)
[error] 3174-3174: 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/gemm/cutedsl/discrete_grouped/dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py` around lines 3166 - 3195, In the store_d_directly branch, replace the undefined thr_copy_t2r reference used to partition gD_sub_loop with a thread slice derived from the existing tiled_copy_t2r binding. Preserve the current tCgD1/tCgD2 indexing and global-store behavior, ensuring the branch traces successfully when store_d_directly is enabled.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py (1)
3187-3196: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Direct-store epilogue references an unbound
thr_copy_t2rin both block-scaled dGLU kernels.epilog_tmem_copy_and_partitionreturns the tiled copy (tiled_copy_t2r), not a per-thread slice, sothr_copy_t2ronly exists as a local inside that helper. Both direct-store branches callthr_copy_t2r.partition_D(...)and would raiseNameErrorat trace time; they are currently dead becauseself.store_d_directlyis hard-codedFalse.moe_grouped_gemm_dglu_dbias.pyline 1870 shows the correct pattern.
python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py#L3187-L3196: addthr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx)before thepartition_Dcall on line 3195.python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py#L3604-L3615: add the sametiled_copy_t2r.get_slice(epi_tidx)binding before line 3614.🧰 Tools
🪛 Ruff (0.16.0)
[error] 3195-3195: Undefined name
thr_copy_t2r(F821)
📍 Affects 2 files
python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py#L3187-L3196(this comment)python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py#L3604-L3615🤖 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/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py` around lines 3187 - 3196, Bind the per-thread copy slice before partitioning D in both direct-store epilogues: in python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py lines 3187-3196 and python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py lines 3604-3615, create thr_copy_t2r from tiled_copy_t2r.get_slice(epi_tidx) before calling partition_D; both sites require the same change.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py (1)
2027-2046: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Compare dgeglu signatures and their call sites across the dGLU kernel family. fd -t f -g '*dglu*.py' python/cudnn/gemm | while read -r f; do echo "=== $f" rg -nP -A8 'def dgeglu\s*\(' "$f" rg -nP -B2 -A6 'self\.dgeglu\(' "$f" doneRepository: NVIDIA/cudnn-frontend
Length of output: 4326
🏁 Script executed:
#!/bin/bash # Inspect the Rubin activation-gradient implementations, their caller, and the # parameter definitions to determine whether alpha/beta scaling is required. sed -n '1800,2190p' python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py sed -n '3380,3475p' python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py rg -n -C 4 'beta_val|square_alpha|alpha|act_func|dgeglu|dswiglu' \ python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py \ python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash python3 - <<'PY' from pathlib import Path import re rubin = Path("python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py").read_text() sibling = Path("python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py").read_text() def body(text, name): m = re.search(rf" def {name}\s*\((.*?)\n `@cute`\.jit", text, re.S) if not m: raise SystemExit(f"missing {name}") return m.group(1) rubin_dgeglu = body(rubin, "dgeglu") rubin_dswiglu = body(rubin, "dswiglu") sibling_dgeglu = body(sibling, "dgeglu") caller = re.search( r"elif cutlass\.const_expr\(self\.act_func == \"dgeglu\"\):\s*" r"d1_vec, d2_vec, dprob_swiglu = self\.dgeglu\((.*?)\)", rubin, re.S, ).group(1) print("Rubin dgeglu parameters:", [x.strip() for x in rubin_dgeglu.splitlines() if ":" in x]) print("Rubin dgeglu mentions scaling:", {name: name in rubin_dgeglu for name in ("beta_val", "square_alpha")}) print("Rubin dgeglu caller arguments:", " ".join(caller.split())) print("Caller passes scaling values:", {name: name in caller for name in ("beta_val", "square_alpha")}) print("Rubin dswiglu scaling operations:", { "square_alpha": "square_alpha" in rubin_dswiglu, "beta_val": "beta_val" in rubin_dswiglu, }) print("Sibling dgeglu accepts scaling:", { "beta_val": "beta_val" in sibling_dgeglu, "square_alpha": "square_alpha" in sibling_dgeglu, }) print("Rubin obtains per-expert scaling:", bool(re.search(r"alpha_val\s*=\s*alpha\[expert_idx\]", rubin)), bool(re.search(r"beta_val\s*=\s*beta\[expert_idx\]", rubin))) PYRepository: NVIDIA/cudnn-frontend
Length of output: 174
🏁 Script executed:
#!/bin/bash python3 - <<'PY' from pathlib import Path import re path = Path("python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py") text = path.read_text() def method(text, name, end): start = text.index(f" def {name}(") finish = text.index(end, start) return text[start:finish] dgeglu = method(text, "dgeglu", "\n\n # GPU device kernel") dswiglu = method(text, "dswiglu", "\n `@cute.jit`\n def dgeglu") caller = re.search( r'elif cutlass\.const_expr\(self\.act_func == "dgeglu"\):\s*' r'd1_vec, d2_vec, dprob_swiglu = self\.dgeglu\(([^)]*)\)', text, ).group(1) print("dgeglu has beta_val:", "beta_val" in dgeglu) print("dgeglu has square_alpha:", "square_alpha" in dgeglu) print("dgeglu caller:", " ".join(caller.split())) print("caller has beta_val:", "beta_val" in caller) print("caller has square_alpha:", "square_alpha" in caller) print("dswiglu scales acc_vec by square_alpha:", bool(re.search(r"acc_vec.*square_alpha", dswiglu))) print("dswiglu scales C operands by beta_val:", dswiglu.count("beta_val") >= 2) print("per-expert alpha/beta loaded:", all( re.search(pattern, text) for pattern in (r"alpha_val\s*=\s*alpha\[expert_idx\]", r"beta_val\s*=\s*beta\[expert_idx\]") )) PYRepository: NVIDIA/cudnn-frontend
Length of output: 486
Thread per-expert scaling through
dgeglu.
alpha_valandbeta_valare loaded per expert, but Rubin’sdgegluneither accepts nor appliessquare_alpha/beta_val; it uses raw accumulator and C operands. Pass and apply these values as the sibling implementations do, otherwiseact_func="dgeglu"produces incorrect gradients whenever either scale differs from 1.0.🧰 Tools
🪛 OpenGrep (1.26.0)
[ERROR] 2036-2036: 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/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py` around lines 2027 - 2046, Update the Rubin dgeglu flow around the dgeglu method to accept the per-expert square_alpha and beta_val values, then apply them to the accumulator and C operands before computing gradients. Thread both values from their existing alpha_val/beta_val loads at the call site, matching sibling implementations, while preserving the current behavior when scales are 1.0.python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py (1)
810-837: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
amax_tensor_bufis honored on cache hits but dropped on the first call.Line 813 applies
amax_tensor_bufunconditionally, while the cache-miss branch only assignsamax_tensorinside thed_dtype in [bfloat16, float16]guard (lines 834-837). With an fp8d_dtypeplus a caller-suppliedamax_tensor_buf, the first invocation runs withamax_tensor=Noneand every subsequent (cached) invocation runs with the caller's buffer — silently inconsistent results for the first step.🐛 Proposed fix
cached_amax_tensor = None - amax_tensor = None + amax_tensor = amax_tensor_buf if d_dtype in [torch.bfloat16, torch.float16]: _logger.debug("grouped_gemm_dswiglu_wrapper_sm100: Detected bf16/float16 d_dtype, constructing amax_tensor") cached_amax_tensor = torch.empty((l, 2, 1), dtype=torch.float32, device=a_tensor.device) amax_tensor = amax_tensor_buf if amax_tensor_buf is not None else cached_amax_tensor📝 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 cache_key in _cache_of_GroupedGemmDswigluSm100Objects: _logger.debug("group_gemm_dswiglu_wrapper_sm100: Using previously cached GroupedGemmDswigluSm100 object") grouped_gemm_dswiglu, cached_amax_tensor, cached_beta_tensor = _cache_of_GroupedGemmDswigluSm100Objects[cache_key] amax_tensor = amax_tensor_buf if amax_tensor_buf is not None else cached_amax_tensor if beta_tensor is not None: effective_beta = beta_tensor elif cached_beta_tensor is not None: effective_beta = cached_beta_tensor else: # Fallback: cache was populated without beta caching (non-NVFP4 path), # but caller now passes None (NVFP4 path). Create ones tensor on-the-fly. effective_beta = torch.ones(l, dtype=torch.float32, device=a_tensor.device) else: _logger.debug( "group_gemm_dswiglu_wrapper_sm100: No previously cached GroupedGemmDswigluSm100 object found, creating new GroupedGemmDswigluSm100 object" ) # For NVFP4 (beta_tensor=None): create and cache a ones tensor — avoids FillFunctor on every step. # For non-NVFP4 (beta_tensor provided): use caller's value directly; don't cache (it changes each step). cached_beta_tensor = torch.ones(l, dtype=torch.float32, device=a_tensor.device) if beta_tensor is None else None effective_beta = cached_beta_tensor if beta_tensor is None else beta_tensor cached_amax_tensor = None amax_tensor = amax_tensor_buf if d_dtype in [torch.bfloat16, torch.float16]: _logger.debug("grouped_gemm_dswiglu_wrapper_sm100: Detected bf16/float16 d_dtype, constructing amax_tensor") cached_amax_tensor = torch.empty((l, 2, 1), dtype=torch.float32, device=a_tensor.device) amax_tensor = amax_tensor_buf if amax_tensor_buf is not None else cached_amax_tensor🤖 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/dswiglu/api.py` around lines 810 - 837, Update the cache-miss initialization in the grouped GEMM wrapper so a caller-provided amax_tensor_buf is assigned to amax_tensor regardless of d_dtype, including fp8 paths. Preserve allocation of cached_amax_tensor only for bfloat16/float16 when no caller buffer is supplied, and ensure the first invocation matches cache-hit behavior.python/cudnn/gemm/cutedsl/grouped/dswiglu/grouped_gemm_dswiglu_quant.py (1)
601-602: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Prefer
not self.generate_sfd.Ruff E712 flags the equality comparison against
False.🧹 Proposed fix
- if cutlass.const_expr(self.generate_sfd == False): + if cutlass.const_expr(not self.generate_sfd):📝 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 cutlass.const_expr(not self.generate_sfd): self.discrete_col_sfd = False🧰 Tools
🪛 Ruff (0.16.0)
[error] 601-601: Avoid equality comparisons to
False; usenot self.generate_sfd:for false checksReplace 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/gemm/cutedsl/grouped/dswiglu/grouped_gemm_dswiglu_quant.py` around lines 601 - 602, Update the condition in the generate_sfd handling block to use the idiomatic negation of self.generate_sfd instead of comparing it explicitly to False, preserving the existing cutlass.const_expr behavior.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py (1)
141-154: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Find m_aligned values passed alongside a 512 MMA tile M for the Rubin GLU kernel. rg -nP -C4 'mma_tiler_mn\s*=\s*\(\s*512' rg -nP -C3 '\bm_aligned\s*=' --type=py -g '!**/moe_blockscaled_grouped_gemm_glu_rubin.py'Repository: NVIDIA/cudnn-frontend
Length of output: 159
🏁 Script executed:
#!/bin/bash set -eu file=$(fd -t f 'moe_blockscaled_grouped_gemm_glu_rubin\.py$' | head -n1) printf '%s\n' "FILE=$file" wc -l "$file" ast-grep outline "$file" --view compact printf '\n-- can_implement and constructor --\n' sed -n '105,190p' "$file" sed -n '220,260p' "$file" printf '\n-- relevant callers and symbols --\n' rg -n -C5 'm_aligned|can_implement|mma_tiler_mn|enable_breuse' "$file" python test docs 2>/dev/null | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 480
🏁 Script executed:
#!/bin/bash set -eu file=python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py printf '%s\n' '-- can_implement and constructor --' sed -n '105,190p' "$file" sed -n '220,260p' "$file" printf '%s\n' '-- relevant callers and symbols --' rg -n -C5 'm_aligned|can_implement|mma_tiler_mn|enable_breuse' "$file" python test docs 2>/dev/null | head -n 700Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash set -eu printf '%s\n' '-- helper implementation --' helper=$(fd -t f 'moe_kernel_helpers\.py$' | head -n1) printf 'FILE=%s\n' "$helper" rg -n -C12 'def can_implement|m_aligned|fix_pad_size' "$helper" | head -n 220 printf '%s\n' '-- callers of this kernel can_implement --' rg -n -C8 'BlockScaledMoEGroupedGemmGluKernel\.can_implement|moe_blockscaled_grouped_gemm_glu_rubin|m_aligned\s*=' \ python test --glob '*.py' --glob '!python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py' \ | head -n 500 printf '%s\n' '-- scheduler/mask use of alignment --' rg -n -C8 'create_mask|padded_offsets|m_aligned|FIX_PAD_SIZE' \ python/cudnn/gemm/cutedsl/grouped/glu --glob '*.py' \ | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash set -eu printf '%s\n' '-- Rubin GLU API and all references --' rg -l 'moe_blockscaled_grouped_gemm_glu_rubin|BlockScaledMoEGroupedGemmGluKernel' python test printf '%s\n' '-- exact m_aligned/tiler configuration references --' rg -n -C5 'moe_blockscaled_grouped_gemm_glu_rubin|BlockScaledMoEGroupedGemmGluKernel|mma_tiler_mn.*512|m_aligned.*512|m_aligned.*256' \ python/cudnn test/python --glob '*.py' | head -n 800 printf '%s\n' '-- scheduler mask implementation in Rubin kernel --' rg -n -C15 'create_mask|padded_offsets|m_aligned' \ python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py \ | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash set -eu api=python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py kernel=python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py dglu=python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py printf '%s\n' '-- GLU API kernel selection and support call --' rg -n -C20 'm_aligned|_kernel|can_implement|check_support|mma_tiler_mn' "$api" | head -n 450 printf '%s\n' '-- DGLU Rubin B-reuse branch --' sed -n '145,190p' "$dglu" sed -n '265,295p' "$dglu" printf '%s\n' '-- GLU scheduler mask and per-expert tile coordinate logic --' rg -n -C12 'create_mask|padded_offsets|mPosition_base|mma_tile_coord_m|tile.*expert|expert.*tile' "$kernel" | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash set -eu api=python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py kernel=python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py printf '%s\n' '-- Rubin selection and check_support control flow --' rg -n -C12 '_is_rubin_kernel|def check_support|can_implement\(' "$api" | head -n 400 printf '%s\n' '-- all direct calls to the Rubin GLU can_implement --' rg -n -C8 'can_implement\(' python/cudnn test/python --glob '*.py' \ | rg -B8 -A8 'Glu|GLU|rubin|GroupedGemm' | head -n 500 printf '%s\n' '-- exact 512 branch comparison --' sed -n '145,185p' python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py sed -n '135,160p' "$kernel"Repository: NVIDIA/cudnn-frontend
Length of output: 28248
Reject non-tile-aligned padding in the B-reuse path.
For
mma_tiler_mn[0] == 512,can_implementacceptsm_aligned=256, despite the constructor requiring B-reuse callers to align M to the 512-row tile. This can let per-expert tiles cross padded segment boundaries. Addif m_aligned % mma_tiler_mn[0] != 0: return False.🤖 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/glu/moe_blockscaled_grouped_gemm_glu_rubin.py` around lines 141 - 154, Update the `use_2cta_instrs` and `mma_tiler_mn[0] == 512` validation branch in `can_implement` to reject `m_aligned` values that are not divisible by the 512-row tile height. Add the alignment check before returning success, while preserving the existing dtype, layout, tensor-alignment, tile-N, and cluster-M validations.python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py (1)
1485-1490: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Duplicated assignment, and
global_scale_bis dereferenced under aglobal_scale_aguard.Line 1489 repeats line 1488 verbatim. Also, the branch only tests
global_scale_a is not Nonebut unconditionally readsglobal_scale_b.iterator, so ana-only call path would fault; make the guard explicit while the file is being relocated.🧹 Proposed cleanup
- if cutlass.const_expr(global_scale_a is not None): + if cutlass.const_expr(global_scale_a is not None and global_scale_b is not None): expert_idx = work_tile_info.expert_idx current_scale_a_iter = global_scale_a.iterator + expert_idx current_scale_b_iter = global_scale_b.iterator + expert_idx - current_scale_b_iter = global_scale_b.iterator + expert_idx alpha = cute.arch.load(current_scale_a_iter.llvm_ptr, cutlass.Float32) * cute.arch.load(current_scale_b_iter.llvm_ptr, cutlass.Float32)📝 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 cutlass.const_expr(global_scale_a is not None and global_scale_b is not None): expert_idx = work_tile_info.expert_idx current_scale_a_iter = global_scale_a.iterator + expert_idx current_scale_b_iter = global_scale_b.iterator + expert_idx alpha = cute.arch.load(current_scale_a_iter.llvm_ptr, cutlass.Float32) * cute.arch.load(current_scale_b_iter.llvm_ptr, cutlass.Float32)🤖 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/wgrad/moe_blockscaled_grouped_gemm_wgrad.py` around lines 1485 - 1490, Remove the duplicated current_scale_b_iter assignment in the block handling global scales. Update the guard around current_scale_a_iter/current_scale_b_iter and alpha calculation to require both global_scale_a and global_scale_b to be present before dereferencing either iterator, preserving safe behavior for an a-only call path.python/cudnn/gemm/ops/moe_grouped_matmul.py (1)
83-89: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Shared per-device handle +
set_streamis not thread-safe.
_get_handlemutates the stream on a handle shared by every caller on that device, and the caller then executes on it (line 345). Two threads submitting on different streams can interleave betweenset_streamandgraph.execute, so work lands on the wrong stream. Handle creation itself also races on the_cudnn_handlesdict.🔒️ Suggested fix
+import threading + +_handles_lock = threading.Lock() + def _get_handle(device: torch.device): """Return a lazily-initialised cuDNN handle with the current CUDA stream.""" - if device not in _cudnn_handles: - _cudnn_handles[device] = cudnn.create_handle() - stream = torch.cuda.current_stream(device).cuda_stream - cudnn.set_stream(handle=_cudnn_handles[device], stream=stream) - return _cudnn_handles[device] + with _handles_lock: + key = (device, threading.get_ident()) + if key not in _cudnn_handles: + _cudnn_handles[key] = cudnn.create_handle() + handle = _cudnn_handles[key] + cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) + return handlePer-thread (or per-stream) handles avoid the stream mix-up; a lock alone only fixes the dict race unless it also covers
execute.📝 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.import threading _handles_lock = threading.Lock() def _get_handle(device: torch.device): """Return a lazily-initialised cuDNN handle with the current CUDA stream.""" with _handles_lock: key = (device, threading.get_ident()) if key not in _cudnn_handles: _cudnn_handles[key] = cudnn.create_handle() handle = _cudnn_handles[key] cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle🤖 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/ops/moe_grouped_matmul.py` around lines 83 - 89, Update _get_handle and its callers to avoid sharing a mutable cuDNN handle across concurrent streams: use a handle scoped per thread or per CUDA stream, and ensure handle creation is synchronized as needed. Preserve the existing current-stream selection and pass the isolated handle through the execution path around graph.execute so concurrent submissions cannot interleave set_stream with execution.
…ss-dsl < 4.8 (#662) * test: fix stale quant kernel paths in Rubin dispatch test The Gemm fusion reorganization (#459) moved the grouped quant kernels to python/cudnn/gemm/cutedsl/grouped/quant/, but test_grouped_gemm_quant_kernels_support_optional_prob still looked for them under grouped_gemm_quant/, failing with FileNotFoundError in oss_tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: skip glu_hadamard_quant tests on cutlass-dsl < 4.8 The glu_hadamard_quant kernel references cutlass.FloatNV8E5M3FNU unconditionally at compile time, and that dtype only exists in cutlass-dsl >= 4.8, so every test in the file failed with an AttributeError on older builds (13 failures on the 4.5.1 CI lane) even when the scale-factor dtype under test is e4m3/e8m0. Gate the module with the same hasattr check _skip_unless_e5m3_supported already uses. Verified on an SM100 box: cutlass-dsl 4.7.0 without the gate reproduces the AttributeError, with the gate all 23 tests skip; 4.8.0a0 runs 17 passed / 6 skipped (the pre-existing Rubin-only e5m3 skips). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2 participants
Summary
All GEMM fusion packages were scattered across the top of
python/cudnn/with redundant, inconsistent prefixes (gemm_swiglu/,grouped_gemm/grouped_gemm_glu/,discrete_grouped_gemm/discrete_grouped_gemm_swiglu/). This PR collects them into a single family tree grouped by operand layout, mirroring the internalfrost_develreorganization:Per-fusion directories drop the now-redundant prefix, so
grouped_gemm/grouped_gemm_glu/becomesgemm/cutedsl/grouped/glu/.This is a pure move — 100 of the 115 changed files are byte-identical renames.
No public API changes
Every symbol keeps its name and stays reachable as
cudnn.<symbol>._LAZY_OPTIONAL_IMPORTSinpython/cudnn/__init__.pyis repointed at the new module paths, so the supported top-level entry point is unchanged.cudnn.grouped_gemmandcudnn.discrete_grouped_gemmstill resolve as attributes through the same table.Two adjacent modules move with the family so the tree is complete, each with a shim so existing imports keep working:
engines/reference_matmul_engine.pygemm/reference/cudnn.engines, socudnn.engines.ReferenceMatmulEngineis unchangedexperimental/ops/moe_grouped_matmul.pygemm/ops/sys.modules, so both the attribute andimport cudnn.experimental.ops.moe_grouped_matmulkeep resolvingGitHub-only APIs carried over
The internal reorg predates several APIs that landed here first. These are carried over in place rather than dropped:
GroupedGemmSm100/grouped_gemm_wrapper_sm100) →gemm/cutedsl/grouped/unfused/GemmProjRopeMxfp8Bf16InSm100,GemmProjRopeMxfp8Mxfp8InSm100)_bf16_api.py/_blockscaled_api.pysplit indglu,glu,quant,wgrad,unfusedNaming note
grouped_gemm/grouped_gemm_utils.py→gemm/cutedsl/grouped/backend_utils.py(backend enum + stream context), keeping it distinct from the kernel-sideutils.pythat sits alongside it. This is the one file that gets a new name rather than just a new path.Drive-by fix
The moved reference engine carried relative imports (
.base,.engine_ids,..graph_types) that only worked fromcudnn/engines/; these are now absolute. ItsTYPE_CHECKING-only..pygraphimport was already pointing at a non-existentcudnn.pygraphmodule before this PR — it now points at the realcudnn._pygraph. (The identical latent import inengines/base.pyis left alone as out of scope.)Verification
python -m compileallclean overpython/cudnnandtest/python.cudnnimports plus all 53_LAZY_OPTIONAL_IMPORTSentries against the file tree. Result: 9 unresolved ondevelop→ 8 on this branch, i.e. zero new breakage and one pre-existing latent import fixed. The 8 remaining are pre-existing and untouched by this PR (BSA namespace-package dir,engines/base.pyTYPE_CHECKING, and BSA/CSA nested re-export chains the static checker can't follow).black --line-length 160clean on every file this PR edits. Eight files still report as unformatted, but all eight are already unformatted ondevelop— reformatting them would have polluted an otherwise pure-move diff.torchinstalled), so CI is the real gate on kernel behavior. Since no kernel logic is touched, the risk surface is entirely module-path resolution, which is what the static pass above covers.Docs, README,
ACKNOWLEDGEMENTS.md,THIRD_PARTY_LICENSES.txt, and thecutedsl-kernel-integrationskill are updated to the new paths.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Chores