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:

Flow

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

❌ 1 New Failure, 1 Pending

As of commit 4fe3daf with merge base 4e25496 ():

NEW FAILURE - The following job has failed:

• 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 e49f49eebb351bd89a9971e99f6d7e39e1de0ca6e42f1429a9cb0ed56861149b /exec failed with exit code 139

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

[drisspg]

@vkuzo updated to use the mm_config

[gau-nernst]

Just some comments

[vkuzo]

curious on why we need both linear and mm flavors instead of picking one? I thought linear was for torch function and mm variants for torch dispatch?

[drisspg]

We only need the mm and addmm in most circumstances, however if you run under inference mdoe that is pre-dispatch you end up see these linear ops

[vkuzo]

I remember this function was being pretty slow when we wrote the emulation code, curious how performance is now? There are likely things we can do to make it faster if this is in the hot path.

[drisspg]

Yeah @gaunerst has a nice triton kernel for doing the cast + pack that I am going to try in a follow up PR