Greptile Summary

This PR implements 4over6 block-scale selection for NVFP4 1D quantization: for each 16-element block the kernel computes two scale candidates (map-to-4 and map-to-6), quantizes with both, and picks the one with lower MSE, using a reduced global-scale ceiling of 256 instead of 448 to give the map-4 branch room to represent larger blocks.

• Adds USE_4OVER6 / kUse4Over6 template parameters to both CUDA quantization kernels (quantize_transpose_nvfp4_tuned_1D.cuh and quantize_transpose_vector_blockwise_fp4.cu) and doubles the compile-time dispatch switch depth.
• Threads use_4over6 through the full Python → C++ → CUDA stack (recipe, quantizer, tensor storage, C++ extensions) with guards disabling incompatible modes (RHT, stochastic rounding, 2D quantization, grouped quantization).
• Ships a matching reference Python implementation in NVFP4QuantizerRef and adds parametrised test coverage for the new mode.

Confidence Score: 4/5

The 4over6 logic is well-guarded with incompatible-mode checks at every entry point; the most notable gap is that NVTE_USE_FAST_MATH is wired only into the split-quantization path and has no effect on the standard single-tensor path.

The core quantization math and scale-selection logic look correct and is backed by a reference implementation with tests. The main concern is the NVTE_USE_FAST_MATH env var being silently ignored for single-tensor 4over6 quantization, and the hardcoded [2] array sizes in the CUDA kernel that rely on an implicit loop-bound assumption without a static assert.

transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu (hardcoded array sizing) and transformer_engine/pytorch/csrc/extensions/cast.cpp (fast-math env var only applied in the split path).

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["NVFP4BlockScaling(enable_4over6=True)"] --> B[NVFP4BlockScalingRecipeState]
    B --> C["NVFP4Quantizer(use_4over6=True)"]
    C --> D{quantize path}
    D -->|single tensor| E["quantize_impl (quantizer.cpp)"]
    D -->|split tensor| F["split_quantize_nvfp4_impl_helper (cast.cpp)"]
    E --> G["QuantizationConfig.set_nvfp4_4over6(true)"]
    F --> G
    F -->|reads env var| H["NVTE_USE_FAST_MATH → config.set_use_fast_math"]
    G --> I{kernel dispatch}
    I -->|tuned 1D| J["quantize_transpose_nvfp4_tuned_1D_kernel<USE_4OVER6=true>"]
    I -->|vector blockwise| K["block_scaled_1d_cast_transpose_kernel<kUse4Over6=true>"]
    J --> L["rowwise_scaling: compute map4+map6 scales, pick lower MSE"]
    K --> M["cvt_fp32_to_fp4_8x_with_mse_rn: err_map4 vs err_map6"]
    L --> N["NVFP4Tensor with _use_4over6=True"]
    M --> N

_{Reviews (1): Last reviewed commit: "Initial implementation" | Re-trigger Greptile}