Optimize NVFP4 4over6 candidate error path by zianglih · Pull Request #3068 · NVIDIA/TransformerEngine

zianglih · 2026-06-01T07:15:37Z

Description

@HumansAnd
This PR adds a fast NVFP4 4over6 candidate-error path that compares map-to-4 and map-to-6 candidates in the E4M3-scaled E2M1 product domain after the E2M1 and E4M3 values are rounded through FP16 conversion.

~~Earlier revisions exposed this as two additional public error modes, MAE_FP16 and MSE_FP16.~~ The interface has been refactored so the public 4over6 error mode remains the selection metric, MAE or MSE, while NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 selects the faster FP16 product-domain implementation for that metric.

Fixes # (issue): N/A

Motivation:

The parent 4over6 PR, NVIDIA/TransformerEngine#2972, showed that 4over6 quantization can become compute-bound. In particular, the previous error path is bottlenecked by E2M1 dequantization and FP32 arithmetic instruction count.
Blackwell exposes dedicated PTX conversion instructions for both FP4 E2M1 and FP8 E4M3 into FP16 (cvt.rn.f16x2.e2m1x2 and cvt.rn.f16x2.e4m3x2). This lets the 4over6 kernel construct the candidate E4M3 x E2M1 product with fewer scalar FP32 operations.
E4M3 x E2M1 products can be represented exactly in FP16 for this use case, while BF16 does not have enough mantissa bits for the same guarantee.
The fast error path compares candidates in the E4M3 x E2M1 scaled range. The final error difference and accumulation are still FP32, matching the previous design, but candidate reconstruction uses the Blackwell FP16 conversion/multiply path instead of the heavier FP32 dequantization expression.
We scale the original input into the E4M3-scaled domain instead of fully decoding each candidate back to the original input domain. That applies the FP32 global scaling once to the original input; fully decoding both candidates would require per-element FP32 scaling on each decoded candidate, increasing the arithmetic cost.

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:

Add NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH to select the fast FP16 product-domain candidate-error implementation for NVFP4 4over6.
Keep NVTE_NVFP4_4OVER6_ERR_MODE focused on the selection metric, MAE or MSE, rather than encoding implementation details in additional public modes.
Thread the 4over6 error fast-math option through NVTEQuantizationConfig as a dedicated boolean config.
Add the FP16 candidate-error path to the NVFP4 4over6 CUDA kernel using Blackwell FP16 conversion PTX with the rn modifier and FP16 multiply.
Extend the PyTorch reference implementation with an emulated FP16 candidate path. PyTorch does not expose this exact FP16 round-to-nearest conversion/multiply sequence, so the reference reconstructs the FP16 result from integer fields and checks bitwise exactness against the kernel.
Extend NVFP4 exact quantization tests to cover 4over6 MAE / MSE, E4M3 max 448 / 256, and error fast-math enabled / disabled.
Update the environment variable docs for NVTE_NVFP4_4OVER6_ERR_MODE and NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH.

Testing note:

The C++ operator test was intentionally not expanded for the FP16 product-domain path. That test suite checks that each 4over6 block matches either map-to-4 or map-to-6 exactly; it does not validate candidate selection. Since this PR changes error computation and candidate selection, the strict PyTorch reference tests provide the meaningful coverage.
Although this path is controlled by NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH, the tests still require bitwise exactness against the PyTorch reference implementation. The reference emulates the FP16 conversion/multiply sequence that PyTorch does not expose directly, and the exact quantization tests use zero tolerance.

Fast-vs-strict agreement check:

To quantify the review concern about how much NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 differs from the regular path, I compared emitted rowwise E4M3 scale bytes with fast error math disabled vs enabled on random M, K = 98304, 7168. The agreement rate is consistently above 99.9% across tensor-scaled 1D, tensor-scaled 2D, row-scaled 1D, E4M3 max 256 / 448, BF16 / FP16 inputs, and MAE / MSE selection modes.

scale	quant	e4m3 max	dtype	err	same E4M3 %	different E4M3	total E4M3
tensor	1d	256	bf16	MAE	99.980177	8730	44040192
tensor	1d	256	bf16	MSE	99.993304	2949	44040192
tensor	1d	256	fp16	MAE	99.995965	1777	44040192
tensor	1d	256	fp16	MSE	99.998940	467	44040192
tensor	1d	448	bf16	MAE	99.931292	30259	44040192
tensor	1d	448	bf16	MSE	99.977966	9704	44040192
tensor	1d	448	fp16	MAE	99.998719	564	44040192
tensor	1d	448	fp16	MSE	99.999782	96	44040192
tensor	2d	256	bf16	MAE	99.999964	16	44040192
tensor	2d	256	bf16	MSE	99.999782	96	44040192
tensor	2d	256	fp16	MAE	100.000000	0	44040192
tensor	2d	256	fp16	MSE	99.999927	32	44040192
tensor	2d	448	bf16	MAE	99.999964	16	44040192
tensor	2d	448	bf16	MSE	99.999455	240	44040192
tensor	2d	448	fp16	MAE	100.000000	0	44040192
tensor	2d	448	fp16	MSE	100.000000	0	44040192
row	1d	256	bf16	MAE	99.977505	9907	44040192
row	1d	256	bf16	MSE	99.993213	2989	44040192
row	1d	256	fp16	MAE	99.996885	1372	44040192
row	1d	256	fp16	MSE	99.999255	328	44040192
row	1d	448	bf16	MAE	99.983867	7105	44040192
row	1d	448	bf16	MSE	99.996428	1573	44040192
row	1d	448	fp16	MAE	99.997593	1060	44040192
row	1d	448	fp16	MSE	99.999505	218	44040192

