fix(gemm): handle linear MXFP8 scales in bmm_mxfp8 cuDNN - #3798
coder-2011 wants to merge 8 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds keyword-only ChangesMXFP8 sf_layout / scale_reordering support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant bmm_mxfp8
participant mxfp8_gemm_sm100
participant cudnn_runner
participant graph_builder
Caller->>bmm_mxfp8: call with sf_layout
bmm_mxfp8->>bmm_mxfp8: validate and map layout
bmm_mxfp8->>mxfp8_gemm_sm100: pass scale_reordering
mxfp8_gemm_sm100->>cudnn_runner: provide reordering for tuning
cudnn_runner->>graph_builder: build graph with reordering
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for configurable scale layouts (sf_layout) in MXFP8 GEMM operations, allowing the cuDNN backend to handle both linear and 128x4 swizzled layouts, while restricting the CUTLASS backend to 128x4 swizzled layouts. Feedback suggests updating _cudnn_bmm_mxfp8_requirement to return False for unsupported layouts to ensure correct backend routing.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@flashinfer/gemm/gemm_base.py`:
- Around line 8723-8728: The bmm_mxfp8 backend gating in the cuDNN suitability
check only validates sf_layout and cuDNN availability, so add a
compute-capability guard in this path to reject non-SM10 devices. Update the
backend selection logic around the sf_layout check to inspect the current device
capability before calling _cudnn_available_or_raise_for_backend, and return
False for SM12/other unsupported devices so backend="auto" cannot fall through
to an unsupported cuDNN path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 569088b7-d049-490e-b3b1-ebcbf2345920
📒 Files selected for processing (2)
flashinfer/gemm/gemm_base.pytests/gemm/test_bmm_mxfp8.py
|
Ran this on an RTX PRO 6000 (SM120, CUDA 13, cudnn-frontend 1.25.0 / cudnn 9.19), since the cudnn path is admitted there at runtime (_cudnn_bmm_mxfp8_requirement allows [100, 103, 110, 120, 121]) even though the test file skips cudnn for anything but SM10x. The fix checks out: with linear scales on the cudnn backend, main produces scrambled output against a reference computed from the dequantized fp8 operands (median rel err ~0.99, cos 0.37) while this branch is exact modulo output rounding (max rel err 0.4%, cos 0.999999). Swizzled is unchanged between main and this branch (cos 0.9993 vs the pre-quant reference) and agrees with the cutlass backend, so no regression there. One thing I had to work around to measure this, worth noting while you're in the test file: plain randn inputs can't see this bug -- every 32-block quantizes to nearly the same e8m0 scale, so main's broken linear path still scores cos 0.93 on my box and clears the test's 0.9 bar (with a median per-element rel error of 31%). I only got discrimination by scaling each 32-element block by a 2^-4..2^4 ladder and checking per-element error against the dequantized-operand reference. That ladder input also exposed an issue in the existing test construction that your PR inherits: the test quantizes mat2 as contiguous [b, k, n], which blocks scales along n, while the graph declares b_descale as [b, k/32, n] with K-major strides (and the bmm_mxfp8 docstring says to pass B column-major). With block-varying magnitudes even the intended swizzled config drops to cos 0.756 using the test's construction, vs 0.9993 when B is quantized as [b, n, k] (blocks along k) and passed transposed -- on both cudnn and cutlass, main and PR alike. So as written the test can't catch scale-layout bugs, which is presumably how the linear-layout one shipped silently. Happy to share the probe script if useful. |
|
Thx, I think the B-axis construction bug can be solved in a future PR, probably by me. |
|
@dhiraj113 @yzh119 @aleozlx @bkryu I woud appreciate if someone can take a look a this! |
|
let me help you escalate. adding @YangXu1990uiuc |
| _check_cudnn_override_shape_availability() | ||
| if policy is None: | ||
| policy = cudnn.build_plan_policy.HEURISTICS_CHOICE | ||
| if scale_reordering is None: |
There was a problem hiding this comment.
Can you do this through the function default argument instead.
There was a problem hiding this comment.
I might be mistaken but, cudnn.* objects are never function defaults in the codebase, bc flashinfer would crash if we dont have cudnn imported.
ex.
def build_cudnn_gemm_fp4_graph(
...,
policy=None,
):
if policy is None:
policy = cudnn.build_plan_policy.HEURISTICS_CHOICEagent said that there arent any cudnn.* function defaults in flashinfer
| block_size, | ||
| o_type, # cudnn.data_type, BF16 or FP16 | ||
| device, | ||
| scale_reordering=None, |
There was a problem hiding this comment.
Same here. Make the default value explicit in the function args.
| out: Optional[torch.Tensor] = None, | ||
| workspace_buffer: torch.Tensor = None, | ||
| tactic: int = -1, | ||
| scale_reordering=None, |
|
|
||
|
|
||
| def _cudnn_gemm_mxfp8_runner(): | ||
| def _cudnn_gemm_mxfp8_runner(scale_reordering): |
There was a problem hiding this comment.
This looks incorrect. Why do you need pass scale_reordering this way when other things are not being passed like this.
There was a problem hiding this comment.
follow the style of def mm_fp4.
| dtype: torch.dtype, | ||
| out: Optional[torch.Tensor] = None, | ||
| backend: Literal["cudnn", "cutlass", "auto"] = "auto", | ||
| *, |
There was a problem hiding this comment.
API shouldn't be extended like this.
|
Also, the test in #3663 doesn't pass with your change. |
|
Here's the code review from Claude. See if it helps - Code Review: PR #3798 — fix(gemm): handle linear MXFP8 scales in bmm_mxfp8 cuDNN I set up git (added coder-2011 remote, checked out fix-bmm-mxfp8-scale-reordering) and verified findings by running on the branch. Findings ranked most-severe first:
ValueError: cuDNN mxfp8 GEMM requires B to be column-major [batch, k, n]
Bottom line: The fix's logic is sound and works for the aligned shapes the test uses, but (1) the shipped test still fails on the issue's own hardware because the unrelated mat2 row-major bug isn't fixed, and (2) the linear path is only correct for 128x4-aligned shapes — a real latent correctness bug the test can't see. I'd ask the author to fix the test's B construction (so the test actually exercises and validates the fix) and either generalize the descriptor for linear layout or explicitly guard/document the alignment restriction. The branch is checked out locally if you want to iterate on it. |
|
Tests do pass on my end, on cuDNN 9.10.2. There was a failure bc of other unrelated issues. Claude explicitly said:
My tests say the same thing. I can fix this in a follow up, or now if u think thats better. I appreciate the thorough review, I learned alot! lmk if I made any stupid mistakes, and I will fix them. I would love to contribute to flashinfer alot more in the future |
9e98766 to
3737f1b
Compare
5 participants
Description
Fixes
bmm_mxfp8cuDNN scale-layout handling for MXFP8 linear scale tensors.Previously,
bmm_mxfp8always built the cuDNN MXFP8 graph with scale tensors declared asF8_128x4. That was correct for swizzled scales, but incorrect when the inputs were quantized withis_sf_swizzled_layout=False, which produces linear scales.This PR adds
sf_layoutargument tobmm_mxfp8and maps it to the cuDNN scale tensor reordering, whereSfLayout.layout_linear->cudnn.tensor_reordering.NONESfLayout.layout_128x4->cudnn.tensor_reordering.F8_128x4The existing default remains
SfLayout.layout_128x4to preserve current behavior. CUTLASS BMM continues to require 128x4 swizzled scales.The BMM MXFP8 test now passes the scale layout used by the test into
bmm_mxfp8, while leaving quantization behavior unchanged.Related Issues
Fixes #3663
Pull Request Checklist
Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.Tests
Tested with:
Random-seed stress for the original failing cuDNN linear-scale shape:
Summary by CodeRabbit
Summary
bmm_mxfp8with a new keyword-onlysf_layoutargument to choose the scale format.sf_layoutoptions across backends.sf_layoutvalues.