NVfp4

[drisspg]

Stacked PRs:

• ->NVfp4 #2408

---

Add NVFP4 Inference flow

Details:

I kept this separate for MX but realistically we should probably merge the two. Basic support for blocksize 16 + e4m3 scales.

Double Quant Update

Ignore previous comments, the double quant is actually really similar to NF4 where you just scale the fp32 scales prior to casting to e4m3 to try and reduce scale quant error.

I have that implemented now in the Nvfp4 code if a tesor_scale is given, just need to figure out how to thread to cublas param scale_in_d or how we want to expose this. We currently don't expose the C matrix to the Python API so we could use alpha as @gau-nernst pointed out to me, however we dont expose alpha either 🙃. However if we wanted to use alpha we would need the value on the host, the sync would likely rule out this option. I might keep this double quant on hold until we have the public api, since I am thinking about adding scale overloads to addmm. However I read the cublas docs many times and it feels as though passing to scale result should work since we don't set the d_mode and its default value should work.

Early Perf

No double quant here

python /home/drisspg/meta/vllm/benchmarks/benchmark_throughput.py \
 --backend vllm \
 --model "data/nvfp4-Qwen3-8B" \
 --dataset-name sharegpt \
 --dataset-path data/ShareGPT_V3_unfiltered_cleaned_split.json \
 --num-prompts 1024 \
 --disable-log-stats \
 --gpu-memory-utilization=0.9 \
 --seed 42

Throughput: 43.23 requests/s, 18347.24 total tokens/s, 8840.47 output tokens/s
Total num prompt tokens:  225190
Total num output tokens:  209407

which is even worse than mxfp4..., will profile later

Micro Bench

LLama 70B mlp no TP:

Model Configuration	Runtime (μs/iteration)	Speedup vs BF16
BF16	1353.09	1.00x
mxfp8	766.76	1.76x
mxfp4	638.00	2.12x
nvfp4	540.41	2.50x

Diffusers

# Bf16 Compile
|           ckpt_id            |   batch_size |  fuse  |  compile  |  compile_vae  |  quantization  |  sparsify  |   model_memory |   inference_memory |   time |
|:----------------------------:|-------------:|:------:|:---------:|:-------------:|:--------------:|:----------:|---------------:|-------------------:|-------:|
| black-forest-labs/FLUX.1-dev |            1 | False  |   True    |     False     |      None      |   False    |         31.438 |             33.827 |  3.286 |

Errors

Annoyingly we are getting an error due to the view as fp4x2 + packing https://fburl.com/cd92w431 because this is trying to be bitcast iside inside triton kernel which is very annoying. Not sure how this didn't show up until vllm / w/ mxfp4
^ similar to this: triton-lang/triton#6054 but make the same changes in _inductor/utils.py as we did for float8em0

Numerics

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2408

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures

As of commit d85d39a with merge base 101c039 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/pt2e/test_numeric_debugger.py::TestNumericDebuggerInfra::test_prepare_for_propagation_comparison
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  test/quantization/pt2e/test_numeric_debugger.py::TestNumericDebuggerInfra::test_prepare_for_propagation_comparison
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t 669bec87c8a304d6313218af7fca8566a24b04d3a50828b7a168acc1effef16e /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t db150226c32ca760871208e019ac234aaa6f0499148bb3947e4d0525b64e570c /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[drisspg]

@vkuzo Curious if you agree we should roll this into the existing mx tensor?

[vkuzo]

can we just write the logic here instead of using to_linear_activation_quantized? I remember same feedback on the mxfp4 inference tensor.

[drisspg]

So you can't just move the logic out here, the entirety of the forward behavior has to be "wrapped" by the subclass. Currently there are two ways to do that, without changing nn.modules.

1. Like above; this is subclass composition
2. The other is to copy the same behavior into the implementations of the ops,
   e.g.

NVFP4's dispatch would need to copy:

ao/torchao/quantization/linear_activation_quantized_tensor.py

Lines 135 to 186 in 7192edf

	@implements([torch.nn.functional.linear, aten.linear.default])
	def _(func, types, args, kwargs):
	input_tensor, weight_tensor, bias = (
	args[0],
	args[1],
	args[2] if len(args) > 2 else None,
	)
	if isinstance(weight_tensor, LinearActivationQuantizedTensor):
	return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias)
	raise NotImplementedError(
	"LinearActivationQuantizedTensor: No specialized dispatch found for linear op"
	)
	@implements([aten.mm.default, aten.addmm.default])
	def _(func, types, args, kwargs):
	if not args[0].is_floating_point():
	raise NotImplementedError(
	"LinearActivationQuantizedTensor: expecting a floating point input"
	)
	if func == aten.addmm.default:
	assert args[1].shape[-1] == args[2].shape[0], (
	f"need mat1 shape: {args[1].shape} final"
	f"dim to match mat2 shape: {args[2].shape} first dim "
	)
	input_tensor, weight_tensor, bias = (
	args[1],
	args[2],
	args[0],
	)
	input_quant_func = weight_tensor.input_quant_func
	original_weight_tensor = weight_tensor.original_weight_tensor
	qtensor = input_quant_func(input_tensor, **weight_tensor.quant_kwargs)
	return func(bias, qtensor, original_weight_tensor)
	else:
	# aten.mm.default
	assert args[0].shape[-1] == args[1].shape[0], (
	f"need mat1 shape: {args[0].shape} final dim"
	f"to match mat2 shape: {args[1].shape} first dim"
	)
	input_tensor, weight_tensor = (
	args[0],
	args[1],
	)
	input_quant_func = weight_tensor.input_quant_func
	original_weight_tensor = weight_tensor.original_weight_tensor
	qtensor = input_quant_func(input_tensor, **weight_tensor.quant_kwargs)
	return func(qtensor, original_weight_tensor)

Not the end of the world. But for some subclasses that serve dual purpose (dyanmic + weight only, + static, + training) it can be alot of switch statements in the ops as opposed to having the base subclass + some sugar

[vkuzo]

lol

[drisspg]

The down up down pattern of casts is probs overkill