[ROCm] Make the Python NFP8 dtype selection architecture-aware - #167
Open
aryaman-gupta wants to merge 3 commits into
Open
[ROCm] Make the Python NFP8 dtype selection architecture-aware#167aryaman-gupta wants to merge 3 commits into
aryaman-gupta wants to merge 3 commits into
Conversation
The C++ side picks the FP8 e4m3 encoding per device arch (getNFP8ScalarType: fnuz on gfx94x/gfx90a, OCP fn elsewhere), but Python still selected on "is this a ROCm build" and always returned fnuz. On gfx950 that labels the weight buffer fnuz while the device decodes OCP fn, so every PyTorch-native read of those bytes is off by one exponent bias - a factor of two. Add _nfp8_is_fnuz()/nfp8_dtype(), mirroring the C++ arch check, and use them in sparse_type_int_to_dtype, SparseType.as_dtype, the NFP8 weight init, and the two test-harness fp8_dtype globals. No new kernel instantiations: relabel_nfp8_for_dispatch() already views an fn-labeled tensor as fnuz at the host boundary, so the single fnuz instantiation still serves both archs. as_dtype() resolves NFP8 before its lookup table rather than inside it, because the table is rebuilt per call and would otherwise run a device query for every dtype lookup. Fixes test_forward_gpu_nfp8_format_matches_host_decode on gfx950. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Shorten the _nfp8_is_fnuz docstring and the as_dtype note, and invert the guard to an early return so it is evident that non-ROCm builds never reach the device query. Behaviour is unchanged - the `and` already short-circuited. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
avbokovoy
approved these changes
Aug 19, 2026
The ty == 9 branch duplicated nfp8_dtype()'s body. Call it instead, leaving _nfp8_is_fnuz() with a single caller. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
FP8 e4m3 has two encodings:
fnuzon gfx90a/gfx94x (MI300), and OCPfnon gfx950 and CUDA.The C++ side already selects per-arch via
getNFP8ScalarType(), but Python selected on "is this a ROCm build" and always returnedfnuz. On gfx950 that labels the weight tensorfnuzwhile the device decodes OCPfn, so every PyTorch-native read of those bytes (casting, printing, serialization) is off by one exponent bias — a factor of two.Fix
Add
_nfp8_is_fnuz()/nfp8_dtype(), mirroring the C++ arch check, and use them at the dtype-selection sites.No new kernel instantiations.
relabel_nfp8_for_dispatch()already views anfn-labelled tensor asfnuzat the host boundary, so the singlefnuzkernel instantiation still serves both archs. An earlier attempt at this problem emitted both FP8 kernel variants on ROCm and was rolled back over build time; this approach avoids that entirely.as_dtype()resolves NFP8 ahead of its lookup table rather than inside it, because the table is rebuilt on every call and would otherwise run a device query for every dtype lookup.Testing (MI350X / gfx950, ROCm 7.1)
test_forward_gpu_nfp8_format_matches_host_decode: fails (max abs err 26) → passesforward_test.py: 14 passed, 0 failedno_cache_fp8,uvm_cache_fp8,backward_adagrad_fp8_pmSUM) — all pass, which exercises the forward, cache, backward and optimizer relabel paths. Those skips are left untouched in this PR.