Implement 4over6 NVFP4 recipe

[zianglih]

Description

@HumansAnd

Implement 4over6 nvfp4 from:

• Paper: https://arxiv.org/abs/2512.02010
• Code: https://github.com/mit-han-lab/fouroversix

FlashInfer PR:

• Support 4over6 nvfp4 for quantizer and fused MoE flashinfer-ai/flashinfer#3264

Enable per-block map-to-4 versus map-to-6 candidate selection for 1D/2D NVFP4 quantization in the NVFP4BlockScaling recipe. This mode currently requires RHT and stochastic rounding to be disabled. Both original per-tensor scaling and row-scaling NVFP4 introduced by #2931 are supported.

This PR also fixes a few minor bugs for row-scaled NVFP4 from #2931.

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

Please list the changes introduced in this PR:

• Adds scoped NVFP4 4over6 control through NVTE_NVFP4_4OVER6=weights|activations|all, with unset preserving existing behavior, and threads the selected scope through recipes, quantizers, tensor metadata, split quantization, single-tensor quantization, and C++ tensor/config APIs.
• Implements 1D & 2D NVFP4 4over6 quantization in the existing NVFP4 CUDA paths by comparing TE-style map-to-4 and map-to-6 FP4 candidates with the original 4over6 MSE rule, choosing map-to-6 on ties, honoring NVTE_USE_FAST_MATH, and rejecting unsupported combinations such as stochastic rounding, grouped tensors, and RHT.
• Updates dequantization and NVFP4 GEMM scaling to respect per-tensor 4over6 metadata, using 256-based normalization for 4over6 tensors and 448-based normalization for regular NVFP4 tensors without requiring callers to do hidden rescaling.
• Extends the Python reference implementation to mirror the intended ground truth, meaning TE-style candidate quantization plus original 4over6 MSE/compare logic, and uses this reference for bitwise exact tests where fast math is disabled.
• Expands C++ and Python coverage across exact NVFP4 quantization, GEMM, dequantization, recipe scope resolution, quantized tensor handling, numerics, sanity, CUDA graph, torch compile, CPU offload, fusible ops, and backward override paths, while documenting the new environment variable and known unsupported modes.

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 PR implements the NVFP4 4over6 quantization algorithm from the FourOverSix paper, enabling per-block map-to-4 vs. map-to-6 candidate selection for 1D and 2D NVFP4 quantization. The feature is gated behind a new NVTE_NVFP4_4OVER6 environment variable (values: weights, activations, all) and the nvfp4_4over6 field on NVFP4BlockScaling.

• New CUDA kernel path (quantize_4over6_nvfp4.cuh): Evaluates map-to-4 (1.5\u00d7 expanded block scale) and map-to-6 (standard block scale) FP4 candidates for each 1\u00d716 block, selects the lower-MSE result with ties going to map-to-6, and emits tensors using a 256-based global E4M3 scale bound instead of 448.
• Full stack propagation: The use_4over6 flag threads through Python recipe \u2192 NVFP4Quantizer \u2192 NVFP4Tensor/NVFP4TensorStorage \u2192 C++ Tensor/QuantizationConfig \u2192 all CUDA dispatch paths, with guard checks rejecting incompatible combinations (RHT, stochastic rounding, grouped tensors).
• Reference implementation and testing: A Python reference in quantization_ref_nvfp4.py mirrors the CUDA MSE logic, and new tests cover exact quantization, GEMM, dequant, and recipe scope resolution.

Confidence Score: 4/5

Safe to merge with awareness of the pre-existing fast-math gap on the single-tensor path; all new 4over6 code is well-guarded and the core quantization logic aligns with the reference.

The 4over6 feature is implemented thoroughly across all 41 changed files. The flag threads correctly through Python recipe, tensor metadata, C++ quantizer, and all CUDA dispatch paths. Validation checks reject incompatible combinations at multiple layers. The Python reference mirrors the CUDA MSE logic with the correct 256-denominator and 1.5x scale expansion. The only new findings are minor style-level concerns.

