Fix Rubin WGRAD compatibility across Blackwell CUTLASS wheels

[WanZzzzzz]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.

Affected area

Summary

Why

Related issues

Related to #460

API and compatibility impact

Testing

Summary by CodeRabbit

• Bug Fixes

  • Improved grouped GEMM gradient handling across supported GPU architectures.
  • Corrected FP4 layout processing for applicable public configurations.
  • Improved handling of epilogue metadata and helper-kernel execution on Blackwell and Rubin platforms.

• Tests

  • Updated grouped GEMM test coverage for the revised repository layout.
  • Added appropriate availability handling for Rubin-specific validation tests.

[coderabbitai]

📝 Walkthrough

Walkthrough

The grouped WGRAD kernel now detects the internal CUTLASS DSL, narrows FP4 layout recasting, routes C TMA metadata by GPU configuration, and accepts optional helper-kernel metadata. Rubin dispatch tests use the relocated path and skip unavailable Rubin helpers.

Changes

Grouped WGRAD dispatch

Layer / File(s)	Summary
DSL detection and FP4 layout selection python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py	The kernel detects the internal CUTLASS DSL. The FP4 layout workaround excludes the internal DSL and preserves the SM107 exception.
Host and helper-kernel TMA routing python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py	The host passes C TMA and epilogue metadata for discrete non-SM107 configurations. Rubin and SM107 rebuild discrete metadata inside helper_kernel.
Rubin dispatch test alignment test/python/fe_api/test_rubin_kernel_dispatch.py	The test uses the relocated grouped GEMM path and skips Rubin TMEM planning checks when Rubin helpers are unavailable.

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

Sequence Diagram(s)

sequenceDiagram
  participant HostDispatch
  participant helper_kernel
  participant RubinSM107
  HostDispatch->>HostDispatch: Select DSL-specific layouts and TMA metadata
  HostDispatch->>helper_kernel: Pass optional C TMA and epilogue metadata
  helper_kernel->>RubinSM107: Rebuild metadata for Rubin or SM107
  helper_kernel-->>HostDispatch: Use host metadata for other configurations

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#456: Refines the same WGRAD helper-kernel, FP4 layout, and Rubin/Blackwell TMA handling.
• NVIDIA/cudnn-frontend#459: Relocates and modifies the same grouped WGRAD kernel and helper_kernel.

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description includes the license checks and related issue, but leaves the affected area, summary, rationale, compatibility impact, and testing sections incomplete.	Complete the affected area, Summary, Why, API and compatibility impact, and Testing sections with implementation details and exact test commands and results.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title clearly identifies the Rubin WGRAD compatibility fix across Blackwell CUTLASS wheels.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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