Prototype NVFP4 with FP8 UE5M3 block scales by timmoon10 · Pull Request #3325 · NVIDIA/TransformerEngine

timmoon10 · 2026-08-07T09:59:53Z

Description

This is a proof-of-concept for NVFP4 with FP8 UE5M3 block scales. Quantization is supported natively and GEMMs go through the cuDNN Frontend kernels (see NVIDIA/cudnn-frontend#545).

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Support NVFP4 tensors with UE5M3 scales
Support UE5M3 scales in NVFP4 quantize/dequantize kernels
Support UE5M3 scales in NVFP4 quantizers

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Co-authored-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

* [PyTorch] Enable e5m3 fused GEMM kernels from cuDNN Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * have to pad to 256 to use cuDNN Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix: need to pass scale_dtype Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * route wgrad to cuDNN's wgrad API Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * Support grouped linear with NVFP4-UE5M3 NVFP4-UE5M3 grouped GEMM falls back to dense GEMMs. Generalize usage of wgrad kernel and use when tensors sizes are not 256-aligned. Fix inconsistent m,n,k GEMM notation. Remove ue5m3 hacks in op fuser tests. Add ue5m3 to grouped MLP tests. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Fix typos Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Codex <noreply@openai.com>

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Tim Moon <tmoon@nvidia.com>

greptile-apps · 2026-08-17T21:40:10Z

Greptile Summary

The PR prototypes NVFP4 tensors with FP8 UE5M3 block scales across native quantization, GEMM integration, PyTorch/JAX bindings, recipes, storage, and tests.

Adds UE5M3 scale metadata and propagation through NVFP4 quantizers and tensor storage.
Extends quantization, dequantization, transpose, Hadamard, and GEMM paths for the new scale format.
Updates grouped-MLP fusion compatibility and optional-kernel fallback handling.
Adds native and PyTorch numerical and integration coverage.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains in the previously reviewed paths.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Propagates UE5M3 scale semantics through fused grouped-MLP execution and fixes the previously reported unsupported-kernel fallback paths.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Extends NVFP4 tensor and quantizer metadata to preserve the selected block-scale dtype.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Carries NVFP4 scale dtype through internal storage and reconstruction paths.
transformer_engine/common/cast/nvfp4/core_nvfp4.cuh	Generalizes native NVFP4 scaling primitives for E4M3 and UE5M3 scale representations.
transformer_engine/common/gemm/config.cpp	Extends GEMM configuration and dispatch for NVFP4 operands using UE5M3 scales.
tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py	Adds exact quantization validation for the new NVFP4 scale format.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    R[NVFP4 recipe with UE5M3 scales] --> Q[Framework quantizer]
    Q --> T[NVFP4 tensor and grouped storage]
    T --> B[PyTorch or JAX native bindings]
    B --> K[Quantize and dequantize kernels]
    B --> G[cuDNN frontend GEMM kernels]
    K --> O[Encoded FP4 values and UE5M3 scales]
    G --> O
    O --> V[Exact and integration tests]

_{Reviews (30): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 · 2026-08-28T09:15:57Z

/te-ci L1

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 · 2026-08-30T01:24:29Z

/te-ci L1

…ard compatibility Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

timmoon10 · 2026-09-01T02:04:55Z

/te-ci L1

* Make NVFP4-UE5M3 grouped GEMM CUDA-graph safe The UE5M3 grouped GEMM builds its one-element int32 offsets tensor inline at two call sites: torch.tensor([tokens], dtype=torch.int32, device=a_tensor.device) # wgrad torch.tensor([N], dtype=torch.int32, device=device) # fprop Each of these materializes an unpinned CPU tensor and copies it host-to-device. CUDA rejects that during graph capture: RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor is pinned. Please use tensor.pin_memory() or allocate the tensor with pin_memory=True. Any model capturing these GEMMs in a CUDA graph therefore fails at capture time. This was hit end-to-end on DeepSeek-V3 671B, where every convergence config captures full forward/backward graphs; it aborts during warmup capture and no training step completes. Route both sites through a cached helper, mirroring get_cached_ones_tensor in this same file -- which both functions already use for their `ones` tensor a few lines earlier, and whose docstring notes it keeps "stable data pointers across CUDA graph replays". Caching by (value, device) moves the host-to-device copy to warmup and keeps the pointer stable across replays. The offsets are constant for a given shape, so this is numerically neutral. Signed-off-by: Michal Futrega <mfutrega@nvidia.com> * Review suggestions Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Michal Futrega <mfutrega@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

tdophung · 2026-09-01T23:38:35Z

/te-ci L1

b3415c1 gated the NVFP4 per-tensor scale branch in cublas_gemm() on have_nvfp4_amax, an OR over the two operands' amax pointers. That branch is also what moves alpha/beta from host into device workspace, and CUBLASLT_POINTER_MODE_DEVICE is set unconditionally for every NVFP4 GEMM. So when both operands had a null amax the branch was skipped, leaving alpha/beta as host stack pointers that cuBLAS then dereferenced as device addresses -- garbage output, not a mis-scaled result. That is exactly the shape of the CI failure: only the both_unit_global_scale parametrization of test_gemm_with_missing_nvfp4_amax went red, while the x_/w_ cases passed because the OR still held. Note that nvte_nvfp4_compute_per_tensor_scale with a null amax is an exact identity on alpha (it substitutes scale_max * fp4_max), so skipping it never saved meaningful work -- its real job on that path is materializing the scalars on device. Making the call is always safe; skipping it is not. Replace the heuristic with kNVTEMatmulConfigAlphaBetaOnDevice, which states pointer residency directly instead of inferring it. Default false, so PyTorch (nvte_cublas_gemm/_scaled pass host pointers, and generic_gemm reaches nvte_cublas_gemm_v2 without setting the flag) is restored to its pre-b3415c19 behavior. JAX sets it when is_nvfp4_scaling(), matching the device buffers it binds for alpha/beta, and folds the per-tensor scale into alpha itself. Also document in the JAX helper that its block-scale bound is hardcoded to E4M3's 448 while the C++ get_nvfp4_scale_max() returns 114688 for UE5M3, which will need reconciling once UE5M3 reaches the JAX path. Verified on GB200: tests/pytorch/nvfp4 9310 passed / 0 failed / 6852 skipped (CI at b3415c1 was 9309 passed / 1 failed / 6852 skipped), tests/pytorch/mxfp8 371 passed, and JAX nvfp4+gemm 382 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tdophung · 2026-09-02T18:51:46Z

/te-ci L1

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

…EMM (#5) * Do not swizzle operand scales in place in the NVFP4-UE5M3 GEMM The cuDNN UE5M3 GEMM wrapper swizzled A/B scales in place and marked the tensors as swizzled. For weights that mark persisted on the parameter. The master-weight cast (_cast_master_weights_to_nvfp4_2d, used by Megatron's fp4 param gather after every optimizer step) then rewrote the scales in the unswizzled layout without clearing the mark, so every later GEMM and dequantize read unswizzled scales as swizzled. Gaussian weights hide it (near-uniform scales, ~0.2 relative error); real DeepSeek-V3 weights are destroyed (2.6 relative error, loss 12.98 vs 7.86 for E4M3 under identical routing). Swizzle clones of the scale tensors instead and leave the operands untouched, matching what the cuBLAS path does in C++. * Test that the NVFP4-UE5M3 GEMM leaves operand scales untouched * Avoid suppressing errors in swizzle function. Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 · 2026-09-02T21:02:00Z

/te-ci L1

for more information, see https://pre-commit.ci

jberchtold-nvidia

Approving for code owner permissions for the two TE/JAX files. Thanks!

Prototype NVFP4 with UE5M3 scales

25d5b4b

Co-authored-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 added the enhancement New feature or request label Aug 7, 2026

[pre-commit.ci] auto fixes from pre-commit.com hooks

dede7c7

for more information, see https://pre-commit.ci

This comment was marked as outdated.

Sign in to view

timmoon10 added the 2.19 label Aug 10, 2026

ptrendx assigned ptrendx and unassigned ptrendx Aug 11, 2026

vthumbe1503 mentioned this pull request Aug 11, 2026

[Pytorch][Common] Option to Disable 2nd level Scale in NVFP4 #3297

Closed

13 tasks

kainzhong reviewed Aug 13, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/cpp_extensions/gemm.py Outdated

timmoon10 commented Aug 13, 2026

View reviewed changes

Comment thread transformer_engine/common/include/transformer_engine/recipe.h

timmoon10 added 2 commits August 13, 2026 23:06

Merge branch 'main' into nvfp4-ue5m3-prototype

8feb5e9

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Use custom recipe for NVFP4-UE5M3 tests

ab3a9b3

Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 force-pushed the nvfp4-ue5m3-prototype branch from 0987860 to ab3a9b3 Compare August 14, 2026 12:09

[pre-commit.ci] auto fixes from pre-commit.com hooks

92b1063

for more information, see https://pre-commit.ci

timmoon10 commented Aug 14, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/cpp_extensions/gemm.py Outdated

timmoon10 and others added 5 commits August 14, 2026 12:44

Add grouped MLP kernel for GGEMM+SwiGLU+RHT+quant

3d25d20

Signed-off-by: Tim Moon <tmoon@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

f030e40

for more information, see https://pre-commit.ci

Merge branch 'main' into nvfp4-ue5m3-prototype

535ea4f

Signed-off-by: Tim Moon <tmoon@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

e366e88

for more information, see https://pre-commit.ci

Debug integration with GGEMM+GLU+RHT+quant

252ae94

Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 marked this pull request as ready for review August 17, 2026 21:35

timmoon10 requested review from Oleg-Goncharov, ksivaman and ptrendx as code owners August 17, 2026 21:35

greptile-apps Bot reviewed Aug 17, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py

timmoon10 and others added 2 commits August 18, 2026 02:26

Remove scale max helper functions from NVFP4 cast utils

f2b9c98

Signed-off-by: Tim Moon <tmoon@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

8685556

for more information, see https://pre-commit.ci

Address some review comments from @ptrendx

93cbd0a

Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 force-pushed the nvfp4-ue5m3-prototype branch from 98792fe to 93cbd0a Compare August 28, 2026 08:54

pre-commit-ci Bot and others added 2 commits August 28, 2026 08:55

[pre-commit.ci] auto fixes from pre-commit.com hooks

d00cbeb

for more information, see https://pre-commit.ci

Respect no-tensor-scaling in grouped tensor helper function

3a62970

Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 and others added 2 commits August 29, 2026 23:25

Debug C++ test compilation error

e404db8

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Make sure GEMM alpha/beta scales are on GPU, even without amaxes

adef2dc

Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

ptrendx added 2.20 and removed 2.19 labels Aug 31, 2026

timmoon10 and others added 5 commits August 31, 2026 23:10

Create new Mcore DDP integration functions rather than breaking backw…

a971552

…ard compatibility Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

Remove unrelated tests from NVFP4 dequantize C++ unit tests

bb59631

Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

Add reference impl and test for NVFP4-UE5M3

85cf34f

Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

Merge branch 'main' into nvfp4-ue5m3-prototype

d73bb04

[pre-commit.ci] auto fixes from pre-commit.com hooks

bdc1d97

for more information, see https://pre-commit.ci

tdophung reviewed Sep 1, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/cublaslt_gemm.cu

Fix JAX NVFP4 GEMM device scalar handling

b3415c1

tdophung requested a review from jberchtold-nvidia as a code owner September 2, 2026 18:49

jberchtold-nvidia reviewed Sep 2, 2026

View reviewed changes

Comment thread transformer_engine/jax/csrc/extensions/gemm.cpp

timmoon10 commented Sep 2, 2026

View reviewed changes

Comment thread transformer_engine/common/include/transformer_engine/gemm.h Outdated

Comment thread transformer_engine/common/gemm/cublaslt_gemm.cu

timmoon10 and others added 2 commits September 2, 2026 14:01

Document bugs with alpha_beta_on_device GEMM config

ee6244e

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

a7ca1a1

for more information, see https://pre-commit.ci

jberchtold-nvidia approved these changes Sep 2, 2026

View reviewed changes

Navigation Menu

Uh oh!

Prototype NVFP4 with FP8 UE5M3 block scales - #3325

Prototype NVFP4 with FP8 UE5M3 block scales#3325
timmoon10 wants to merge 57 commits into
NVIDIA:mainfrom
timmoon10:nvfp4-ue5m3-prototype

timmoon10 commented Aug 7, 2026 •

edited

Loading

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot commented Aug 17, 2026 •

edited

Loading

Uh oh!

Uh oh!

timmoon10 commented Aug 28, 2026

Uh oh!

timmoon10 commented Aug 30, 2026

Uh oh!

timmoon10 commented Sep 1, 2026

Uh oh!

Uh oh!

tdophung commented Sep 1, 2026

Uh oh!

tdophung commented Sep 2, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Sep 2, 2026

Uh oh!

jberchtold-nvidia left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

timmoon10 commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

timmoon10 commented Aug 28, 2026

Uh oh!

timmoon10 commented Aug 30, 2026

Uh oh!

timmoon10 commented Sep 1, 2026

Uh oh!

Uh oh!

tdophung commented Sep 1, 2026

Uh oh!

tdophung commented Sep 2, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Sep 2, 2026

Uh oh!

jberchtold-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

timmoon10 commented Aug 7, 2026 •

edited

Loading

greptile-apps Bot commented Aug 17, 2026 •

edited

Loading