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

This proof-of-concept adds UE5M3 block-scale support throughout NVFP4 recipes, quantization kernels, tensor metadata, GEMM integration, and grouped-MLP execution.

Extends native and PyTorch NVFP4 scale handling to E4M3 and UE5M3.
Propagates scale dtype and global-scale semantics through quantization, storage, reconstruction, and GEMM paths.
Adds exact quantization, dequantization, GEMM, grouped-quantization, recipe, and grouped-MLP coverage.

Confidence Score: 5/5

The PR appears safe to merge.

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 correctly repairs the previously reported fallback and one-group reconstruction paths.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Extends NVFP4 tensor and quantizer metadata to distinguish E4M3 and UE5M3 block scales.
transformer_engine/pytorch/quantization.py	Selects the NVFP4 scale dtype from the configured recipe format.
transformer_engine/pytorch/cpp_extensions/gemm.py	Routes UE5M3-scaled NVFP4 operands through the compatible GEMM integration.
transformer_engine/common/cast/nvfp4/core_nvfp4.cuh	Generalizes native NVFP4 scaling logic for the added UE5M3 representation.
tests/pytorch/test_grouped_mlp.py	Expands grouped-MLP coverage for the new NVFP4 scale configurations and repaired execution paths.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  R[NVFP4 recipe] --> Q[Quantizer with scale dtype]
  Q --> K[Native NVFP4 quantization]
  K --> T[FP4 data and block scales]
  T --> G[Scale-aware GEMM]
  G --> O[PyTorch and grouped-MLP output]

_{Reviews (24): Last reviewed commit: "Respect no-tensor-scaling in grouped ten..." | 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>

@ptrendx

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

timmoon10 · 2026-08-25T09:34:18Z

/te-ci L1

ptrendx · 2026-08-25T10:42:57Z

                                const NVTETensor global_scale, size_t h, size_t w,
                                size_t scale_stride_h, size_t scale_stride_w, size_t start_offset,
-                                size_t block_len, cudaStream_t stream);
+                                size_t block_len, NVTEDType scale_dtype, cudaStream_t stream);


Why do we need that API break? We have scale tensor there which has dtype inside it, then why do we need to pass it also as an additional argument?

In general, we should think very carefully here why we would do this (and if really needed then I would much rather have v2 variants rather than break the APi).

We did think about maintaining backward compatibility, but I'm not sure if it's worth the effort since these are advanced functions used specifically for Mcore DDP integration: #3325 (comment)

ptrendx · 2026-08-25T11:13:19Z

+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("return_transpose", [False, True], ids=["rowwise", "with_columnwise"])
+@pytest.mark.parametrize("use_4over6", [False, True], ids=["standard", "4over6"])
+def test_disable_second_level_scale_uses_only_block_scale(


Not really for this test, but for this file - we need to actually test E5M# here - this includes creating the reference implementation just like we have for regular nvFP4 recipe.

ptrendx · 2026-08-25T11:22:45Z

    def expected_e4m3_max(tensor_type):
        if not expected_use_4over6(tensor_type):
-            return 448
+            return 0


Why? Is the result of this function something else than e4m3 max? If so, the name of the function should change too.

There are two changes to how we're interpreting e4m3_max:

It is now an override for the maximum scale dtype value. A value of zero means we use the actual dtype max. This is only used for 4over6, which is not supported with e5m3.

We can no longer assume that the NVFP4 scales are e4m3. I agree it would be better to rename the variable to something like scale_max so it's no longer married to e4m3. However, I'm trying to limit the amount this PR touches 4over6 logic.

ptrendx · 2026-08-25T11:31:49Z

      .value("kBFloat16", transformer_engine::DType::kBFloat16)                                    \
      .value("kFloat8E4M3", transformer_engine::DType::kFloat8E4M3)                                \
      .value("kFloat8E5M2", transformer_engine::DType::kFloat8E5M2)                                \
+      .value("kFloat8E8M0", transformer_engine::DType::kFloat8E8M0)                                \


Why do we put it in the middle of this enum type? That would be breaking for places that have this enum as int.

This matches the dtype order in the C API: https://github.com/timmoon10/TransformerEngine/blob/58dfc41046b7068dffc29eba80ecf65005988009/transformer_engine/common/include/transformer_engine/transformer_engine.h#L25
Also, we explicitly specify the value for each Python enum, so the order doesn't matter.

ptrendx · 2026-08-25T11:44:25Z

                )
                columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device)
        elif compatible_recipe.float8_block_scaling():
+            scale_inv_dtype = DType.kFloat32


Earlier in this file amax buffers are always created, whereas it should check whether the second level scale should be disabled.

zhongbozhu · 2026-08-28T00:44:21Z



 @dataclass(repr=False)
 class NVFP4BlockScaling(Recipe):


If it's just one UE5M3 toggle, we can actually extend it to NVFP4BlockScaling here so that it automatically inherits the Mcore --fp4-param-gather feature. I agree that real hybrid recipes like mixing with mxfp8 should continue to live in CustomRecipe structure.

You're right that basic UE5M3 support just requires setting recipe.fp8_format = Format.UE5M3. This can already be done in Mcore by setting config.fp8_recipe_attrs = {"fp8_format": Format.UE5M3} (once NVIDIA/Megatron-LM#6341 merges).

However, the real practical benefit is that UE5M3 allows us to change the tensor scaling. We skip tensor scaling in the forward pass, while keeping current scaling in the backward pass. This is all a research question and I figure the recipe recommendations might change at any moment, so I'm resistant to hard-coding it in TE. CustomRecipe is a much better approach for unstable research recipes.

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

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

timmoon10 and others added 2 commits August 24, 2026 23:47

Debug minor test failures

30b8af8

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

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

d9b633d

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

This comment was marked as outdated.

Sign in to view

Fix incorrect scale dtypes in grouped tensor builder method

75bdfe4

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

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread transformer_engine/common/cast/dispatch/quantize.cuh Outdated

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread transformer_engine/common/cast/dispatch/quantize.cuh Outdated

timmoon10 commented Aug 25, 2026

View reviewed changes

Comment thread transformer_engine/common/cast/dispatch/quantize.cuh Outdated

Comment thread transformer_engine/common/cast/dispatch/quantize.cuh Outdated

Avoid redundant amax ptr check in row-scaled NVFP4 quantize

58dfc41

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

ptrendx reviewed Aug 25, 2026

View reviewed changes

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

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread tests/cpp/operator/test_cast_nvfp4_transpose.cu

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread tests/cpp/test_common.cu Outdated

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread tests/cpp/test_common.cu Outdated

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread tests/pytorch/test_fusible_ops.py Outdated

ptrendx reviewed Aug 25, 2026

View reviewed changes

Comment thread tests/pytorch/test_grouped_mlp.py Outdated

ptrendx reviewed Aug 25, 2026

View reviewed changes

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

This comment was marked as resolved.

Sign in to view

zhongbozhu reviewed Aug 28, 2026

View reviewed changes

timmoon10 commented Aug 28, 2026

View reviewed changes

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

timmoon10 added 2 commits August 28, 2026 07:38

Clean up C++ unit tests

b1a4483

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

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>

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!

This comment was marked as outdated.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Aug 25, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

timmoon10 commented Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

timmoon10 commented Aug 7, 2026 •

edited

Loading

greptile-apps Bot commented Aug 17, 2026 •

edited

Loading

timmoon10 Aug 27, 2026 •

edited

Loading

timmoon10 Aug 28, 2026 •

edited

Loading