[xpu][mx] Fix NaN scale propagation in RCEIL triton kernel

[ugolowic]

In the context of #3576.

Fix NaN scale propagation in RCEIL triton kernel for SPIR-V backends

On XPU's SPIR-V backend, tl.clamp and .to(uint8) silently destroy NaN values:

• tl.clamp uses NaN-ignoring fmin/fmax (returning -127.0 instead of NaN) - see intel/intel-xpu-backend-for-triton#5003 and KhronosGroup/SPIRV-LLVM-Translator#3282
• fptoui(NaN), used by .to(uint8) is undefined behavior per the LLVM spec (returning 0 instead of 255) - see The LLVM LangRef fptoui semantics section

The existing NaN workaround in _triton_calculate_scale_rceil correctly detected NaN blocks via x != x + tl.max(nan_mask), but then injected NaN into max_abs, which was subsequently killed by tl.clamp further down the chain.

The proposed solution is to move the NaN correction after the log2 -> ceil -> clamp -> +127 -> to(uint8) chain.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull workflow for PyTorch trunk commits

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

[anmyachev]

Hi @ugolowic,

tl.clamp uses NaN-ignoring fmin/fmax (returning -127.0 instead of NaN) - see intel/intel-xpu-backend-for-triton#5003 and KhronosGroup/SPIRV-LLVM-Translator#3282

This is the default behavior. Did you try tl.clamp(..., propagate_nan=tl.PropagateNan.ALL)?

[danielvegamyhre]

do our existing special value tests in test_kernels.py pass with this? the desired behavior here is that if any value in the 1x32 block is nan, the scale becomes nan and the whole quantized block becomes nan.

we did a lot of work to fix a nan loss issue, and make torch vs triton special value handling consistent and correct: #4201

so i am hesitant to make any numerical changes now that we have confirmed in production the fix works and users are using the kernel again.