[1.28.0-rc] Gate the packed-FP4 wgrad layout workaround on cutlass-dsl < 4.8 - #764
Merged
Anerudhan merged 3 commits intoAug 26, 2026
Merged
Conversation
Public cutlass-dsl wheels before 4.8 interpret packed sub-byte from_dlpack layouts in byte units, so BlockScaledMoEGroupedGemmWgradKernel recasts the FP4 A/B layouts to element units. The 4.8.0a0 public wheels adopted the internal wheel's native sub-byte layout semantics, so the recast now double-corrects: the MMA consumes byte-aliased data and every fp4 wgrad test fails with ~94% mismatched output (and occasionally an illegal memory access from the corrupted discrete-pointer TMA path). Gate the workaround on the cutlass-dsl version instead of only on internal-wheel presence. Verified on sm100 (torch 2.13, dense compile_execute fp4 sf_e4m3, mma 128x128/256x128, cluster 1x1/2x1): bit-exact vs the torch reference on both cutlass-dsl 4.7.0 (gate on) and 4.8.0a0+20260823210556.ac70faa (gate off); before this change 4.8 fails with ~94% mismatch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
The pre-commit GitHub Action runs black on --all-files and this pre-existing file was not black-clean, failing the check for every PR. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-pick of all three commits from #763 onto
1.28.0-rc: the gate fix, the comment trim, and the black formatting oftest_grouped_gemm_glu_hadamard_quant.py(so the pre-commit--all-filescheck passes on this branch too).Problem / root cause
BlockScaledMoEGroupedGemmWgradKernelapplies a packed-FP4from_dlpacklayout workaround (cute.recast_layout(4, 8, ...)on the A/B layouts) gated only on the internal cutlass-dsl wheel not being installed. Public cutlass-dsl wheels < 4.8 need it, but the 4.8.0a0 public wheels adopted native sub-byte layout semantics, so the recast double-corrects: every fp4 wgrad test fails with ~94% mismatched output (21 tests on the Blackwell nightlyoss:rel:cutlass_dsl_4.8shard1), occasionally escalating to an illegal memory access that cascades into unrelated test ERRORs.Fix
Gate the workaround on the cutlass-dsl version (
< 4.8) instead of only on internal-wheel presence. sm_107 and internal-wheel behavior unchanged; behavior on 4.5–4.7 public wheels unchanged (gate stays ON).Verification (done on #763, same patch)
sf_e4m3, mma 128x128/256x128, cluster 1x1/2x1: bit-exact (0/491520 mismatched) on cutlass-dsl 4.7.0 (gate ON) and 4.8.0a0+20260823210556.ac70faa (gate OFF); without the fix, 4.8 fails at ~94% mismatch.oss:relpointed at the 4.8.0a0 prerelease: all four jobs green; Blackwell shard1 went from 21 failed to 910 passed, 0 failed.🤖 Generated with Claude Code