Validation:

python3 -m pytest --tb=auto tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py -k 4over6
# 1824 passed, 1632 skipped, 432 deselected, 2 warnings

pre-commit run --all-files
# Passed

Performance:

No backend implementation changed in the interface refactor, so the benchmark values below are unchanged from the original FP16 error-path implementation.

Standalone quantization-kernel benchmark:

I also ran a local direct-quantizer benchmark that times NVFP4Quantizer.update_quantized with CUDA events, isolating quantization-kernel latency from grouped-linear dispatch and GEMM work.

python3 -u benchmarks/benchmark_4over6.py --shapes profile --warmup 20 --iters 2000

Shape: (16384, 6144). Slowdown is relative to the current non-4over6 NVFP4 kernel in the same 1D or 2D mode.

mode	err	non-4over6 us	strict 4over6 us	strict slowdown	fast 4over6 us	fast slowdown	fast vs strict speedup
1d	MAE	103.259	800.892	7.756x	289.100	2.800x	2.770x
1d	MSE	103.259	829.693	8.035x	289.148	2.800x	2.869x
2d	MAE	126.737	847.070	6.684x	306.012	2.415x	2.768x
2d	MSE	126.737	867.081	6.842x	305.842	2.413x	2.835x

The direct kernel numbers show the main effect of this PR clearly: the fast 4over6 error path substantially reduces quantization-kernel overhead compared with strict original-domain error compute, while preserving bitwise exactness against the PyTorch reference.

Commands for the 2D activation sweep:

# Baseline, no 4over6.
NVTE_NVFP4_DISABLE_RHT=1 \
NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1 \
python3 benchmarks/linear/benchmark_grouped_linear.py --recipe nvfp4

# 4over6 activation modes.
NVTE_NVFP4_4OVER6=activations \
NVTE_NVFP4_4OVER6_ERR_MODE=<MAE|MSE> \
NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=<0|1> \
NVTE_NVFP4_DISABLE_RHT=1 \
NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1 \
python3 benchmarks/linear/benchmark_grouped_linear.py --recipe nvfp4

Raw 2D activation grouped forward/backward timings, in ms per microbatch:

err_mode	err_use_fast_math	m	k	n	recipe	num_gemms	grouped_fwd_bwd_time_ms
baseline	-	16384	7168	2048	nvfp4	4	0.768440
baseline	-	32768	7168	2048	nvfp4	4	1.246045
baseline	-	65536	7168	2048	nvfp4	4	2.226334
baseline	-	98304	7168	2048	nvfp4	4	3.220651
baseline	-	16384	7168	2048	nvfp4	8	0.999235
baseline	-	32768	7168	2048	nvfp4	8	1.428313
baseline	-	65536	7168	2048	nvfp4	8	2.400536
baseline	-	98304	7168	2048	nvfp4	8	3.387845
MAE	0	16384	7168	2048	nvfp4	4	1.638763
MAE	0	32768	7168	2048	nvfp4	4	2.870079
MAE	0	65536	7168	2048	nvfp4	4	5.330869
MAE	0	98304	7168	2048	nvfp4	4	7.818854
MAE	0	16384	7168	2048	nvfp4	8	1.914825
MAE	0	32768	7168	2048	nvfp4	8	3.168478
MAE	0	65536	7168	2048	nvfp4	8	5.612699
MAE	0	98304	7168	2048	nvfp4	8	8.076360
MSE	0	16384	7168	2048	nvfp4	4	1.672350
MSE	0	32768	7168	2048	nvfp4	4	2.936419
MSE	0	65536	7168	2048	nvfp4	4	5.464571
MSE	0	98304	7168	2048	nvfp4	4	8.020864
MSE	0	16384	7168	2048	nvfp4	8	1.944005
MSE	0	32768	7168	2048	nvfp4	8	3.234871
MSE	0	65536	7168	2048	nvfp4	8	5.745626
MSE	0	98304	7168	2048	nvfp4	8	8.274712
MAE	1	16384	7168	2048	nvfp4	4	0.979436
MAE	1	32768	7168	2048	nvfp4	4	1.648099
MAE	1	65536	7168	2048	nvfp4	4	2.966817
MAE	1	98304	7168	2048	nvfp4	4	4.395422
MAE	1	16384	7168	2048	nvfp4	8	1.176896
MAE	1	32768	7168	2048	nvfp4	8	1.848563
MAE	1	65536	7168	2048	nvfp4	8	3.148987
MAE	1	98304	7168	2048	nvfp4	8	4.552140
MSE	1	16384	7168	2048	nvfp4	4	0.978000
MSE	1	32768	7168	2048	nvfp4	4	1.644019
MSE	1	65536	7168	2048	nvfp4	4	2.974658
MSE	1	98304	7168	2048	nvfp4	4	4.401849
MSE	1	16384	7168	2048	nvfp4	8	1.176988
MSE	1	32768	7168	2048	nvfp4	8	1.846523
MSE	1	65536	7168	2048	nvfp4	8	3.150669
MSE	1	98304	7168	2048	nvfp4	8	4.554252

2D activation slowdown relative to baseline:

