Seeing a lot of FLUX.2 “insufficient memory” and “it’s slow” posts here, so I
want to share something that fixes both — and clear up a quantization gotcha
that trips everyone up.
I added the first image-generation model to @eugr’s spark-vllm-docker
(everything in it so far is an LLM via vLLM). It runs FLUX.2-dev as a
headless OpenAI Images-API server, with the transformer quantized to
NVFP4 on the fly using torchao.
Measured on my DGX Spark GB10 (sm_121a), FLUX.2-dev, 28 steps @ 1024²:
| BF16 | NVFP4 | speedup | |
|---|---|---|---|
| text-to-image | ~2.3 min | ~45 s | ~3× |
| single-ref edit | ~4 min 20 s | ~1 min 51 s | ~2.3× |
| VRAM (steady) | ~112 GB | ~66 GB | ~40% less |
The gotcha — “NVFP4” doesn’t automatically mean faster:
- Most quantized FLUX files floating around (fp8-mixed, gguf, naively-loaded
“nvfp4” checkpoints) are weight-only: weights are stored small but the
matmul upcasts back to BF16 to compute. You get the memory saving but
little-to-no speedup — sometimes slower. - This uses torchao W4A4 (activations quantized too), so the matmul actually
runs in FP4 on the Blackwell tensor cores via Triton kernels. That’s where
the ~3× comes from. Same model, same steps — the only variable is compute.
A couple of things I learned that match what others hit here:
- Don’t use the
modelopt_fp4/ prequantized-NVFP4-checkpoint path on
sm_121a — it hits a diffusers unpack/shape bug (others have reported the same
“don’t pass modelopt_fp4 on SM121A”). On-the-fly torchao is the working route. mslkis the missing dependency — without it torchao errors “mslk is required
for NVFP4 triton quantization.”- Memory: the ~112 GB BF16 peak only happens on the first quantize. It
saves the quantized weights, and later boots load them directly at ~66 GB
with no BF16 spike — so after a one-time quant it co-hosts fine next to an
LLM. - CUDA 13 + Blackwell (sm_120a/121a) required for the FP4 kernels.
It’s complementary to ComfyUI — this is the API-first/headless route (POST to
/v1/images/generations), not a GUI.
Based on the PyTorch team’s Blackwell diffusion write-up:
@eugr — would love your take on the servers/<model>/ layout I used for
non-vLLM servers; happy to restructure however fits the repo.