NVIDIA / NVIDIA/TransformerEngine
Preserve NaN through half-precision MXFP8 amax reductions
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 831
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 65
Description
Follow-up from #3459.
The merged half-precision MXFP8 rowwise path still reduces BF16/FP16 pairs through ptx::abs_max_2x before converting the block maximum to E8M0. A mixed finite + NaN pair can therefore lose the NaN during the reduction and produce a finite scale instead of the exceptional E8M0 value 255.
The same issue affects the bidimensional path because thread_amax is already derived from that reduction before the row/column block maxima are combined.
Current merged code still contains:
- ptx::abs_max_2x(thread_amax_2x, ...) in the half-precision rowwise reduction;
- no focused mixed finite+NaN regression for rowwise or bidimensional MXFP8 scale generation.
Suggested fix:
- make the half-precision reduction explicitly NaN-preserving before scale generation;
- add mixed finite+NaN regressions for both rowwise and bidimensional scaling and compare the produced scale/payload behavior with the established MXFP8 path.
This was the remaining review blocker on #3459 when it merged.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the half-precision MXFP8 rowwise reduction at ptx::abs_max_2x(thread_amax_2x, ...), then trace how thread_amax feeds rowwise and bidimensional block maxima and E8M0 scale generation. Add mixed finite+NaN regressions for both paths and compare their scale and payload behavior with the established MXFP8 path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100