NVfp4

[drisspg]

Stacked PRs:

• add-to-benchmarks #2427
• ->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

Script: https://gist.github.com/drisspg/4024ed055a6db911495102614c674c4c -> still emulating till we fix this bug in cublaslt bindings
Double quant really helps w/ tensor that have very small amax values, likely by reducing the amount of underflows will verify:

[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.

✅ No Failures

As of commit d5bded3 with merge base 4e25496 ():
💚 Looks good so far! There are no failures yet. 💚

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

[drisspg]

Weight only fails w/ compile and bisected to:

Likely from the work around to get triton to not error on e2m1

Disabling lowerings fixed the issue.
Starting bisect by getting upper bound.
Upper bound of 38 found for inductor.
Bisecting inductor - lowerings (Range: [0, 38], Midpoint: 19)
Bisecting inductor - lowerings (Range: [20, 38], Midpoint: 29)
Bisecting inductor - lowerings (Range: [30, 38], Midpoint: 34)
Bisecting inductor - lowerings (Range: [35, 38], Midpoint: 36)
Bisecting inductor - lowerings (Range: [35, 36], Midpoint: 35)
Binary search completed for inductor - lowerings. The bisect number is 36. Debug info: convert_element_type_5
Bisection status deleted.
   Bisection result: BisectionResult(backend='inductor', subsystem='lowerings', bisect_number=36, debug_info='convert_element_type_5')

6. Testing inductor config workarounds for WEIGHT_ONLY:
   {'inductor.coordinate_descent_tuning': False}           ERROR
   {'inductor.force_fuse_int_mm_with_mul': False}          ERROR
   {'inductor.post_grad_passes': False}                    ERROR
   {'inductor.pattern_matcher': False}                     ERROR
   {'inductor.epilogue_fusion': False}                     ERROR
   {'inductor.max_autotune': False}                        ERROR
   {'triton.autotune_pointwise': False}                    ✗ 3.1dB
   {'inductor.benchmark_kernel': False}                    ERROR
   {'inductor.aggressive_fusion': False}                   ERROR

7. Testing other compile backends:
   Backend 'eager': SQNR = 20.00 dB ✓
   Backend 'aot_eager': SQNR = 20.00 dB ✓
skipping cudagraphs due to skipping cudagraphs due to cpu device (_tensor_constant0). Found from : 
   File "/home/drisspg/.conda/envs/nightly/lib/python3.12/site-packages/torch/_dynamo/external_utils.py", line 70, in inner
    return fn(*args, **kwargs)

   Backend 'cudagraphs': SQNR = 20.00 dB ✓

[drisspg]

@vkuzo updated to use the mm_config

[gau-nernst]

Just some comments

[gau-nernst]

just curious. what's the point of exposing all of these when only a specific value is accepted.

[drisspg]

This is a good point, my original intent was to not make another subclass, and just merge in w/ mxfp, cc @vkuzo I imagine that we wan't this separated? I started to work on a observer for this since without it this just a worse mxfp4.

[gau-nernst]

I have wondered about this for MX-dtype as well. It makes sense for MX-dtype to have scale swizzling here since we may want the layout to be vendor-neutral. But NVFP4 is specific to NVIDIA, so why not put this under to_nvfp4()?

[drisspg]

This is a good question, for mx I recently confirmed that AMD does not require swizziling, this is a good point about NVFP4, I am actually going update this code path to cache the swizzled layout for the weight