[Quantization] Add FLUX.2-dev Mistral TE online FP8 support (#5136)

Signed-off-by: diyapeng <50388704+diyapeng@users.noreply.github.com>
Co-authored-by: diyapeng <50388704+diyapeng@users.noreply.github.com>

docs/user_guide/quantization/fp8.md

@@ -88,6 +88,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
 | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
 | FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
+| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
 | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
 | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |

tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py

@@ -0,0 +1,50 @@
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
+    Flux2Pipeline,
+    _resolve_component_quant_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_resolve_component_quant_config_routes_components():
+    text_encoder_config = object()
+    component_config = Mock()
+    component_config.resolve.side_effect = {
+        "text_encoder": text_encoder_config,
+        "transformer": None,
+    }.get
+
+    assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
+    assert _resolve_component_quant_config(component_config, "transformer") is None
+
+
+def test_resolve_component_quant_config_preserves_global_config():
+    global_config = object()
+    assert _resolve_component_quant_config(global_config, "transformer") is global_config
+    assert _resolve_component_quant_config(None, "text_encoder") is None
+
+
+def test_offload_transformer_before_text_encoder_weights():
+    pipeline = object.__new__(Flux2Pipeline)
+    pipeline.transformer = Mock()
+    weights = [
+        ("transformer.layer.weight", torch.empty(1)),
+        ("text_encoder.layer.weight", torch.empty(1)),
+        ("text_encoder.layer.bias", torch.empty(1)),
+    ]
+
+    iterator = iter(pipeline._offload_transformer_before_text_encoder(weights))
+    assert next(iterator)[0] == "transformer.layer.weight"
+    pipeline.transformer.to.assert_not_called()
+
+    with patch("vllm_omni.diffusion.models.flux2.pipeline_flux2.current_omni_platform.empty_cache") as empty_cache:
+        assert next(iterator)[0] == "text_encoder.layer.weight"
+        assert next(iterator)[0] == "text_encoder.layer.bias"
+
+    pipeline.transformer.to.assert_called_once_with("cpu")
+    empty_cache.assert_called_once_with()

tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py

@@ -5,6 +5,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 
 from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import (
     MistralEncoderModel,
@@ -491,3 +492,57 @@ class TestComputeLogits:
         assert logits.dim() == 3
         assert logits.shape[0] == 1
         assert logits.shape[1] == 4
+
+
+class TestQuantConfig:
+    """Verify quant_config is threaded into weight linears only, with correct prefixes."""
+
+    def test_quant_config_on_weight_linears(self, mocker):
+        quant_config = mocker.MagicMock(name="QuantizationConfig")
+        # Keep construction on the CPU-safe unquantized path while still
+        # verifying LinearBase consults the supplied quant_config.
+        quant_config.get_quant_method.return_value = UnquantizedLinearMethod()
+
+        config = _make_config(num_hidden_layers=1)
+        model = MistralEncoderModel(config, prefix="text_encoder", quant_config=quant_config)
+
+        layer = model.language_model.model.layers[0]
+        for linear in (
+            layer.self_attn.qkv_proj,
+            layer.self_attn.o_proj,
+            layer.mlp.gate_up_proj,
+            layer.mlp.down_proj,
+        ):
+            assert linear.quant_config is quant_config
+
+        # get_quant_method should be consulted for the four weight linears.
+        assert quant_config.get_quant_method.call_count == 4
+        seen_prefixes = set()
+        for call in quant_config.get_quant_method.call_args_list:
+            if "prefix" in call.kwargs:
+                seen_prefixes.add(call.kwargs["prefix"])
+            elif len(call.args) > 1:
+                seen_prefixes.add(call.args[1])
+        assert any(p.endswith("self_attn.qkv_proj") for p in seen_prefixes)
+        assert any(p.endswith("self_attn.o_proj") for p in seen_prefixes)
+        assert any(p.endswith("mlp.gate_up_proj") for p in seen_prefixes)
+        assert any(p.endswith("mlp.down_proj") for p in seen_prefixes)
+
+        # Embeddings / lm_head stay unquantized (quant_config never passed).
+        embed = model.language_model.model.embed_tokens
+        lm_head = model.language_model.lm_head
+        assert getattr(embed, "quant_config", None) is None
+        assert getattr(lm_head, "quant_config", None) is None
+
+    def test_layer_prefixes_include_text_encoder(self):
+        config = _make_config(num_hidden_layers=2)
+        model = MistralEncoderModel(config, prefix="text_encoder")
+
+        layer0 = model.language_model.model.layers[0]
+        assert layer0.self_attn.qkv_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
+        assert layer0.self_attn.o_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
+        assert layer0.mlp.gate_up_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
+        assert layer0.mlp.down_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
+
+        layer1 = model.language_model.model.layers[1]
+        assert "layers.1" in layer1.self_attn.qkv_proj.prefix

tests/diffusion/quantization/test_quantization_quality.py

@@ -82,6 +82,7 @@ class QualityTestConfig:
     gpu: str = "H100"  # minimum GPU requirement
     negative_prompt: str = ""
     guidance_scale: float | None = None
+    enable_cpu_offload: bool = False
 
     def baseline_ref(self) -> str:
         return self.baseline_model or self.model or ""
@@ -181,6 +182,18 @@ QUALITY_CONFIGS = [
         max_lpips=0.20,
         num_inference_steps=10,
     ),
+    QualityTestConfig(
+        id="fp8_flux2_dev_text_encoder",
+        model="black-forest-labs/FLUX.2-dev",
+        quantization={"text_encoder": "fp8", "transformer": None, "vae": None},
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.15,
+        num_inference_steps=10,
+        enable_cpu_offload=True,
+        height=1024,
+        width=1024,
+    ),
     QualityTestConfig(
         id="fp8_qwen_image",
         model="Qwen/Qwen-Image",
@@ -459,7 +472,10 @@ def test_quantization_quality(config: QualityTestConfig):
     generate_fn = _generate_video if config.task == "t2v" else _generate_image
 
     # --- BF16 baseline ---
-    omni_bl = Omni(model=config.baseline_ref())
+    bl_kwargs: dict = {"model": config.baseline_ref()}
+    if config.enable_cpu_offload:
+        bl_kwargs["enable_cpu_offload"] = True
+    omni_bl = Omni(**bl_kwargs)
     baseline_out, bl_mem = generate_fn(omni_bl, config)
     omni_bl.shutdown()
     del omni_bl
@@ -468,10 +484,13 @@ def test_quantization_quality(config: QualityTestConfig):
 
     # --- Quantized ---
     quantization = config.quantization_ref()
+    qt_kwargs: dict = {"model": config.quantized_ref()}
+    if config.enable_cpu_offload:
+        qt_kwargs["enable_cpu_offload"] = True
     if quantization is None:
-        omni_qt = Omni(model=config.quantized_ref())
+        omni_qt = Omni(**qt_kwargs)
     else:
-        omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization)
+        omni_qt = Omni(**qt_kwargs, quantization_config=quantization)
     quant_out, qt_mem = generate_fn(omni_qt, config)
     omni_qt.shutdown()
     del omni_qt

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -37,10 +37,17 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
 
+def _resolve_component_quant_config(quant_config, component: str):
+    if hasattr(quant_config, "resolve"):
+        return quant_config.resolve(component)
+    return quant_config
+
+
 class Flux2ImageProcessor(VaeImageProcessor):
     """Image processor to preprocess the reference image for Flux2."""
 
@@ -393,10 +400,14 @@ class Flux2Pipeline(
         text_encoder_config = AutoConfig.from_pretrained(
             model, subfolder="text_encoder", local_files_only=local_files_only
         )
+        te_quant_config = _resolve_component_quant_config(od_config.quantization_config, "text_encoder")
         self.text_encoder = MistralEncoderModel(
             text_encoder_config,
             prefix="text_encoder",
-        ).to(self._execution_device)
+            quant_config=te_quant_config,
+        )
+        if not any(param.is_meta for param in self.text_encoder.parameters()):
+            self.text_encoder.to(self._execution_device)
         self.tokenizer = PixtralProcessor.from_pretrained(
             model, subfolder="tokenizer", local_files_only=local_files_only
         )
@@ -409,8 +420,9 @@ class Flux2Pipeline(
             self._execution_device
         )
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
+        transformer_quant_config = _resolve_component_quant_config(od_config.quantization_config, "transformer")
         self.transformer = Flux2Transformer2DModel(
-            quant_config=od_config.quantization_config, od_config=od_config, **transformer_kwargs
+            quant_config=transformer_quant_config, od_config=od_config, **transformer_kwargs
         )
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
@@ -1148,5 +1160,26 @@ class Flux2Pipeline(
         return DiffusionOutput(output=image)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        text_encoder_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "text_encoder")
+        transformer_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "transformer")
+        if (
+            self.od_config.enable_cpu_offload
+            and text_encoder_quant_config is not None
+            and transformer_quant_config is None
+        ):
+            weights = self._offload_transformer_before_text_encoder(weights)
+
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
+
+    def _offload_transformer_before_text_encoder(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        transformer_offloaded = False
+        for name, weight in weights:
+            if not transformer_offloaded and name.startswith("text_encoder."):
+                self.transformer.to("cpu")
+                current_omni_platform.empty_cache()
+                transformer_offloaded = True
+                logger.info("Offloaded the unquantized transformer before text encoder quantization")
+            yield name, weight

vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
@@ -37,6 +37,9 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 logger = logging.getLogger(__name__)
 
 
@@ -118,6 +121,7 @@ class MistralEncoderAttention(nn.Module):
         num_kv_heads: int,
         head_dim: int,
         prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         self.hidden_size = hidden_size
@@ -140,6 +144,7 @@ class MistralEncoderAttention(nn.Module):
             total_num_heads=num_heads,
             total_num_kv_heads=num_kv_heads,
             bias=False,
+            quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
         )
 
@@ -147,6 +152,7 @@ class MistralEncoderAttention(nn.Module):
             input_size=num_heads * head_dim,
             output_size=hidden_size,
             bias=False,
+            quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
 
@@ -194,18 +200,21 @@ class MistralEncoderMLP(nn.Module):
         hidden_size: int,
         intermediate_size: int,
         prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         self.gate_up_proj = MergedColumnParallelLinear(
             input_size=hidden_size,
             output_sizes=[intermediate_size, intermediate_size],
             bias=False,
+            quant_config=quant_config,
             prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = RowParallelLinear(
             input_size=intermediate_size,
             output_size=hidden_size,
             bias=False,
+            quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
         )
         self.act_fn = SiluAndMul()
@@ -227,6 +236,7 @@ class MistralEncoderLayer(nn.Module):
         intermediate_size: int,
         rms_norm_eps: float,
         prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         self.self_attn = MistralEncoderAttention(
@@ -235,11 +245,13 @@ class MistralEncoderLayer(nn.Module):
             num_kv_heads=num_kv_heads,
             head_dim=head_dim,
             prefix=f"{prefix}.self_attn",
+            quant_config=quant_config,
         )
         self.mlp = MistralEncoderMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             prefix=f"{prefix}.mlp",
+            quant_config=quant_config,
         )
         self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
@@ -295,7 +307,12 @@ class MistralEncoderModel(nn.Module):
     parallel layers for TP but simple SDPA for attention (no PagedAttention).
     """
 
-    def __init__(self, config: PretrainedConfig, prefix: str = ""):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ):
         super().__init__()
         self._processor = None
         self._system_message_t2i: str | None = None
@@ -339,6 +356,7 @@ class MistralEncoderModel(nn.Module):
         m = self.language_model.model
 
         m.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size)
+        layer_prefix_root = f"{prefix}.language_model.model.layers" if prefix else "language_model.model.layers"
 
         m.layers = nn.ModuleList(
             [
@@ -349,7 +367,8 @@ class MistralEncoderModel(nn.Module):
                     head_dim=self.head_dim,
                     intermediate_size=self.intermediate_size,
                     rms_norm_eps=self.rms_norm_eps,
-                    prefix=f"language_model.model.layers.{i}",
+                    prefix=f"{layer_prefix_root}.{i}",
+                    quant_config=quant_config,
                 )
                 for i in range(self.num_layers)
             ]