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 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

[greptile-apps]

Greptile Summary

This proof-of-concept adds UE5M3 block-scale support throughout NVFP4 quantization, tensor metadata, bindings, GEMM paths, and tests.

• Adds UE5M3 scale dtype and availability/configuration plumbing.
• Updates NVFP4 quantization, dequantization, grouped GEMM, and fused grouped-MLP paths.
• Extends C++ and PyTorch coverage for UE5M3-scaled NVFP4 tensors.

Confidence Score: 5/5

The pull request 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 NVFP4 scale dtype through single-group reconstruction, safely retains unsupported-fusion fallback, and adds UE5M3-aware cuDNN grouped-GEMM handling.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Extends NVFP4 quantizer and tensor metadata with validated E4M3 or UE5M3 scale types and second-level-scale configuration.
transformer_engine/pytorch/quantization.py	Adds UE5M3 capability detection and creates NVFP4 quantizers with the recipe-selected block-scale dtype.
transformer_engine/common/cast/nvfp4/core_nvfp4.cuh	Generalizes core NVFP4 scale processing for UE5M3 alongside existing E4M3 scaling.
transformer_engine/pytorch/csrc/quantizer.cpp	Carries the new NVFP4 scale dtype across the Python-to-C++ quantizer and tensor-allocation boundary.
tests/pytorch/test_grouped_mlp.py	Extends grouped-MLP coverage for the updated NVFP4 recipe and scale-format behavior.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Recipe["NVFP4 recipe<br/>UE5M3 scale format"] --> Quantizer["NVFP4Quantizer<br/>scale_dtype"]
  Quantizer --> Tensor["NVFP4 tensor/storage<br/>FP4 data + UE5M3 scales"]
  Tensor --> Bindings["PyTorch/C++ bindings"]
  Bindings --> Kernels["Quantize/dequantize kernels"]
  Bindings --> GEMM["cuDNN grouped and dense GEMM"]

Loading

_{Reviews (15): Last reviewed commit: "Enable cuDNN GGEMM+GLU+RHT+quant kernel" | Re-trigger Greptile}

[timmoon10]

/te-ci L1

[timmoon10]

/te-ci L1