[FP8 Quantization] Add FP8 quantization support for Flux transformer

[zzhuoxin1508]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

This PR introduces online FP8 quantization for DiT models in vLLM-Omni, starting with Flux (text-to-image)., Online FP8 converts BF16/FP16 weights to FP8 at model load time with dynamic activation scaling — no calibration dataset or pre-quantized checkpoint needed.

files change

• flux_transformer.py
  • Dual-stream blocks (FluxTransformerBlock ×19)
    • Pass quant_config to FluxAttention: applied to to_qkv, add_kv_proj (QKVParallelLinear), to_out, to_add_out (RowParallelLinear)
    • Pass quant_config to FeedForward (both image and text streams): applied to ColumnParallelLinear (up-projection) and RowParallelLinear (down-projection)
  • Single-stream blocks (FluxSingleTransformerBlock ×38)
    • Pass quant_config to FluxAttention: applied to to_qkv (QKVParallelLinear)
    • Replace nn.Linear with ReplicatedLinear for proj_mlp and proj_out, since nn.Linear does not support quant_config
• vllm_omni/diffusion/models/flux/pipeline_flux.py — Extract and pass quant_config to transformer via get_vllm_quant_config_for_layers

Not quantized: context_embedder, x_embedder, final proj_out — these are small input/output embedding layers where quantization error has direct impact on output and parameter savings are negligible.

Benchmark (H100 x1)

Model	Task	Memory BF16	Memory FP8	Reduction	Time BF16	Time FP8	Speedup
FLUX.1-dev (12B)	Text-to-Image 1024×1024, 20 steps	31.43 GiB	23.47 GiB	25%	6.30s	5.15s	18%

---

Image Quality Comparison

Prompt	BF16	FP8
"a cup of coffee on the table"
"a cat sitting on a windowsill at sunset."
"Cinematic macro shot, cozy sunlit reading nook, soft sheepskin rug, antique wooden bookshelf with blurred old books background, steaming mug of tea, realistic knit blanket texture, volumetric sun rays, potted succulent, intricate gold dust particles in air, warm colors, masterpiece, sharp focus, comforting atmosphere."

command:

python examples/offline_inference/text_to_image/text_to_image.py \
  --model black-forest-labs/FLUX.1-dev \
  --prompt "a cup of coffee on the table." \
  --seed 42 \
  --num-inference-steps 20 \
  --height 1024 \
  --width 1024 \
  --guidance-scale 3.5 \
  --cfg-scale 1.0 \
  --quantization fp8 \
  --output outputs/coffee.png

All parameters kept constant except for the prompt and quantization

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[zzhuoxin1508]

@lishunyang12

[hsliuustc0106]

PR #1640 Review: 为Flux transformer添加FP8量化支持

📊 总体评分：7.5/10

这是一个结构良好的PR，正确实现了FP8量化支持，实现了显著的性能提升。

---

✅ 优点

1. 代码质量和架构设计

• ✅ 参数传递一致：quant_config正确通过类层次结构传播（FluxTransformer2DModel → FluxTransformerBlock → FluxAttention/FeedForward）
• ✅ 遵循现有模式：实现风格与其他模型（Z-Image, Qwen-Image）保持一致
• ✅ 正确的导入处理：使用TYPE_CHECKING避免循环导入

2. FP8量化实现

• ✅ 正确的层类型替换：在FluxSingleTransformerBlock中将nn.Linear替换为ReplicatedLinear
• ✅ 合理的层排除策略：输入/输出嵌入层（context_embedder, x_embedder, proj_out）未量化
• ✅ 完整的线性层覆盖：
  • QKVParallelLinear: to_qkv, add_kv_proj
  • RowParallelLinear: to_out, to_add_out, FFN下投影
  • ColumnParallelLinear: FFN上投影
  • ReplicatedLinear: proj_mlp, proj_out (single-stream blocks)

3. 性能优化

• ✅ 25%内存节省：31.43 GiB → 23.47 GiB
• ✅ 18%速度提升：6.30s → 5.15s
• ✅ 选择性量化策略：小嵌入层保持BF16避免输出质量下降

---

⚠️ 需要改进

🔴 必须修复

1. 测试覆盖不足

• ❌ 缺少Flux-specific FP8测试
• ❌ 缺少端到端集成测试
• ❌ 缺少权重加载测试

建议添加测试：

# tests/diffusion/models/flux/test_flux_fp8.py

