[Quantization] Add FLUX.2-dev Mistral TE online FP8 support

[dpeng123]

PLEASE FILL IN THE PR DESCRIPTION HERE.

RFC #1854

Purpose

Add online FP8 quantization support for the Mistral text encoder used by
black-forest-labs/FLUX.2-dev.

This branch:

• Threads the resolved text-encoder quantization configuration through the
  Mistral attention and MLP weight-bearing linear layers while leaving embeddings
  and the language-model head unquantized.
• Preserves the full text_encoder module prefix so component quantization rules
  select the intended Mistral layers.
• Resolves text-encoder and transformer quantization independently, enabling
  text-encoder-only FP8 with an unquantized BF16 transformer and VAE.
• Supports FP8 meta-device parameter initialization without prematurely moving
  unresolved parameters to CUDA.
• With CPU offload enabled, moves the unquantized transformer to CPU before
  online text-encoder quantization to avoid the transient load-time GPU OOM.
• Adds FLUX.2-dev FP8 quality coverage and documents FP8 support.

Environment

Item	Value
GPU	NVIDIA H100 SXM, 80 GB HBM3
GPU count / TP	1 / 1
Model	black-forest-labs/FLUX.2-dev
vLLM	0.25.0
vLLM-Omni	0.25.0, editable local branch
PyTorch	2.11.0+cu130
LPIPS	0.1.4, AlexNet v0.1
Image setup	1024x1024, 50 steps, guidance 4.0
CPU offload	Enabled

Test Plan

Unit tests

python -m pytest -q \
  tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py \
  tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py

Result: PASS, 31 passed, 14 warnings in 18.20s.

Online serving smoke test

vllm serve black-forest-labs/FLUX.2-dev \
  --omni \
  --diffusion-quantization-config '{"text_encoder":"fp8","transformer":null,"vae":null}' \
  --enable-cpu-offload \
  --host 0.0.0.0 \
  --port 8091

curl -sS -X POST http://127.0.0.1:8091/v1/images/generations \
  -H 'Content-Type: application/json' \
  -d '{
    "prompt": "a lovely bunny holding a sign that says '\''STOP'\''",
    "size": "1024x1024"
  }'

Field	Value
HTTP status	200
curl total time	58.804734s
server E2E time	58.644981s
generated images	1
image size	1024x1024
image mode	RGB
PNG bytes	3,147,861
peak GPU memory after request	~65.5 GiB

Offline A/B with 10 steps

Config:

• Model: black-forest-labs/FLUX.2-dev
• Prompt: a cup of coffee on a wooden table, morning light
• Size: 1024x1024
• Steps: 10
• Seed: 42
• CPU offload: True
• Quantization config: {"text_encoder": "fp8", "transformer": null, "vae": null}

Case	Result	Generation latency	Peak GPU memory	LPIPS vs BF16
BF16 baseline	PASS	45.105s	~64.6 GiB	0.000000
text_encoder FP8	PASS	51.433s	63.65 GiB	0.007244

The measured LPIPS is 0.007244, which is below the test threshold max_lpips=0.15.

BF16 baseline	text_encoder FP8

Offline A/B with 50 steps

Baseline:

quantization_config = None

Candidate:

quantization_config = {
    "text_encoder": "fp8",
    "transformer": None,
    "vae": None,
}

Each variant loaded the model once and ran 3 measured prompts (seeds 42, 43, 44).
Memory: PyTorch peaks (max_memory_allocated / max_memory_reserved) are recorded per stage in the GPU worker, then counters reset. The text-encoder peak is taken after the prior component is offloaded, so it is not mixed with the BF16 transformer. Stage peaks are also combined for the full-request peak.

Results

Offline benchmark

Prompt	BF16 latency	TE-FP8 latency	FP8 / BF16	BF16 peak	TE-FP8 peak	LPIPS
a cup of coffee on a wooden table, morning light	59.05 s	54.51 s	0.923x	66,070 MB	66,102 MB	0.1230
a lovely bunny holding a sign that says 'STOP'	27.43 s	26.40 s	0.962x	66,070 MB	66,102 MB	0.0418
a fox waking up in bed and yawning	27.30 s	26.43 s	0.968x	66,070 MB	66,102 MB	0.0726
Mean	37.93 s	35.78 s	0.943x	66,070 MB	66,102 MB	0.0792

