Prototype NVFP4 with FP8 UE5M3 block scales

[timmoon10]

Description

This is a proof-of-concept for NVFP4 with FP8 UE5M3 block scales. Quantization is supported, but GEMMs are not yet available.

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

[timmoon10]

/te-ci L1

[kainzhong]

You can remove this line now since you've done this at line 530 just before "Trivial cases"

[timmoon10]

Won't this break ABI compatibility if accessing as C functions? It's also quite hacky and nonstandard.

We should check, but I believe these are advanced functions that are only used for our Mcore DDP integration. If this is true, I think it's fine to break backward compatibility and just add the scale dtype as a normal function arg.

[tdophung]

yesh it does break ABI. Since I did not want to create separate function for each (like a _v2) with the extra scale_dtype arg, I resorted to this thinking that it is a nvfp4-plus-devel branch specifically anyways so it should be ok to assume that people that intentionally uses it (Mcore) would always pass in the scale_dtype.

But i think you are right, that it is hacky/deceiving, so I'll just make it a normal arg

[timmoon10]

All of this padding logic is unnecessary since unaligned data goes to the wgrad kernel.

[greptile-apps]

Greptile Summary

The PR prototypes NVFP4 quantization with UE5M3 block scales across the CUDA core and PyTorch frontend.

• Adds UE5M3 scale dtype propagation through NVFP4 tensor storage, quantizers, cast/dequantize kernels, and bindings.
• Updates grouped-linear and fused grouped-MLP handling for the new scale format.
• Adds C++ and PyTorch coverage for quantization, dequantization, grouped operations, and exact scale behavior.

Confidence Score: 4/5

The PR is not yet safe to merge because single-group NVFP4 grouped-MLP execution still fails while reconstructing a dense NVFP4 tensor.

The single-group grouped-MLP helper omits the required scale_dtype keyword when constructing NVFP4Tensor, so reachable forward and backward paths raise TypeError before GEMM or scale swizzling.

Files Needing Attention: transformer_engine/pytorch/ops/fused/grouped_mlp.py

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/core_nvfp4.cuh	Generalizes NVFP4 scale handling so quantization kernels can operate with E4M3 or UE5M3 scales.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Propagates the scale dtype through Python NVFP4 tensor metadata and construction.
transformer_engine/pytorch/csrc/quantizer.cpp	Extends the native quantizer bridge to carry UE5M3 NVFP4 scale metadata.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Extends fused grouped-MLP quantization and cuDNN operand handling for UE5M3 scales.
tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py	Adds exact-reference coverage for NVFP4 quantization using the new scale format.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Recipe["NVFP4 recipe<br/>E4M3 or UE5M3 scales"] --> Quantizer["PyTorch NVFP4Quantizer"]
  Quantizer --> Binding["PyTorch C++ binding"]
  Binding --> Core["NVFP4 CUDA quantize/dequantize kernels"]
  Core --> Storage["FP4 data + typed FP8 block scales"]
  Storage --> Tensor["NVFP4Tensor / GroupedTensor"]
  Tensor --> Ops["Grouped-linear and grouped-MLP paths"]

Loading

_{Reviews (2): Last reviewed commit: "Fix compile error" | Re-trigger Greptile}

[greptile-apps]

Unsupported fusion aborts fallback

When the optional fused grouped-MLP implementation is unavailable, this assertion aborts pipeline construction instead of returning the original operations, breaking grouped MLP execution on unsupported architectures, disabled feature configurations, or installations without the required kernels.

Suggested change

	if not fused_op_cls.is_supported():
	assert False ### TODO Remove
	return ops
	if not fused_op_cls.is_supported():
	return ops

Knowledge Base Used: PyTorch Fusible-Operation Framework

[timmoon10]

/te-ci L1