fix: acquire SM100 grouped-wgrad tensor maps before TMA use

[JF-D]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.
• I reviewed the Hard Rules in the AGENTS.md for each directory this PR touches (see root AGENTS.md § Reviewing a PR) and my changes comply, or I explain the exception below.
• I added GitHub labels: one cat-*, one or more area:* / op:*, and one orig-* (see label list).

Affected area

Summary

Fix incorrect SM100 grouped-wgrad results during CUDA graph replay when consecutive calls reuse a tensor-map workspace with different input, scale, or output addresses. The same sequence passes in eager mode on the tested GB300 stack.

Acquire device-written descriptors in the consuming warp when switching experts. Reuse acquires across consecutive tiles to avoid per-tile overhead. Add an eager/graph regression test that changes tensor addresses while reusing a cached plan.

Why

Rewriting a descriptor in global memory does not by itself invalidate stale descriptor state used by TMA. The descriptor-update kernel writes through the generic memory proxy, but the GEMM lacks the required tensor-map acquire, so it can keep using the previous mapping.

Eager launch-time cache invalidation appears to mask this omission on the tested stack. Inserting an ordinary launch of an empty kernel makes the failing sequence correct; inserting the same kernel as a graph node does not. This isolates a launch-path effect, although we have not directly observed the driver's invalidation commands. Graph execution must not rely on that incidental invalidation between nodes.

Add cpasync.fence_tma_desc_acquire() before TMA operations to explicitly make the updated descriptors visible. Acquire once per expert transition rather than every tile to limit overhead.

API and compatibility impact

Only the kernel fix and regression test are included. No public API, workspace-size, or kernel-launch changes.
And slightly perf regression due to additional acquire.

Testing

Summary by CodeRabbit

• Bug Fixes

  • Improved grouped GEMM weight-gradient execution across dense and ragged inputs, including discrete and dense outputs.
  • Improved reliability when reusing operations with different tensors and output or scaling addresses.
  • Ensured correct behavior for empty experts and repeated CUDA Graph replays.

• Tests

  • Added coverage across supported GPU architectures, accumulation modes, input layouts, output formats, and CUDA Graph execution.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a0391330-777d-4a10-946d-9ebda255c127

📥 Commits

Reviewing files that changed from the base of the PR and between 7ac5354 and ba50412.

📒 Files selected for processing (2)

• python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py
• test/python/fe_api/grouped_gemm/test_wgrad_tma_descriptor_reuse.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

The grouped-GEMM weight-gradient kernel now fences TMA descriptor acquisition when processing new experts and reuses descriptors across consecutive tiles. A parametrized CUDA test validates eager execution and repeated CUDA graph replays across supported configurations.

Changes

Wgrad TMA descriptor reuse

Layer / File(s)	Summary
Descriptor acquisition fences python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py	The load path acquires ragged-input and scale descriptors when the expert changes. The discrete-output epilogue acquires each expert’s C descriptor before stores, including empty experts.
Descriptor reuse validation test/python/fe_api/grouped_gemm/test_wgrad_tma_descriptor_reuse.py	A parametrized SM100/SM103 test covers dense and ragged inputs, discrete and dense outputs, and accumulation modes. It validates eager execution and eight CUDA graph replays with distinct tensor addresses.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Suggested reviewers: anerudhan

Merge Risk: ⚪ Minimal · up to ba504

No actionable merge-blocking risk remains from the reviewed changes.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description clearly explains the bug, fix, rationale, compatibility impact, and regression scope. However, it does not select an affected area, does not include a related issue entry, and does not…	Select the applicable affected area, add a related issue or state that none applies, and document the exact test commands with their results. Complete or explain the required GitHub label checklist item.
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main fix: acquiring SM100 grouped-wgrad tensor maps before TMA use. It is concise and specific.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description clearly explains the bug, fix, rationale, compatibility impact, and regression scope. However, it does not select an affected area, does not include a related issue entry, and does not provide exact testing commands and results. The required label checklist item is also unchecked.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

---

_{Comment @coderabbitai help to get the list of available commands.}

[Anerudhan]

@cudnn-ci-bot run oss

[cudnn-ci-bot]

🏁 Pipeline finished

