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 quantization with FP8 UE5M3 block scales across the CUDA core, PyTorch bindings, quantized tensor metadata, GEMM paths, and grouped-MLP integration.

Adds UE5M3 scale-type propagation through NVFP4 recipes, quantizers, storage, and kernels.
Updates GEMM and grouped-GEMM paths to interpret UE5M3 scale buffers correctly.
Extends C++ and PyTorch coverage for quantization, dequantization, GEMM, grouped quantization, and fused operations.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains from the previously reported issues.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Propagates UE5M3 scale metadata through grouped-MLP kernels and completes the previously reported fallback and single-group NVFP4 reconstruction fixes.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Extends NVFP4 tensor and quantizer metadata to represent the block-scale dtype.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Stores and preserves NVFP4 scale dtype alongside packed data, scales, and amax metadata.
transformer_engine/pytorch/cpp_extensions/gemm.py	Routes UE5M3-scaled NVFP4 operands through the corresponding GEMM handling.
transformer_engine/common/cast/nvfp4/core_nvfp4.cuh	Generalizes core NVFP4 scale operations to support UE5M3 in addition to E4M3.

_{Reviews (13): 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>

timmoon10 · 2026-08-21T02:15:53Z

/te-ci L1

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

Signed-off-by: Kaining Zhong <kainingz@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

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

Fix compile error

803dd24

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

This comment was marked as outdated.

Sign in to view

timmoon10 commented Aug 19, 2026

View reviewed changes

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

This comment was marked as outdated.

Sign in to view

tdophung force-pushed the nvfp4-ue5m3-prototype branch from 0ad7060 to cbae89d Compare August 20, 2026 18:10

greptile-apps Bot reviewed Aug 20, 2026

View reviewed changes

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

This comment was marked as outdated.

Sign in to view

tdophung and others added 12 commits August 21, 2026 02:55

Fix NVFP4 scale dtype ABI

27f29d0

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

remove redundant output alloc

c9dad3e

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

no need to pad N now

2bf42eb

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

fix linting errors

b6913b2

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

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

f2c26bd

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

Disable cuDNN GGEMM+GLU+RHT+quant kernel

030e9af

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

Guard NVFP4 alpha scaling by scaling mode

9162e98

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

Restore UE5M3 NVFP4 cast support

db0957b

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

Localize NVFP4 4over6 scale policy

982be7b

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

Tweak arg order in C API functions

01f4edc

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

Rename cuDNN GGEMM helper functions for general_gemm

3a63623

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

Merge branch 'main' into HEAD

26db2ee

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

timmoon10 force-pushed the nvfp4-ue5m3-prototype branch from ef16e25 to 26db2ee Compare August 21, 2026 02:59

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

4e9d0a4

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

Navigation Menu

Uh oh!

Prototype NVFP4 with FP8 UE5M3 block scales - #3325

Prototype NVFP4 with FP8 UE5M3 block scales#3325
timmoon10 wants to merge 27 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!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

timmoon10 commented Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 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

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

timmoon10 commented Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

timmoon10 commented Aug 7, 2026 •

edited

Loading

greptile-apps Bot commented Aug 17, 2026 •

edited

Loading