Component memory profile

The allocated column is the primary component-footprint metric. Reserved memory
is included as allocator context and is used for the whole-request high-water mark.
All three measured prompts produced the same component peaks.

Stage	BF16 allocated	TE-FP8 allocated	Change	BF16 reserved	TE-FP8 reserved
Text encoder	45,466.62 MB	24,309.70 MB	-46.53%	45,498 MB	24,350 MB
Transformer	62,503.96 MB	62,535.96 MB	+0.05%	62,988 MB	63,020 MB
VAE decode stage	64,113.77 MB	64,145.77 MB	+0.05%	66,070 MB	66,102 MB
Whole request	64,113.77 MB	64,145.77 MB	+0.05%	66,070 MB	66,102 MB

Offline A/B (BF16 | TE-FP8)

Prompt	BF16	TE-FP8
prompt_0
prompt_1
prompt_2

Takeaways:

• Text-encoder FP8 reduces the encoder-stage allocated peak by 21,156.92 MB
  (46.53%), from 45,466.62 MB to 24,309.70 MB.
• Average end-to-end latency improves by 5.67%, from 37.93 s to 35.78 s. The
  first measured request includes compile/cold-path cost in both variants.
• Transformer, VAE-stage, and whole-request peaks are effectively unchanged
  (+32 MB) because the diffusion transformer remains unquantized BF16.
• Mean LPIPS is 0.0792. Per-prompt drift ranges from 0.0418 to 0.1230.
• Prompts two and three use the same already-loaded model instance.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[dpeng123]

Hi @Gaohan123 @Isotr0py @RuixiangMa @SamitHuang @ZJY0516 @david6666666 @hsliuustc0106 @princepride @wtomin @xuechendi @yenuo26 , Can someone please take look at the PR? Thanks!

[david6666666]

why we use enable_cpu_offload here?

[dpeng123]

enable_cpu_offload=True is here for VRAM, not for the FP8 quality comparison itself.

FLUX.2-dev (Mistral TE + BF16 DiT + VAE) exceeds ~80 GiB on a single H100 without offload — we hit OOM when running this case with offload disabled. That matches our docs and the TE-FP8 validation recipe:

https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/examples/offline_inference/text_to_image.md — “FLUX.2-dev … >80 (CPU offload required)”
For this test we only quantize the text encoder ({"text_encoder": "fp8", "transformer": None, "vae": None}), so the DiT still loads in BF16 and dominates memory.

Happy to add a one-line comment on the config if that would help future readers.

[david6666666]

Makes sense, thanks!

[david6666666]

vllm serve black-forest-labs/FLUX.2-dev
--omni
--quantization fp8
--host 0.0.0.0
--port 8091

curl -sS -X POST http://127.0.0.1:8091/v1/images/generations
-H 'Content-Type: application/json'
-d '{
"prompt": "a lovely bunny holding a sign that says '''STOP'''",
"size": "1024x1024"
}'

Please post the smoke test results for your online service. Also, I have a question regarding the --quantization fp8 parameter: does this mean the Transformer model itself is also using FP8 quantization?

[david6666666]

Was this exact 10-step case run? The PR only reports 50-step A/B results, so the 0.15 threshold isn't validated yet.

[dpeng123]

Just did offline test with 10 steps, and put in PR description, can you help check?

[david6666666]

Looks good, thanks!

[dpeng123]

vllm serve black-forest-labs/FLUX.2-dev --omni --quantization fp8 --host 0.0.0.0 --port 8091

curl -sS -X POST http://127.0.0.1:8091/v1/images/generations -H 'Content-Type: application/json' -d '{ "prompt": "a lovely bunny holding a sign that says '''STOP'''", "size": "1024x1024" }'

Please post the smoke test results for your online service. Also, I have a question regarding the --quantization fp8 parameter: does this mean the Transformer model itself is also using FP8 quantization?

Thanks for point out this, just did the online serving test again, with --diffusion-quantization-config '{"text_encoder":"fp8","transformer":null,"vae":null}, updated PR description with the new test result

[david6666666]

LGTM now. Thanks for your contribution