[DylanWanggg]

Description

compute_scale_from_amax returned scale 1 only for a zero, infinite, or NaN amax. A positive amax below max_fp8 * FLT_MIN missed that check, so scale = max_fp8 / amax overflowed and scale_inv was stored as an FP32 subnormal.

max_fp8 * FLT_MIN is exact in binary32. This change floors a positive amax at that value before the division. scale is then at most 2^126 and scale_inv is at least FLT_MIN. An amax of 7.7e-37 is amplified to about 65 and kept as a non-zero FP8 code instead of flushing the tile to zero. Zero, inf, and NaN still return scale 1. An amax at or above the threshold, including ordinary training magnitudes, stays bitwise identical.

The E4M3 threshold is 448 * FLT_MIN (0x04e00000). The E5M2 threshold is 57344 * FLT_MIN (0x08600000).

No linked issue.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• transformer_engine/common/recipe/recipe_common.cuh: Floor a positive amax at max_fp8 * FLT_MIN before dividing, and still return scale 1 for zero, inf, and NaN.
• tests/pytorch/references/quantize_scale_calc.py: Floor a positive tiny amax in the PyTorch reference.
• transformer_engine/pytorch/custom_recipes/reference_current_scaling.py: Floor a positive tiny amax in the custom-recipe reference.
• transformer_engine/jax/quantize/quantizer.py: Floor a positive tiny amax in the JAX quantizer.
• tests/cpp/operator/test_cast_current_scaling.cu: Floor a positive tiny amax in the current-scaling reference.
• tests/cpp/operator/test_cast_transpose_current_scaling.cu: Floor a positive tiny amax in the transpose reference.
• tests/cpp/operator/test_cast_current_scaling_grouped.cu: Floor a positive tiny amax in the grouped reference.
• tests/cpp/operator/test_cast_float8blockwise.cu: Floor a positive tiny block amax in the block-scaling reference.
• tests/cpp/operator/test_tiny_amax_scale_host.cpp: Check that a positive amax below the floor uses scale 2^126 and scale_inv FLT_MIN.
• tests/cpp/operator/test_tiny_amax_scale_host.py: Check that floored tiny blocks stay representable in FP8 and that ordinary values stay bitwise identical.

Delayed scaling _default_sf_compute already clamps an infinite scale to the FP32 maximum and is unchanged.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new blocking issue was found.

Summary

Floors tiny positive amax values before computing FP8 scales, keeping scale_inv at least FLT_MIN.

• Tiny positive FP8 maxima now get a safe scale inverse.

_{Reviews (4) · Last reviewed commit: "Merge upstream main into fix/fp8-tiny-am..." · Reviewed by Greptile}

[DylanWanggg]

@ptrendx Merged current main into this branch (one CMakeLists conflict, both sides kept). Ready for review when you have time. Greptile's earlier notes are addressed, and DCO / pre-commit were green before this merge.