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 to NVFP4 quantization and routes compatible GEMMs through cuDNN Frontend.

• Propagates the scale dtype through native kernels, PyTorch quantizers, tensor storage, reconstruction, and grouped operations.
• Adds UE5M3-aware scaling, dequantization, GEMM dispatch, and grouped-MLP execution.
• Extends native and PyTorch numerical coverage for quantization, GEMM, grouped execution, and graph-safe paths.

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 metadata and scaling through grouped-MLP forward and backward kernels while preserving unfused and optional-kernel fallbacks.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Extends NVFP4 quantizer and tensor metadata to retain the selected block-scale dtype.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Updates NVFP4 storage allocation and reconstruction for E4M3 or UE5M3 block scales.
transformer_engine/pytorch/quantization.py	Adds UE5M3 capability checks and creates NVFP4 quantizers with recipe-selected scale dtypes.
transformer_engine/common/cast/nvfp4/core_nvfp4.cuh	Generalizes native NVFP4 scale calculations and storage handling for UE5M3.
transformer_engine/common/gemm/cublaslt_gemm.cu	Updates GEMM dispatch boundaries to account for the new NVFP4 scale representation.
tests/pytorch/test_grouped_mlp.py	Expands grouped-MLP coverage for UE5M3 NVFP4 configurations and execution paths.
tests/cpp/operator/test_cast_nvfp4_transpose.cu	Adds exact native reference coverage for UE5M3 NVFP4 cast and transpose behavior.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    R[NVFP4 recipe] --> Q[UE5M3-aware quantizer]
    Q --> T[FP4 data and UE5M3 block scales]
    T --> S[Scale layout preparation]
    S --> G[cuDNN Frontend GEMM]
    G --> O[High-precision or quantized output]
    T --> D[Dequantization]
    D --> H[High-precision tensor]

Loading

_{Reviews (21): Last reviewed commit: "Avoid redundant amax ptr check in row-sc..." | Re-trigger Greptile}

[timmoon10]

/te-ci L1