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

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