Gate the packed-FP4 wgrad layout workaround on cutlass-dsl < 4.8

[Anerudhan]

Problem

The nightly oss:rel:cutlass_dsl_4.8 [Blackwell] CI jobs fail all fp4 test_grouped_gemm_wgrad tests (21 on shard1, e.g. pipeline 64271661 / job 409620600) with ~94% mismatched output — and, run-dependently, an illegal memory access that poisons the CUDA context and cascades into dozens of unrelated test ERRORs. cutlass-dsl 4.7 is green on the same commit.

Root cause

BlockScaledMoEGroupedGemmWgradKernel carries a packed-FP4 from_dlpack layout workaround (cute.recast_layout(4, 8, ...) on the A/B layouts) that public cutlass-dsl wheels < 4.8 require, gated on the internal cutlass-dsl wheel not being installed. The 4.8.0a0 public wheels adopted the internal wheel's native sub-byte layout semantics, so on 4.8 the recast double-corrects: the MMA consumes byte-aliased data (effectively halved-K / wrong strides), producing garbage output, and the corrupted addressing in the discrete-pointer TMA path can cross into unmapped memory (the IMA).

A/B proof matrix from the standalone repro (dense compile_execute fp4 sf_e4m3 vs the torch scaled_grouped_mm reference, sm100):

	workaround ON	workaround OFF
cutlass-dsl 4.7.0	pass, bit-exact	fail, ~94% wrong
cutlass-dsl 4.8.0a0 dailies	fail, ~94% wrong	pass, bit-exact

Fix

Gate the workaround on the cutlass-dsl version (< 4.8) instead of only on internal-wheel presence. The sm_107 and internal-wheel exemptions are unchanged.

Verification

On an sm100 box (torch 2.13, CUDA 13), dense compile_execute fp4 sf_e4m3, mma 128x128 and 256x128, cluster 1x1 and 2x1:

• cutlass-dsl 4.7.0 + this fix → gate ON → 0/491520 mismatched (bit-exact)
• cutlass-dsl 4.8.0a0+20260823210556.ac70faa (the exact daily the failing CI job installed) + this fix → gate OFF → 0/491520 mismatched (bit-exact)

A full CI replay (internal MR !2355, oss:rel pointed at the 4.8.0a0 prerelease) is running on Blackwell shard0/shard1; results will be posted as a comment here.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved compatibility for grouped GEMM operations across different CUTLASS DSL versions.
  • Newer versions now use native packed-FP4 layouts, while older versions continue using the appropriate compatibility path.
  • Preserved correct scale-factor indexing for grouped GEMM quantization workflows.

[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: 923a241a-2843-4282-b5c1-ff800651fa2b

📥 Commits

Reviewing files that changed from the base of the PR and between a8b3c7d and 521b19b.

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (1)

• python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py

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

---

📝 Walkthrough

Walkthrough

The grouped WGRAD GEMM now selects the packed-FP4 layout workaround based on the CUTLASS DSL version. Internal DSL and public versions 4.8+ use native layouts; older public versions use recasting.

Changes

FP4 layout compatibility

Layer / File(s)	Summary
Version-aware workaround selection python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py, test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py	The module parses CUTLASS DSL versions and enables recasting only for older public versions. The runtime path uses this flag while preserving the SM107 exclusion. The blocked scale-factor index expression is reformatted without changing its behavior.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 521b1

This localized change selects the packed-FP4 layout workaround by cutlass-dsl version and is reported bit-exact on both supported versions; no actionable merge-blocking risk remains after normal checks and review.

Suggested reviewers: timmoon10

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description clearly explains the problem, root cause, fix, and verification results, but it does not follow the repository template. It omits the required checklist, affected area, explicit summar…	Update the description to include all template sections. Mark or explain the Before submitting checklist, select the affected area, add explicit Summary and Why sections, state Related issues or None, document API and compatibility impact o…
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 and concisely describes the main change: applying the packed-FP4 wgrad layout workaround only for cutlass-dsl versions below 4.8.
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 problem, root cause, fix, and verification results, but it does not follow the repository template. It omits the required checklist, affected area, explicit summary, related issues, API and compatibility impact, and exact testing commands.

Resolution

Update the description to include all template sections. Mark or explain the Before submitting checklist, select the affected area, add explicit Summary and Why sections, state Related issues or None, document API and compatibility impact or None, and list the exact testing commands with their results.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[Anerudhan]

CI replay results (internal MR !2355, pipeline 64703783, oss:rel pointed at the 4.8.0a0 prerelease with this fix applied):

Job	Result
oss:rel [Blackwell, shard1] (the failing configuration)	✅ 910 passed, 682 skipped, 0 failed — all 21 previously-failing test_grouped_gemm_wgrad fp4 sf_e4m3 tests pass
oss:rel [Blackwell, shard0]	✅ success
oss:rel [Hopper]	✅ success
oss:rel [Ampere]	✅ success

The replay installed cutlass-dsl 4.8.0a0+20260825210638.952ea40 (newest daily; the version spec floats), and the local sm100 verification in the PR description pinned 4.8.0a0+20260823210556.ac70faa — the exact daily from the failing nightly job 409620600 — so both dailies are covered, plus 4.7.0 for the gate-ON side.

[Anerudhan]

@cudnn-ci-bot run oss

[cudnn-ci-bot]

🏁 Pipeline finished

SHA: 521b19b
Targets: oss
Branch: cudnn-gh/pr-763-521b19b
Pipeline: 64709442
Last updated: 2026-08-26 19:39 UTC