m	k	n	recipe	num_gemms	MAE, fast=0	MSE, fast=0	MAE, fast=1	MSE, fast=1	fast vs strict speedup
16384	7168	2048	nvfp4	4	2.133x	2.176x	1.275x	1.273x	MAE 1.673x / MSE 1.709x
32768	7168	2048	nvfp4	4	2.303x	2.357x	1.323x	1.319x	MAE 1.741x / MSE 1.787x
65536	7168	2048	nvfp4	4	2.394x	2.455x	1.333x	1.336x	MAE 1.796x / MSE 1.838x
98304	7168	2048	nvfp4	4	2.428x	2.490x	1.365x	1.367x	MAE 1.779x / MSE 1.822x
16384	7168	2048	nvfp4	8	1.916x	1.945x	1.178x	1.178x	MAE 1.626x / MSE 1.651x
32768	7168	2048	nvfp4	8	2.218x	2.265x	1.294x	1.293x	MAE 1.714x / MSE 1.752x
65536	7168	2048	nvfp4	8	2.338x	2.393x	1.312x	1.312x	MAE 1.782x / MSE 1.824x
98304	7168	2048	nvfp4	8	2.384x	2.442x	1.344x	1.344x	MAE 1.774x / MSE 1.817x

The fast error path consistently reduces 4over6 overhead compared with the default MAE / MSE original-domain error path in this sweep. We also see the same speedup trend on another grouped-linear NVFP4 recipe.

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

Signed-off-by: Ziang Li <ziangli@umich.edu>

greptile-apps · 2026-06-01T07:28:38Z

Greptile Summary

This PR adds a fast NVFP4 4over6 candidate-error path controlled by NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1. Instead of decoding candidates fully back to the original input domain in FP32, the new path uses Blackwell PTX (cvt.rn.f16x2.e2m1x2, cvt.rn.f16x2.e4m3x2, mul.rn.f16x2) to compute the E2M1 × E4M3 product in FP16 and compare candidates in the E4M3-scaled domain, roughly halving 4over6 overhead.

CUDA kernel (quantize_4over6_nvfp4.cuh): introduces FP16ErrorScalePair, compute_fp16_error_scales, f16x2_scaled_to_float2, and accumulate_fp16_scaled_error_pair; removes the old non-rn fast_math arithmetic branch from accumulate_dequant_error and folds it cleanly into a compile-time if constexpr (Cfg::err_use_fast_math) dispatch.
Python reference (quantization_ref_nvfp4.py): emulates the FP16 multiplication sequence via integer bit manipulation in _ref_nvfp4_4over6_fp16_candidate; adds _sum_4over6_2d_error with an assert tile_len_y == 16 guard; threads nvfp4_4over6_err_use_fast_math through the quantizer constructor and all call sites.
Tests (test_nvfp4_quantize_exact.py): consolidates parametrization into a NVFP44Over6TestConfig dataclass; covers MAE/MSE × e4m3_max=448/256 × fast_math on/off; uses an env-var context manager that correctly saves and restores NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH.

Confidence Score: 5/5

Safe to merge; all paths verified correct and covered by zero-tolerance bitwise tests.

The PTX register layout in compute_fp16_error_scales is verified: map4 is packed into bits [7:0] of fp8_pair, cvt.rn.f16x2.e4m3x2 converts lower→lower/upper→upper, and the two mov.b32 replications correctly broadcast each FP16 scale into an f16x2 for mul.rn.f16x2. The Python reference _ref_nvfp4_4over6_fp16_candidate correctly decomposes E2M1 and E4M3 bit fields, and the subnormal FP16 branch is dead code in practice (minimum product exponent is -10, well above the -14 threshold). All error accumulation differences and the tree-reduction ordering match between CUDA and Python. The 1824-passed, zero-tolerance test run provides strong bitwise correctness evidence.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh	Adds FP16 fast-error path using Blackwell PTX (cvt.rn.f16x2.e4m3x2, cvt.rn.f16x2.e2m1x2, mul.rn.f16x2); removes the old non-rn fast_math branch from accumulate_dequant_error; bit-packing and PTX register layout verified correct.
transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py	Adds _ref_nvfp4_4over6_fp16_candidate and _sum_4over6_2d_error helpers; threads nvfp4_4over6_err_use_fast_math through the reference quantizer; NaN E4M3 guard is dead code in practice (scale values are always ≤ 0x7E/448).
tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py	Replaces three separate parametrize decorators with a consolidated NVFP44Over6TestConfig dataclass; adds environment-variable context manager; correctly covers MAE/MSE × e4m3_max=448/256 × fast_math=on/off combinations.
docs/envvars.rst	Updates NVTE_NVFP4_4OVER6_ERR_MODE and NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH descriptions to accurately reflect the refactored interface.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["cvt_fp32_to_fp4_8x_with_error()"] --> B{err_use_fast_math?}
B -->|false| C["accumulate_dequant_error()\nFP32 decode: dequant x sf x global_amax / denom\n__fdiv_rn / __fmul_rn / __fsub_rn"]
B -->|true| D["accumulate_fp16_scaled_error_pair()\nFP16: mul.rn.f16x2(q_h2, scale_h2)\nthen cvt.f32.f16 to FP32 diff"]
D --> E["f16x2_scaled_to_float2()\nmul.rn.f16x2 then cvt.f32.f16 x2"]
D --> F["original = x x global_encode_scale\n(__fmul_rn)"]
E --> G["diff = candidate minus original"]
F --> G
G --> H["compute_error_rn: abs diff or diff squared"]
C --> H
H --> I["FP32 error accumulation\n(__fadd_rn)"]

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