SHA: ba50412
Targets: oss
Branch: cudnn-gh/pr-1126-ba50412
Pipeline: 68473550
Last updated: 2026-09-17 22:33 UTC

23 passed, 1 failed, 6 manual

manual

• manual:python_samples - Manual
• manual:python_tests - Manual
• manual:frost - Manual
• manual:oss - ✅ Passed
• manual:pycudnn - Manual
• manual:multi_gpu - Manual
• manual:backend - Manual

analysis

• analysis:cudnn_clang_disable_exception - ✅ Passed
• analysis:cudnn_v9_no_half_conversion - ✅ Passed
• analysis:cudnn_clang - ✅ Passed
• analysis:check-relative-includes - ✅ Passed
• analysis:check-CUDNN_FRONTEND_SKIP_JSON_LIB - ✅ Passed
• analysis:guardwords_scan - ✅ Passed
• analysis:jax-import-guard - ✅ Passed
• san:build - ✅ Passed

build

• analysis:api_index - ❌ Old Failure (nightly failed)
• build:dev:linux:amd64 - ✅ Passed
• build:rel:linux:amd64 - ✅ Passed
• build:dev:linux:arm64 - ✅ Passed
• build:rel:linux:arm64 - ✅ Passed
• build:rel:win:amd64 - ✅ Passed

oss_tests

• oss:cutlass-rel:sm80 - ✅ Passed
• oss:cutlass-rel:sm90 - ✅ Passed
• oss:cutlass-rel:sm100-shard0 - ✅ Passed
• oss:cutlass-rel:sm100-shard1 - ✅ Passed
• oss:cutlass-4.8:sm107 - ✅ Passed

sanitizer_tests

• san:cpp_test:sm80 - ✅ Passed
• san:cpp_test:sm90 - ✅ Passed
• san:cpp_test:sm100 - ✅ Passed

triage

• triage:ai - ✅ Passed

[YangXu1990uiuc]

Codex bot review · model: gpt-6-astra

Reviewed ba50412200dd against base d34a6909c9d1.

The fix makes sense, and I found no additional issue. The acquire is in each warp that actually issues the relevant TMA operation; reusing it across consecutive tiles is valid because the descriptor is immutable during this GEMM launch. The discrete-output path also acquires for empty experts. Existing approval is already in place, so this is a follow-up comment rather than a duplicate approval.

Validation: Independent SM100 run at ba50412: all 8 new eager/graph regression cases passed, with no skips. Running the same test against merge-base 7ac5354 reproduced a graph-replay failure in the first dense/2D/overwrite case: 65536/65536 elements differed, max absolute error 1792. The baseline run deliberately stopped at that first failure. Checked descriptor construction, consuming warps and persistent expert transitions against the PTX tensor-map proxy ordering rule: ordinary inter-kernel ordering does not substitute for the required tensor-map acquire. See the PTX memory model. Correctness-gated head/base/base/head timing on the same 68-SM SM100 GPU, CUDA 13.2 and CuTe DSL 4.8: CPU steady-state wrapper enqueue stayed about 14 us for E=4/features=256/tokens=512 and 15 us for E=8/features=1024/tokens=8192. GPU graph replay was 10.23-10.25 us before versus 10.24 us after for dense/2D, and 27.13-27.20 us before versus 28.14-28.18 us after for discrete/ragged (about 3.7%). This small measured cost is acceptable for the correctness fix. Both versions passed exact output checks, poisoned-output replay and changed-input replay in these stable-address timing cases. Reused maintainer-triggered pipeline 68473550 after verifying its mirror records the full reviewed GitHub SHA. OSS SM80/SM90/SM100/SM107 jobs, builds and C++ sanitizer jobs passed. The remaining api_index failure is the missing SdpaFwdDslSm100.template_params entry, outside this two-file change; no duplicate CI run was requested.

Limitations: Local hardware coverage is SM100; SM103 was not independently exercised. The performance check covers two FP8 block-scaled shapes and the complete captured wrapper path, not a broad tuning sweep or an isolated GEMM-only measurement. The overall CI pipeline is not fully green because of the separate api_index failure.

Comment only: this PR already has an active approval; preserving existing approvals.