transformer_engine/pytorch/csrc/extensions/cast.cpp - the use_fast_math env-var read is only wired into the split-quantize helper, not into the single-tensor quantize_impl, so the fast-math variant of the 4over6 MSE kernel is unreachable for ordinary single-tensor calls.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh	New file implementing 4over6 device-side quantization helpers: scale computation, MSE accumulation via cvt_fp32_to_fp4_8x_with_mse_rn, candidate selection, and shared-memory scratch for the 2D block decision. Core logic is sound and well-templated.
transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh	Moves row_scaled_nvfp4 and use_4over6 from runtime booleans to compile-time template parameters; updates the dequant factor denominator from 448 to 256 for 4over6 tensors using a constexpr ternary. Change is correct and clean.
transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu	Adds kUse4Over6 and kUseFastMath template params to the main 1D cast-transpose kernel; always allocates a static FourOverSixScratch shared-memory struct (minimized to depth-1 when 4over6 is disabled), adding minor shared-memory overhead to non-4over6 instantiations.
transformer_engine/common/recipe/init.py	Adds nvfp4_4over6 field with validation in __post_init__; correctly requires disable_rht and disable_stochastic_rounding for activations/all scopes but not for weights.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Guards grouped and RHT split-quantize against use_4over6; wires use_4over6 and use_fast_math through split_quantize_nvfp4_impl_helper. The use_fast_math env-var read is only performed here, not in the single-tensor quantize_impl path.
transformer_engine/pytorch/csrc/quantizer.cpp	Reads use_4over6 from the Python quantizer and propagates it to tensor wrappers and the quantization config. Validation checks for incompatible RHT/SR combinations are in place in quantize_impl.
transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py	Adds _quantize_blockwise_4over6_reference mirroring the CUDA MSE logic with the correct 256-denominator and 1.5x scale expansion. The use_4over6 attribute is resolved with a multi-level getattr fallback, which is slightly fragile.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Correctly propagates use_4over6 through __new__, copy, view/reshape autograd functions, sharding metadata, and __reduce_ex__; the nvfp4_shard_metadata tuple grows from 5 to 7 elements with matching unpack in nvfp4_unshard.
transformer_engine/common/recipe/nvfp4.cu	Per-tensor GEMM scale kernel now accepts per-tensor fp8_max_A/fp8_max_B instead of a hardcoded constant, correctly using 256 for 4over6 tensors and 448 for standard ones.
transformer_engine/pytorch/quantization.py	Correctly maps the nvfp4_4over6 scope string to per-quantizer use_4over6 booleans: all True everywhere, weights True only for weight tensors, activations True for non-weight tensors.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["NVFP4BlockScaling recipe\nnvfp4_4over6: weights/activations/all/None"] --> B["RecipeState._make\nresolves use_4over6 per tensor_type"]
    B --> C["NVFP4Quantizer with use_4over6"]
    C --> D["create_tensor / quantize_impl"]
    D --> E{use_4over6?}
    E -->|No| G["Standard NVFP4 path\n448-based global scale"]
    E -->|Yes| F{valid combo?}
    F -->|RHT or SR or grouped| ERR["NVTE_CHECK error"]
    F -->|OK| H["compute_global_encode_scaling_factor 256-based"]
    H --> I["compute_4over6_decoding_scaling_factors\nmap6 and map4 candidates"]
    I --> J["cvt_fp32_to_fp4_8x_with_mse_rn x2\naccumulate MSE for both candidates"]
    J --> K{err_map4 < err_map6?}
    K -->|Yes| L["rOut_map4 selected"]
    K -->|No - ties go to map6| M["rOut_map6 selected"]
    L --> N["NVFP4Tensor _use_4over6=True\nglobal E4M3 bound = 256"]
    M --> N
    N --> O["Dequant and GEMM scale use 256 denominator"]

Loading

_{Reviews (7): Last reviewed commit: "Drop write back lifting" | Re-trigger Greptile}

[zianglih]

Functionality has been verified by internal RL experiments.
We may want to allow separate 4over6 config for weights and activations, maybe NVTE_NVFP4_ENABLE_4OVER6=weights|activations|all.

[zianglih]

Need to rebase.

[timmoon10]

