Repository navigation
fix(quant): guard the fp8-rounded SF against underflow in block-scale quantization #6996
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -360,9 +360,12 @@ struct FP4Converter< | |
| QuantizationSFLayout::SWIZZLED); | ||
| *SFOffset = fp8SFVal; | ||
|
|
||
| // Guard the post-rounding SFValue, not vecMax: the fp8 SF conversion can underflow small | ||
| // vecMax values to zero, and reciprocal(0) = inf turns zero elements into NaN (see #2441). | ||
| float outputScale = | ||
| vecMax != 0.f ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) | ||
| : 0.0f; | ||
| SFValue != 0.f | ||
| ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) | ||
| : 0.0f; | ||
| float2 fp2Vals[ELTS_PER_THREAD_T / 2]; | ||
| #pragma unroll | ||
| for (int i = 0; i < static_cast<int>(ELTS_PER_THREAD_T); i += 2) { | ||
|
|
@@ -441,7 +444,9 @@ struct FP8Converter< | |
| tmpSFVal.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf); | ||
| SFValue = static_cast<float>(tmpSFVal); | ||
| fp8SFVal = tmpSFVal.__x; | ||
| float outputScale = vecMax != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f; | ||
| // Guard the post-rounding SFValue, not vecMax: the fp8 SF conversion can underflow small | ||
| // vecMax values to zero, and reciprocal(0) = inf turns zero elements into NaN (see #2441). | ||
| float outputScale = SFValue != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f; | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win 🔎 Supported by static analysis🏁 Script executed: sed -n 340,470p include/flashinfer/norm/fused_dit_layernorm.cuhRepository: flashinfer-ai/flashinfer Length of output: 5558 🏁 Script executed: rg -n -F --glob '*.{cuh,cu,h,hpp,cc,cpp}' -e 'reciprocal_approximate_ftz' -e '__nv_cvt_float_to_e8m0' -e 'static_cast<float>(tmpSFVal)' -e '__nv_fp8_e8m0' include csrc 2>/dev/null || test "$?" -eq 1
printf '\n--- relevant helper definitions ---\n'
rg -n -F --glob '*.{cuh,cu,h,hpp,cc,cpp}' -e '__device__.*reciprocal_approximate_ftz' -e '__forceinline__.*reciprocal_approximate_ftz' include csrc 2>/dev/null || test "$?" -eq 1Repository: flashinfer-ai/flashinfer Length of output: 8939 🌐 Web query:
💡 Result: 🏁 Script executed: printf '%s\n' '--- helper and conversion references ---'
rg -n -F --glob '*.cuh' --glob '*.cu' --glob '*.h' --glob '*.hpp' --glob '*.cc' --glob '*.cpp' \
-e 'reciprocal_approximate_ftz' \
-e '__nv_cvt_float_to_e8m0' \
-e '__nv_fp8_e8m0' \
include csrc 2>/dev/null || test "$?" -eq 1
printf '%s\n' '--- CUDA headers available in checkout ---'
rg --files | rg '(^|/)(cuda_fp8|cuda_runtime|math_functions|helper_math).*\\.(h|hpp|cuh)$' || trueRepository: flashinfer-ai/flashinfer Length of output: 8783 🏁 Script executed: printf '%s\n' '--- reciprocal helper ---'
sed -n '225,240p' include/flashinfer/norm/fused_dit_layernorm.cuh
printf '%s\n' '--- reviewed E8M0 block ---'
sed -n '437,455p' include/flashinfer/norm/fused_dit_layernorm.cuh
printf '%s\n' '--- base comparison for this block ---'
git diff 6a2f2661d24db6400466db51f482cd3c95faffaa -- include/flashinfer/norm/fused_dit_layernorm.cuh | sed -n '/FP8Converter/,/^[^+ -]/p' || trueRepository: flashinfer-ai/flashinfer Length of output: 2366 Preserve the For Suggested fix--- "a/include/flashinfer/norm/fused_dit_layernorm.cuh"
+++ "b/include/flashinfer/norm/fused_dit_layernorm.cuh"
@@ -444,9 +444,9 @@
tmpSFVal.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf);
SFValue = static_cast<float>(tmpSFVal);
fp8SFVal = tmpSFVal.__x;
- // Guard the post-rounding SFValue, not vecMax: the fp8 SF conversion can underflow small
- // vecMax values to zero, and reciprocal(0) = inf turns zero elements into NaN (see #2441).
- float outputScale = SFValue != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f;
+ // E8M0 has no zero encoding. Preserve zero output for all-zero blocks.
+ float outputScale =
+ vecMax != 0.f && SFValue != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f;
auto SFOffset = cvt_quant_get_sf_out_offset<uint32_t, NUM_THREADS_PER_SF>(
batchIdx, rowIdx, colIdx, numRows, numCols / SF_VEC_SIZE, SFout,🤖 Prompt for AI Agents
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. e8m0 has no zero encoding, so an all-zero block gets byte 0x00 (2^-127, an fp32 denormal), and |
||
|
|
||
| auto SFOffset = cvt_quant_get_sf_out_offset<uint32_t, NUM_THREADS_PER_SF>( | ||
| batchIdx, rowIdx, colIdx, numRows, numCols / SF_VEC_SIZE, SFout, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The stock NVFP4 quantizer still guards
vecMaxatcsrc/nv_internal/tensorrt_llm/kernels/quantization_utils.cuh:505,:723and:734, so on an underflowing blocknvfp4_quantize_smoothnow emits zero codes whilenvfp4_quantizeemits saturated ones, which breaks the bit-equivalencetest_nvfp4_quantize_smoothrelies on. Could you apply the same guard there in this PR?