fix(quantization): nvfp4_quantize(backend='cuda') silently corrupts scale factors when global_scale is not float32 - #3497
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughNormalize ChangesFP4 Global Scale Dtype Fix
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 addresses issue #3398 by normalizing the global_scale tensor to float32 in the Python wrapper before passing it to the CUDA kernel, preventing byte-wise misinterpretation of bf16 or fp16 scales. It also adds a C++ type check for globalScale and introduces a regression test covering different scale dtypes. The reviewer noted that the normalization logic is bypassed for the cute-dsl backend and recommended moving it to the beginning of the function, while also ensuring that global_scale is cast to the correct device to prevent device mismatch errors.
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.
| if global_scale is not None and global_scale.dtype != torch.float32: | ||
| global_scale = global_scale.to(torch.float32) |
There was a problem hiding this comment.
The normalization of global_scale to float32 is currently placed after the cute-dsl backend early return (line 842). This means if a user calls fp4_quantize with backend="cute-dsl" and a non-float32 global_scale, it will bypass this normalization.
Additionally, we should also ensure global_scale is on the same device as input (i.e., global_scale = global_scale.to(input.device)) to prevent device mismatch errors or host-to-device copy issues during kernel execution.
To address both issues, we should update this block to handle both device and dtype normalization, and ideally move it to the very beginning of the fp4_quantize function (e.g., right after the sf_vec_size check) so that both backends benefit from it.
if global_scale is not None:
if global_scale.device != input.device:
global_scale = global_scale.to(input.device)
if global_scale.dtype != torch.float32:
global_scale = global_scale.to(torch.float32)There was a problem hiding this comment.
Good catches on both — agreed, fixed in the latest commit.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
flashinfer/quantization/fp4_quantization.py (1)
842-850: ⚡ Quick winConsider adding global_scale dtype normalization in the cute-dsl path for consistency.
The CUDA backend path normalizes
global_scaleto float32 at lines 862-865, but the cute-dsl dispatch in_fp4_quantize_cute_dsldoesn't apply the same conversion before invoking the cute-dsl kernel. While the cute-dsl kernel may handle non-float32 scales correctly, applying the same normalization here would ensure consistent behavior and prevent potential future issues.♻️ Proposed fix
def _fp4_quantize_cute_dsl( input: torch.Tensor, global_scale: Optional[torch.Tensor], sf_vec_size: int, sf_use_ue8m0: bool, is_sf_swizzled_layout: bool, is_sf_8x4_layout: bool, enable_pdl: Optional[bool], ) -> Tuple[torch.Tensor, torch.Tensor]: """CuTe-DSL dispatch for fp4_quantize. Maps parameters to the appropriate kernel.""" from ..cute_dsl import is_cute_dsl_available if not is_cute_dsl_available(): raise RuntimeError( "CuTe-DSL backend requested but CuTe-DSL is not available. " "Please install the required dependencies." ) + + # Normalize global_scale dtype to match CUDA backend behavior + if global_scale is not None and global_scale.dtype != torch.float32: + global_scale = global_scale.to(torch.float32) if sf_vec_size == 16 and not sf_use_ue8m0:🤖 Prompt for 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. In `@flashinfer/quantization/fp4_quantization.py` around lines 842 - 850, The cute-dsl dispatch path is missing normalization of the global_scale dtype; before calling _fp4_quantize_cute_dsl convert/cast the global_scale tensor to float32 (matching the CUDA branch's float32 normalization) so both backends receive the same dtype; update the call site that currently passes global_scale to first ensure global_scale = global_scale.to(torch.float32) (or equivalent) and then pass it into _fp4_quantize_cute_dsl.tests/utils/test_fp4_quantize.py (1)
1506-1560: 💤 Low valueWell-designed regression test with comprehensive failure-mode coverage.
The test correctly validates both symptoms of issue
#3398:
- Scale factors becoming all-zero (line 1539-1542)
- Incorrect magnitude scaling via median ratio check (line 1554-1559)
The median-based magnitude validation is particularly good—it catches under-scaling that a direction-only (cosine similarity) check would miss.
Optional enhancement: Consider adding a test case that directly exercises
fp4_quantizerather than going throughnvfp4_quantize, since the fix is infp4_quantize. This would more precisely test the layer where the fix lives and provide better isolation if the fix is refactored in the future.🤖 Prompt for 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. In `@tests/utils/test_fp4_quantize.py` around lines 1506 - 1560, Add a sibling regression test that calls fp4_quantize directly (instead of nvfp4_quantize) using the same parameterization (m and scale_dtype) and the same input setup (x with dtype bfloat16, global_scale built in scale_dtype) found in test_nvfp4_quantize_global_scale_dtype_regression; reproduce the two assertions: (a) that the scale-factor bytes (from the sf output of fp4_quantize) are not all zero and (b) that the dequantized magnitude (using e2m1_and_ufp8sf_scale_to_float or the appropriate fp4 dequant helper) yields a median |deq/x| between 0.5 and 2.0. Use fp4_quantize and its returned sf layout/format exactly (refer to fp4_quantize, nvfp4_quantize, and e2m1_and_ufp8sf_scale_to_float) so the test targets the fp4_quantize implementation directly and mirrors the existing nvfp4 test logic.
🤖 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.
Nitpick comments:
In `@flashinfer/quantization/fp4_quantization.py`:
- Around line 842-850: The cute-dsl dispatch path is missing normalization of
the global_scale dtype; before calling _fp4_quantize_cute_dsl convert/cast the
global_scale tensor to float32 (matching the CUDA branch's float32
normalization) so both backends receive the same dtype; update the call site
that currently passes global_scale to first ensure global_scale =
global_scale.to(torch.float32) (or equivalent) and then pass it into
_fp4_quantize_cute_dsl.
In `@tests/utils/test_fp4_quantize.py`:
- Around line 1506-1560: Add a sibling regression test that calls fp4_quantize
directly (instead of nvfp4_quantize) using the same parameterization (m and
scale_dtype) and the same input setup (x with dtype bfloat16, global_scale built
in scale_dtype) found in test_nvfp4_quantize_global_scale_dtype_regression;
reproduce the two assertions: (a) that the scale-factor bytes (from the sf
output of fp4_quantize) are not all zero and (b) that the dequantized magnitude
(using e2m1_and_ufp8sf_scale_to_float or the appropriate fp4 dequant helper)
yields a median |deq/x| between 0.5 and 2.0. Use fp4_quantize and its returned
sf layout/format exactly (refer to fp4_quantize, nvfp4_quantize, and
e2m1_and_ufp8sf_scale_to_float) so the test targets the fp4_quantize
implementation directly and mirrors the existing nvfp4 test logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8741613f-8614-4813-8ea8-ba5f447a39cd
📒 Files selected for processing (3)
csrc/nv_internal/tensorrt_llm/thop/fp4Quantize.cppflashinfer/quantization/fp4_quantization.pytests/utils/test_fp4_quantize.py
|
/bot run |
3 participants
📌 Description
Summary
Fixes #3398.
flashinfer.mm_fp4/nvfp4_quantize(cuda backend) silently produced all-zero or magnitude-wrong outputs for certain batch sizes M when global scales are not provided as bfloat16. The root cause in the issue is that the CUDA quantize kernel reads the global scale asfloat32, but callers commonly pass a bfloat16 global scale — e.g.(448 * 6) / x.abs().max()inheritsx's bf16 dtype. The kernel then misreads it byte-wise, and a dtype guard that would have caught this was commented out.The PR fixes the issue by converting non-fp32 global scale factors to fp32 and adding checks. Also adds unit tests.
🔍 Related Issues
#3398
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ 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
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Bug Fixes
Tests