We are specifying this redundantly in NVTETensor and NVTEQuantizationConfig. If this option can be isolated to quantization, then we should not add clutter to the tensor. If the option is needed for downstream consumers (dequantization, GEMM), then it should be treated as part of the tensor data. I'm not especially familiar, but 4over6 seems like it should be specific to quantization.

[zianglih]

4over6 changes the decode convention from 1 / (6 * 448) to 1 / (6 * 256). Therefore, for our current representation 4over6 is part of the tensor data contract, not just a quantization option.

[timmoon10]

How much benefit does changing the FP8 scale have on convergence? If we don't see a clear benefit, then it would be nicer to use the same scale for 4over6 and non-4over6. That way keep can keep this logic confined to quantization, and downstream consumers are completely unaffected.

If there is an impact on training quality, we should still consider disentangling the FP8 scaling from 4over6. I don't see why other NVFP4 recipes might not benefit from tweaking the scaling.

[zianglih]

From the original paper:

Finally, we make one modification to the computation of the tensor scale α (Equation 1) when
quantizing to NVFP4 with 4/6. When MFP4 ×MFP8 is used to compute the tensor scale, it ensures
that all quantized values will be less than 6 ×448. However, this makes it impossible to select a scale
of 4 for the blocks that contain a tensor’s largest values, because the block’s scale would need to be
448 × 6/4 = 672, which would overflow since 448 is the maximum value that can be represented by
E4M3. As a result, when computing the tensor scale, we replace MFP8 to 256 in Equation 1, since
256 is the largest E4M3 that can be multiplied by 6/4 and represented without error in E4M3, as 384.

Also:

In Section 3.1, we propose calculating the FP32 global tensor scale using 256 as the maximum FP8
E4M3 value rather than the default of 448, as this allows blocks with a tensor’s largest value to have
the option to have a largest FP4 value of 4. In Figure 6, we find that this provides a marginal benefit
over using the standard tensor scale calculation. Even though this adjustment only affects a small
number of large values, this performance gain may come from the fact that larger activation values
can have an outsize impact on model performance. This adjustment is incorporated into the remaining
experiments in this section.

[negvet]

Not sure if there are internal or external studies about the convergence. But this is required to make it work. We need the largest value that is smaller than 448/1.5 and which is itself, and its multiplication by 1.5 is represented by E4M3 exactly. This would help to avoid quantization noise on both map to 4 and map to 6 paths.

[kwyss-nvidia]

We did find the use of 256 to calculate the second level scaling factor helped convergence vs 448, but only slightly.

It's possible that the premise of the paper's argument (prevent saturations when 4 scaling effectively multiplies the block decode scale by 1.5) is sound, but a value larger than 256 can achieve this and the perfect representation of the block with the global amax value with both scalings is not worth the extra range loss.

[zianglih]

let me make 256 scaling a separate env var disabled by default

[zianglih]

448, 320, 288, 256 are all potential candidates for map-to-6:

• 448: effectively disable map-to-4 option above 256, preserve range
• 320, 288: map-to-4 uses 448, no precise 1.5x
• 256: map-to-4 uses 384, precise 1.5x

For now let me refactor the interface to NVTE_NVFP4_4OVER6_E4M3="448"|"256", default to "448" and dispatches to a number in template parameter in C++ code instead of a boolean toggle. People can add support for other values or make it more generic (like directly parsing the env var digits) in the future.

[zianglih]

NVTE_NVFP4_4OVER6_E4M3_USE_256=weights|activations|all is a cleaner pattern and allows separate configuration.

[timmoon10]

This test is okay, but it would provide much more confidence if the NVFP4 quantization tests compared against a CPU reference impl.

[zianglih]

Extended tests/cpp/operator/test_cast_nvfp4_transpose.cu coverage in 3bb42b1.

[negvet]

We need both MSE (better for post-training?) and MAE (better for pre-training as per our internal studies) to be supported, with MAE as the default.

[negvet]

Not sure if there are internal or external studies about the convergence. But this is required to make it work. We need the largest value that is smaller than 448/1.5 and which is itself, and its multiplication by 1.5 is represented by E4M3 exactly. This would help to avoid quantization noise on both map to 4 and map to 6 paths.