NVIDIA / NVIDIA/TransformerEngine

Preserve NaN through half-precision MXFP8 amax reductions

Open
#3,550 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.