greptile-apps · 2026-06-01T07:28:42Z

+ def _sum_4over6_2d_error(err: torch.Tensor, tile_len_y: int) -> torch.Tensor:
+ """Reduce 16 row errors in the same tree order as the CUDA warp reduction."""
+ rows = err.view(err.shape[0] // tile_len_y, tile_len_y, err.shape[1], 1)
+ rows = rows.squeeze(-1)
+ rows = rows[:, 0:8, :] + rows[:, 8:16, :]
+ rows = rows[:, 0:4, :] + rows[:, 4:8, :]
+ rows = rows[:, 0:2, :] + rows[:, 2:4, :]
+ return (rows[:, 0, :] + rows[:, 1, :]).unsqueeze(-1)


Hardcoded 16-row tree reduction without enforcement
The function accepts tile_len_y as a parameter and uses it for the initial view, but the subsequent binary-tree steps unconditionally slice at [8:16], [4:8], [2:4], [0:2], which are only correct when tile_len_y == 16. A caller passing tile_len_y = 32 would silently sum only the first 16 rows and discard the remaining 16; a caller passing tile_len_y = 8 would produce a shape mismatch on the first add. Adding an assert tile_len_y == 16 at the top would make the contract explicit and catch future regressions immediately.

Addressed in 5136b58. _sum_4over6_2d_error now asserts tile_len_y == 16 before the tree reduction, and I also updated the reference docstring/comment to distinguish original-domain MAE/MSE from E4M3-scaled FP16 product-domain MAE_FP16/MSE_FP16.

Signed-off-by: Ziang Li <ziangli@umich.edu>

zianglih · 2026-06-02T01:58:39Z

Since this is identical to the orginal mode in the algebraic sense, may be a better interface design it to let this mode replace the original fast math modes, instead of extending error modes.

Signed-off-by: Ziang Li <ziangli@umich.edu>

zianglih · 2026-06-02T05:36:19Z

Interface refactor update:

I refactored the core 4over6 error interface so the selection metric and implementation fast path are separate knobs:

NVTE_NVFP4_4OVER6_ERR_MODE remains the candidate-selection metric and only accepts MAE / MSE.
NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 enables the FP16 product-domain error path for the selected metric.
The C++ quantization config mirrors this split with NVTENVFP44Over6Mode for disabled / min-MAE / min-MSE, plus a separate nvfp4_4over6_err_use_fast_math boolean.

This replaces the earlier public MAE_FP16 / MSE_FP16 mode design. The backend FP16 error implementation and benchmark numbers did not change in this interface refactor; the PR body now relabels the commands and tables around ERR_MODE + ERR_USE_FAST_MATH instead.

zianglih · 2026-06-02T05:50:57Z

@greptile-apps I have updated PR body, review again.

zianglih · 2026-06-02T12:13:53Z

FlashInfer PR that implements the same contract:

Add CuTe DSL NVFP4 quantization with 4over6 FP16 scoring flashinfer-ai/flashinfer#3448

ptrendx · 2026-06-02T21:00:24Z

Do you have any data on how different the result is between the regular path and this new fast path? Maybe we could just make it a default path rather than introduce another env variable?

zianglih · 2026-06-02T21:27:39Z

Hi @ptrendx ,

I have some results using this script:

"""Compare NVFP4 4over6 E4M3 scales with and without error fast math."""importosfromcontextlibimportcontextmanagerimporttorchfromtransformer_engine.pytorchimportNVFP4Quantizerimporttransformer_engine_torchastexM, K=98304, 7168@contextmanagerdef_error_fast_math(enabled: bool):
old_value=os.environ.get("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH")
os.environ["NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"] ="1"ifenabledelse"0"try:
yieldfinally:
ifold_valueisNone:
os.environ.pop("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", None)
else:
os.environ["NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"] =old_valuedef_quantize_scale_bytes(
x: torch.Tensor,
err_mode: str,
err_fast_math: bool,
row_scaled: bool,
with_2d_quantization: bool,
nvfp4_e4m3_max: int,
) ->torch.Tensor:
quantizer=NVFP4Quantizer(
fp4_dtype=tex.DType.kFloat4E2M1,
rowwise=True,
columnwise=False,
with_amax_reduction=False,
amax_reduction_group=None,
with_rht=False,
with_post_rht_amax=False,
with_2d_quantization=with_2d_quantization,
row_scaled_nvfp4=row_scaled,
nvfp4_use_4over6=True,
nvfp4_e4m3_max=nvfp4_e4m3_max,
nvfp4_4over6_err_mode=err_mode,
)
with_error_fast_math(err_fast_math):
quantized=quantizer(x)
assertquantized._rowwise_scale_invisnotNonereturnquantized._rowwise_scale_inv.contiguous().view(torch.uint8)
def_compare_e4m3(
x: torch.Tensor,
dtype_name: str,
scale_mode: str,
row_scaled: bool,
quant_mode: str,
with_2d_quantization: bool,
nvfp4_e4m3_max: int,
) ->None:
forerr_modein ("MAE", "MSE"):
regular=_quantize_scale_bytes(
x, err_mode, False, row_scaled, with_2d_quantization, nvfp4_e4m3_max
)
fast=_quantize_scale_bytes(
x, err_mode, True, row_scaled, with_2d_quantization, nvfp4_e4m3_max
)
same=torch.count_nonzero(regular==fast).item()
total=regular.numel()
print(
f"{scale_mode:>6}{quant_mode:>5}{nvfp4_e4m3_max:8d} "f"{dtype_name:>5}{err_mode:>3} "f"{100.0*same/total:12.6f}{total-same:15d}{total}"
)
defmain():
torch.set_grad_enabled(False)
print(f"shape=({M}, {K}), 1d_e4m3_values={M*K//16}")
print("scale quant e4m3_max dtype mode same_e4m3_pct different_e4m3 total_e4m3")
forscale_mode, row_scaled, quant_mode, with_2d_quantizationin (
("tensor", False, "1d", False),
("tensor", False, "2d", True),
("row", True, "1d", False),
):
fornvfp4_e4m3_maxin (256, 448):
fordtype, dtype_namein ((torch.bfloat16, "bf16"), (torch.float16, "fp16")):
torch.manual_seed(1234)
x=torch.randn((M, K), dtype=dtype, device="cuda")
_compare_e4m3(
x,
dtype_name,
scale_mode,
row_scaled,
quant_mode,
with_2d_quantization,
nvfp4_e4m3_max,
)
delxtorch.cuda.empty_cache()
if__name__=="__main__":
main()

The agreement rate is consistently >99.9% on random M, K = 98304, 7168:

shape=(98304, 7168), 1d_e4m3_values=44040192
scale quant e4m3_max dtype mode same_e4m3_pct different_e4m3 total_e4m3
tensor 1d 256 bf16 MAE 99.980177 8730 44040192
tensor 1d 256 bf16 MSE 99.993304 2949 44040192
tensor 1d 256 fp16 MAE 99.995965 1777 44040192
tensor 1d 256 fp16 MSE 99.998940 467 44040192
tensor 1d 448 bf16 MAE 99.931292 30259 44040192
tensor 1d 448 bf16 MSE 99.977966 9704 44040192
tensor 1d 448 fp16 MAE 99.998719 564 44040192
tensor 1d 448 fp16 MSE 99.999782 96 44040192
tensor 2d 256 bf16 MAE 99.999964 16 44040192
tensor 2d 256 bf16 MSE 99.999782 96 44040192
tensor 2d 256 fp16 MAE 100.000000 0 44040192
tensor 2d 256 fp16 MSE 99.999927 32 44040192
tensor 2d 448 bf16 MAE 99.999964 16 44040192
tensor 2d 448 bf16 MSE 99.999455 240 44040192
tensor 2d 448 fp16 MAE 100.000000 0 44040192
tensor 2d 448 fp16 MSE 100.000000 0 44040192
row 1d 256 bf16 MAE 99.977505 9907 44040192
row 1d 256 bf16 MSE 99.993213 2989 44040192
row 1d 256 fp16 MAE 99.996885 1372 44040192
row 1d 256 fp16 MSE 99.999255 328 44040192
row 1d 448 bf16 MAE 99.983867 7105 44040192
row 1d 448 bf16 MSE 99.996428 1573 44040192
row 1d 448 fp16 MAE 99.997593 1060 44040192
row 1d 448 fp16 MSE 99.999505 218 44040192

We did not introduce any extra env var in this PR. NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH was already there previously but for controlling rounding modifiers in dequant arithmetic instrcutions (which did not lead to noticeable speedup). We just replace that backend with a new implementation.

I do think we can make NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH =1 the default but I am not sure if we should make this specific contract the only kept 4over6 implementation. Also in the future we may land other performance improvements which may change the contract numerics. I think it is better to support the canonical 4over6 contract from the original paper by NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=0, and use NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH =1 to land this contract and potential future improvements.

ptrendx · 2026-06-02T23:04:16Z

-template <typename Cfg, int E4M3_MAX>
-__device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error(const float (&x)[8],
- const float block_scale_inverse,
+__device__ __forceinline__ float2 e2m1x2_scaled_e4m3_to_float2(const uint32_t e2m1_byte,


Why do we only compute 2 values in this function? Also, we don't use half of the scale_h2 here. Why don't we instead try to convert here values from both of the branches (so both 4 and 6 would be there, the scaling factors for both of these branches would be converted in a single instruction). Ideally we would then reuse those scaling factors rather than recasting them for every element in a block - considering we are math bound here, we need to make sure that we eliminate as many instructions as possible.

I tried a refactored implementation in 54797b3 but it did not show meaningful speedup:
NVTE_NVFP4_4OVER6=activations \ NVTE_NVFP4_4OVER6_ERR_MODE=<MAE|MSE> \ NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 \ NVTE_NVFP4_DISABLE_RHT=1 \ NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1 \ python3 benchmarks/linear/benchmark_grouped_linear.py --recipe nvfp4
Extended Fast-Path Timing Table:
m k n recipe num_gemms baseline_ms old_MAE_fast1_ms refactor_MAE_fast1_ms MAE_refactor_speedup old_MSE_fast1_ms refactor_MSE_fast1_ms MSE_refactor_speedup
16384 7168 2048 nvfp4 4 0.768440 0.979436 0.984238 0.995x 0.978000 0.984061 0.994x
32768 7168 2048 nvfp4 4 1.246045 1.648099 1.643575 1.003x 1.644019 1.645252 0.999x
65536 7168 2048 nvfp4 4 2.226334 2.966817 2.977467 0.996x 2.974658 2.990481 0.995x
98304 7168 2048 nvfp4 4 3.220651 4.395422 4.396148 1.000x 4.401849 4.409521 0.998x
16384 7168 2048 nvfp4 8 0.999235 1.176896 1.182225 0.995x 1.176988 1.258703 0.935x
32768 7168 2048 nvfp4 8 1.428313 1.848563 1.853778 0.997x 1.846523 1.857125 0.994x
65536 7168 2048 nvfp4 8 2.400536 3.148987 3.152150 0.999x 3.150669 3.153690 0.999x
98304 7168 2048 nvfp4 8 3.387845 4.552140 4.549876 1.000x 4.554252 4.560287 0.999x
the refactor-vs-old geomean was:
MAE: 0.9983x MSE: 0.9889x combined: 0.9936x
This refactor is essentially common instruction lifting/reuse, and it keeps the same core PTX instructions (cvt.rn.f16x2.e4m3x2, cvt.rn.f16x2.e2m1x2, mul.rn.f16x2) rather than introducing a different PTX operation. I think the compiler can already do this in the old implementation but I am not sure.

Kernel level benchmark does not show speedup either:
mode kernel err metric old us refactor us old/refactor
1d nvfp4 - strict 103.260 103.259 1.000x
1d 4over6 MAE strict 799.852 800.892 0.999x
1d 4over6 MAE fast 288.834 289.100 0.999x
1d 4over6 MSE strict 829.294 829.693 1.000x
1d 4over6 MSE fast 287.426 289.148 0.994x
2d nvfp4 - strict 126.692 126.737 1.000x
2d 4over6 MAE strict 847.700 847.070 1.001x
2d 4over6 MAE fast 306.286 306.012 1.001x
2d 4over6 MSE strict 866.949 867.081 1.000x
2d 4over6 MSE fast 306.748 305.842 1.003x
Script: 83e2308 , shape (16384, 6144), with --warmup 20 --iters 2000

I keep the refactoring in beaed67 .The perf behavior of the explicit pattern is more robust to compiler optimizations.

Signed-off-by: Ziang Li <ziangli@umich.edu>

This reverts commit 54797b3. Signed-off-by: Ziang Li <ziangli@umich.edu>

Signed-off-by: Ziang Li <ziangli@umich.edu>

## 📌 Description @HumansAnd Adds a CuTe DSL NVFP4 quantization backend with support for per-token activation quantization and NVFP4 4over6 scale-candidate selection. The 4over6 configuration remains environment-driven to match the existing kernel configuration flow: - `FLASHINFER_NVFP4_4OVER6=1` - `FLASHINFER_NVFP4_4OVER6_ERR_MODE=MAE|MSE` - `FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH=0|1` - `FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=0|1` - `FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=0|1` - C++ dispatch also accepts `TRTLLM_DISABLE_FP4_QUANT_FAST_MATH=1` as a compatibility fallback. Python/CuTe DSL paths use the FlashInfer env name. The public 4over6 error-mode surface is now only `MAE` and `MSE`. The former FP16-specific public modes were removed. Instead, `FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH` selects the candidate-error contract: - `0`: strict full-dequant candidate scoring with explicit RN arithmetic. - Per-token NVFP4 uses the online row amax and scores candidates with the TE-style expression `(e2m1 * sf * global_amax) / (6 * e4m3_max)`. - Per-tensor NVFP4 uses the global decode scale path derived from the provided global scale, rather than materializing a separate global amax in the Python API. This avoids adding an always-on input reduction for per-tensor callers. - CUDA and CuTe DSL follow the same per-token/per-tensor split. - `1`: FP16-domain candidate scoring. Candidates are decoded through the E2M1 x E4M3 FP16 conversion path and compared against `x * global_encode_scale`; this path also uses explicit RN arithmetic for candidate error scoring. ~~Earlier revisions treated the strict full-dequant path as a single TE-style expression for both per-token and per-tensor NVFP4.~~ The current contract only requires the TE expression for the per-token path where amax is already computed online; per-tensor uses the existing global decode scale expression to avoid extra overhead. CUDA, CuTe DSL, and fused MoE dispatch all route through the shared `NVFP44Over6Config<e4m3_max, err_mode, err_use_fast_math>` configuration. The old 4over6 candidate-error paths that used unqualified fast-math arithmetic were removed. ## 🔍 Related Issues TE PR that implements the same contract: - NVIDIA/TransformerEngine#3068 Other relevant PRs: - #3387 - NVIDIA/TransformerEngine#2972 - #3264 - #3027 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Unit tests extended: - `tests/test_helpers/utils_fp4.py` extends the TE-style NVFP4/4over6 reference path for both full-dequant scoring and FP16-domain scoring. - `tests/utils/test_fp4_quantize.py` extends `test_nvfp4_quantize_te_reference` to cover strict bitwise CUDA and CuTe DSL agreement for NVFP4, per-token activation, scale-factor layouts, zero inputs, and 4over6 MAE/MSE configurations. Commands run: ```bash pre-commit run --all-files ``` Passed. ```bash python3 -m compileall \ flashinfer/quantization/nvfp4_quantization_utils.py \ tests/test_helpers/utils_fp4.py ``` Passed. ```bash CUDA_VISIBLE_DEVICES=0 \ python3 -m pytest -q --tb=short tests/utils/test_fp4_quantize.py::test_nvfp4_quantize_te_reference ``` Passed: `9504 passed in 756.79s (0:12:36)`. ```bash CUDA_VISIBLE_DEVICES=0 \ python3 -m pytest -q --tb=short tests/moe/test_trtllm_cutlass_fused_moe.py ``` Passed: `75 passed, 25 skipped, 1 warning in 74.79s (0:01:14)`. ```bash CUDA_VISIBLE_DEVICES=0 \ python3 -m pytest -q --tb=short tests/moe/test_trtllm_gen_per_token_moe.py ``` Passed: `108 passed in 785.08s (0:13:05)`. ## Benchmarking Updated benchmark run compares PR head (`b44e847b`) against the earliest PR commit with CUDA 4over6 support (`e30a0b14`). The denominator for every 4over6 speedup is the baseline CUDA 4over6 MSE no-fast time for the same `(M, K)` case. Common config: - dtype: `bfloat16` - scale-factor layout: `swizzled_128x4` - `FLASHINFER_NVFP4_4OVER6=1` - `FLASHINFER_NVFP4_4OVER6_ERR_MODE=MSE` - `FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=1` - `FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1` - CUDA graph disabled, CUPTI timing enabled, cold L2 cache enabled Per-token runs add `--per-token-activation`; per-tensor runs omit it. `FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH` is set to `0` for no-fast rows and `1` for fast rows. Each table cell is `geomean / min / max` speedup over the scanned M/K cases. | Activation scale | CUDA no fast math | CUDA fast math | CuTe DSL no fast math | CuTe DSL fast math | | --- | ---: | ---: | ---: | ---: | | Per-tensor | 2.83x / 1.65x / 5.81x | 3.22x / 1.67x / 7.25x | 2.17x / 1.21x / 8.26x | 4.12x / 2.45x / 20.92x | | Per-token | 0.89x / 0.71x / 0.98x | 1.41x / 1.05x / 3.40x | 2.05x / 1.03x / 2.89x | 3.00x / 1.80x / 4.17x | Pure per-token NVFP4 without 4over6, same dtype/layout/quant-fast-math setting: | Mode | CuTe DSL vs CUDA geomean/min/max | | --- | ---: | | Pure per-token no 4over6 | 2.24x / 1.04x / 3.76x | Heat maps: ![NVFP4 4over6 MSE per-tensor speedup heat map](https://gist.githubusercontent.com/zianglih/1e1fdd42d27244692d66cd6a5b2b904f/raw/2f4f7965c0b57e7b148acdde64a632daa11072a6/nvfp4_4over6_mse_speedup_vs_baseline_per_tensor.svg) ![NVFP4 4over6 MSE per-token speedup heat map](https://gist.githubusercontent.com/zianglih/1e1fdd42d27244692d66cd6a5b2b904f/raw/23fc9c84605795b9c17febac47f0175fc14c3a07/nvfp4_4over6_mse_speedup_vs_baseline_per_token.svg) ![NVFP4 pure per-token no-4over6 CuTe DSL vs CUDA heat map](https://gist.githubusercontent.com/zianglih/1e1fdd42d27244692d66cd6a5b2b904f/raw/34e161c89dc0a23378c5d23041567953b8f7c9b2/nvfp4_pure_per_token_no_4over6_cute_vs_cuda.svg) Previous benchmark snapshot kept for review context: | Activation scale | Baseline CUDA no-fast geomean (ms) | CUDA no-fast | CUDA fast | CuTe DSL no-fast | CuTe DSL fast | | --- | ---: | ---: | ---: | ---: | ---: | | Per-tensor | ~~0.074336~~ | ~~1.000x~~ | ~~1.668x~~ | ~~1.434x~~ | ~~1.738x~~ | | Per-token | ~~0.065345~~ | ~~0.974x~~ | ~~1.554x~~ | ~~2.297x~~ | ~~3.267x~~ | - ~~Previous benchmark run compared current PR head (`e5349fc2`) against a pre-FP16-refactor baseline (`99a21fcd`).~~ - ~~Previous common config used `FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=0`.~~ - ~~Previous table reported only geomean speedup.~~ Full 22x15 per-case data was generated separately and kept out of the PR body. ## Reviewer Notes The 4over6 path can become compute-bound at large inputs. The benchmark path keeps the built-in sweep and uses environment variables for 4over6 mode selection rather than adding 4over6-specific shape or mode CLI arguments.  ## Summary by CodeRabbit * **New Features** * Added 4over6 quantization mode for NVFP4 with configurable error metrics (MAE/MSE) * Added per-token activation support for NVFP4 quantization * Extended CuTe-DSL backend capabilities for improved quantization operations * **Documentation** * Updated quantization benchmarks with new 4over6 mode configuration examples  --------- Signed-off-by: Ziang Li <ziangli@umich.edu>

timmoon10 · 2026-06-08T19:30:33Z

/te-ci

zianglih · 2026-06-08T22:42:30Z

pytorch failure is attention:

=========================== short test summary info ============================
FAILED ../../tests/pytorch/attention/test_attention.py::test_dpa_fp8_vs_f16[mxfp8-True-True-bshd_bshd_bshd-fp8_9-dtype1] - AssertionError: fused_attn_bwd_fp8[0] vs fused_attn_bwd_f16[0] RMSE nan is over tolerance nan (0.11000 * nan)
FAILED ../../tests/pytorch/attention/test_attention.py::test_dpa_fp8_vs_f16[mxfp8-True-True-bshd_bshd_bshd-fp8_11-dtype1] - AssertionError: fused_attn_bwd_fp8[0] vs fused_attn_bwd_f16[0] RMSE nan is over tolerance nan (0.11000 * nan)
FAILED ../../tests/pytorch/attention/test_attention.py::test_dpa_fp8_vs_f16[mxfp8-True-True-sbhd_sbhd_sbhd-fp8_9-dtype1] - AssertionError: fused_attn_bwd_fp8[0] vs fused_attn_bwd_f16[0] RMSE nan is over tolerance nan (0.11000 * nan)
==== 3 failed, 2197 passed, 1554 skipped, 31 warnings in 971.42s (0:16:11) =====
Error: sub-test failed: NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py

This PR does not touch JAX

timmoon10

LGTM

zianglih added 2 commits May 31, 2026 23:20

Add FP16 NVFP4 4over6 error modes
f151dd6
Signed-off-by: Ziang Li <ziangli@umich.edu>

Clean up NVFP4 4over6 FP16 mode changes
71ba85e
Signed-off-by: Ziang Li <ziangli@umich.edu>

github-actionsBot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Jun 1, 2026

zianglih marked this pull request as ready for review June 1, 2026 07:19

zianglih requested review from Oleg-Goncharov, ksivaman and ptrendx as code owners June 1, 2026 07:19

greptile-appsBot reviewed Jun 1, 2026
View reviewed changes

Clarify NVFP4 4over6 reference error modes
5136b58
Signed-off-by: Ziang Li <ziangli@umich.edu>

zianglih added 4 commits June 1, 2026 22:08

Refactor NVFP4 4over6 error mode interface
0a6e0d6
Signed-off-by: Ziang Li <ziangli@umich.edu>

Clean up NVFP4 4over6 exact test configs
25861a1
Signed-off-by: Ziang Li <ziangli@umich.edu>

Lift NVFP4 exact test quantizer construction
6fce6b9
Signed-off-by: Ziang Li <ziangli@umich.edu>

Clean up doc changes
f47fb7b
Signed-off-by: Ziang Li <ziangli@umich.edu>

zianglih changed the title ~~Add FP16 error modes for NVFP4 4over6~~Add fast FP16 error path for NVFP4 4over6Jun 2, 2026

zianglih changed the title ~~Add fast FP16 error path for NVFP4 4over6~~Optimize NVFP4 4over6 candidate error pathJun 2, 2026

zianglih mentioned this pull request Jun 2, 2026
Add CuTe DSL NVFP4 quantization with 4over6 FP16 scoring flashinfer-ai/flashinfer#3448
Merged
5 tasks

ptrendx reviewed Jun 2, 2026
View reviewed changes

zianglih added 4 commits June 2, 2026 16:36

Add script and experimental refactor
54797b3
Signed-off-by: Ziang Li <ziangli@umich.edu>

Revert "Add script and experimental refactor"
8091a59
This reverts commit 54797b3. Signed-off-by: Ziang Li <ziangli@umich.edu>

Preserve explicit common instruction lifting
beaed67
Signed-off-by: Ziang Li <ziangli@umich.edu>

Drop scripts
a7d1a6c
Signed-off-by: Ziang Li <ziangli@umich.edu>

timmoon10 approved these changes Jun 9, 2026
View reviewed changes

timmoon10 merged commit b972fa8 into NVIDIA:mainJun 9, 2026
29 of 33 checks passed

zianglih mentioned this pull request Jun 27, 2026
[Roadmap] Blackwell MXFP8 and NVFP4 RL training radixark/miles#615
Open
63 tasks

coderabbitaiBot mentioned this pull request Jul 21, 2026
feat(phyai): humming quant kernels mingti-org/phyai#48
Open

zianglih mentioned this pull request Aug 24, 2026
Fuse NVFP4 fake-QAT quantize-dequantize with CuTe DSL radixark/Megatron-LM#88
Draft
6 tasks

mode	kernel	err	metric	old us	refactor us	old/refactor
1d	nvfp4	-	strict	103.260	103.259	1.000x
1d	4over6	MAE	strict	799.852	800.892	0.999x
1d	4over6	MAE	fast	288.834	289.100	0.999x
1d	4over6	MSE	strict	829.294	829.693	1.000x
1d	4over6	MSE	fast	287.426	289.148	0.994x
2d	nvfp4	-	strict	126.692	126.737	1.000x
2d	4over6	MAE	strict	847.700	847.070	1.001x
2d	4over6	MAE	fast	306.286	306.012	1.001x
2d	4over6	MSE	strict	866.949	867.081	1.000x
2d	4over6	MSE	fast	306.748	305.842	1.003x

Uh oh!

Conversation

zianglih commented Jun 1, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-appsBot commented Jun 1, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

greptile-appsBotJun 1, 2026

Choose a reason for hiding this comment

Uh oh!

zianglihJun 1, 2026

Choose a reason for hiding this comment

Uh oh!

zianglih commented Jun 2, 2026

Uh oh!

zianglih commented Jun 2, 2026

Uh oh!

zianglih commented Jun 2, 2026

Uh oh!

zianglih commented Jun 2, 2026

Uh oh!

ptrendx commented Jun 2, 2026

Uh oh!

zianglih commented Jun 2, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ptrendxJun 2, 2026

Choose a reason for hiding this comment

Uh oh!

zianglihJun 2, 2026

Choose a reason for hiding this comment

Uh oh!

zianglihJun 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zianglihJun 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

timmoon10 commented Jun 8, 2026

Uh oh!

zianglih commented Jun 8, 2026

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zianglih commented Jun 1, 2026•
edited
Loading

greptile-appsBot commented Jun 1, 2026•
edited
Loading

zianglih commented Jun 2, 2026•
edited
Loading

zianglihJun 3, 2026•
edited
Loading

zianglihJun 3, 2026•
edited
Loading