def test_flux_transformer_fp8_initialization():
    """Test that FluxTransformer2DModel initializes correctly with FP8 config."""
    from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
    from vllm_omni.diffusion.quantization import get_diffusion_quant_config

    diff_config = get_diffusion_quant_config("fp8")
    quant_config = get_vllm_quant_config_for_layers(diff_config)

    model = FluxTransformer2DModel(
        od_config=create_test_od_config(),
        quant_config=quant_config
    )

    # Verify quantization is applied to linear layers
    for block in model.transformer_blocks:
        assert hasattr(block.attn.to_qkv, 'quant_method')

def test_flux_fp8_forward_pass():
    """Test forward pass works with FP8 quantization."""
    # Create model with FP8
    # Run forward pass
    # Verify output shape and dtype

2. 文档需要更新

• ❌ docs/user_guide/diffusion/quantization/fp8.md中的支持模型表需要添加Flux

建议更新表格：

Model	HF Models	Recommendation	ignored_layers
Flux.1-dev	black-forest-labs/FLUX.1-dev	All layers	None

---

🟡 建议修复

1. 添加代码注释

flux_transformer.py:478-481
建议为packed_modules_mapping添加注释说明用途：

# Mapping for LoRA weight loading: packed module name -> logical sub-projections
# Used by LoRA manager to expand checkpoints trained on individual Q/K/V projections
packed_modules_mapping = {
    "to_qkv": ["to_q", "to_k", "to_v"],
    "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],
}

2. 更新docstring

flux_transformer.py:443-472
建议在FluxTransformer2DModel的docstring中添加quant_config参数文档：

Args:
    ...
    quant_config (`QuantizationConfig`, *optional*):
        Quantization config for linear layers. Pass FP8 config to enable
        online quantization from BF16 weights.

3. 验证张量并行兼容性

需要确认FP8量化在TP>1时的正确性。

---

🟢 可选优化

1. 添加更详细的性能benchmark数据

   • 吞吐量（images/sec）
   • 不同分辨率下的性能
   • 不同batch size下的内存占用

2. 验证权重加载逻辑
   确认load_weights方法在FP8模式下正常工作

---

🐛 潜在风险

中等风险

1. packed_modules_mapping未在测试中验证

   • LoRA权重与量化后的基础模型可能不兼容
   • 建议添加LoRA + FP8测试或文档说明限制

2. TP兼容性未明确

   • 需要确认FP8量化与TP>1时的正确性

低风险

1. 硬编码的LayerNorm
   • nn.LayerNorm保持不变是正确的，但建议添加注释说明原因

---

📝 合并建议

建议在合并前完成：

1. ✅ 添加Flux FP8集成测试
2. ✅ 更新FP8文档添加Flux支持
3. ⚠️ （可选但推荐）添加关键代码注释

总体评价：
这是一个高质量的PR，架构设计合理，性能提升显著。主要改进空间在于测试覆盖和文档更新。完成上述必须修复项后即可合并。

---

🦐 Reviewed by AI Assistant

[david6666666]

add supported models in vllm-omni/docs/user_guide/diffusion/quantization/fp8.md

[zzhuoxin1508]

add supported models in vllm-omni/docs/user_guide/diffusion/quantization/fp8.md

Support list added.

[zzhuoxin1508]

Regrading the ai assitant suggestions:
All points follow the same pattern as the existing Z-Image FP8 implementation, which has none of these tests or comments either. I'd prefer to keep this PR consistent with the existing codebase style and defer the tests to a future PR. As for TP compatibility, the underlying layers (QKVParallelLinear, ColumnParallelLinear, RowParallelLinear) are vLLM's built-in TP-aware primitives, so correctness under TP>1 is guaranteed by the framework.

[hsliuustc0106]

lgtm

[lishunyang12]

Looks good. The quant_config plumbing is straightforward and ReplicatedLinear return semantics check out.

[WaterKnight1998]

Are you planning on support Flux 2 Klein??

[lishunyang12]

Are you planning on support Flux 2 Klein??

Quantization team will distributed the work to land the support once the unified framework is integrated. Also, welcome community contributions

[WaterKnight1998]

Are you planning on support Flux 2 Klein??

Quantization team will distributed the work to land the support once the unified framework is integrated. Also, welcome community contributions

I am happy to contribute, just let me know when unified framework is in place.

[lishunyang12]

Are you planning on support Flux 2 Klein??

Quantization team will distributed the work to land the support once the unified framework is integrated. Also, welcome community contributions

I am happy to contribute, just let me know when unified framework is in place.

It is good hear that. I will let you know then. I will updaste this Q1 quantization roadmap later #1057.