[mxfp8 training] update triton_to_mxfp8_dim0 nan handling

[danielvegamyhre]

Summary

• Received reports of NaN loss with MXFP8 training that resolved when opting out of using Triton kernel for dim0 quantization (triton_to_mxfp8_dim0)
• Added unit tests with various special values (nan, inf, -inf, subnormals, extremely large/small normal values, etc) to find discrepancies between torch impl and triton
• To make a long story short, I became suspicious that the torch reference impl was also not handling certain cases correctly. To get triton to match, i had to do special sweeps of NaN first, then inf, then -inf (in that order), which killed perf. All of which the CUDA code doesn't have to do! (suspicious). So I updated both torch and triton to match the same TE RCEIL logic, which doesn't need any of this special handling .

Changes

• Computing amax: When computing amax, do NOT propagate NaNs, ignore NaNs
  • TE uses: __hmax which treats nans as missing data and returns the numeric value.
  • TE ref: https://github.com/NVIDIA/TransformerEngine/blob/f4debf6648a080c47eeb2213a3a040b4b2638adb/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh#L225
  • CUDA ref: https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH____HALF__COMPARISON.html#group__cuda__math____half__comparison_1ga964cd0e2994141acc0fcb7a64792faf1
• Computing reciprocal fp32 scale: handle biased e8m0 scale NaN (255) and inf (254) as in TE here
• NaNs in input stay NaNs: applying scale to nan = nan

Benchmarks

(torch) dev@gpu-dev-8951ebdf:~/ao$ CUDA_VISIBLE_DEVICES=1 PYTHONPATH=/home/$USER/ao:$PYTHONPATH python benchmarks/mx_formats/cast_bench.py --mode dim0_mxfp8_triton_rceil --M 32768 --K 7168
M 32768 K 7168 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.11.0+cu130
triton version: 3.6.0
mode: dim0_mxfp8_triton_rceil
time_us 123.9359974861145
mem_bw_gbps 5744.76438195262

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4201

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 8 Unrelated Failures

As of commit 82e7fc9 with merge base 3ad1067 ():

NEW FAILURE - The following job has failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio mslk --index-url https://download.... / linux-job (gh)
  RuntimeError: Command docker exec -t de7492a23c61f033a329aecdb8ce2a460c282e78feda95ce7e3caf745ddab6c1 /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test (CPU 2.10, linux.4xlarge, torch==2.10.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test (CPU 2.9, linux.4xlarge, torch==2.9.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test (CUDA 2.10, linux.g5.12xlarge.nvidia.gpu, torch==2.10.0, cuda, 12.6) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test (CUDA 2.8, linux.g5.12xlarge.nvidia.gpu, torch==2.8.0, cuda, 12.6) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test (CUDA 2.9, linux.g5.12xlarge.nvidia.gpu, torch==2.9.1, cuda, 12.6) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

edit: nevermind, managed to recover performance