（Phase 1）Add ModelOpt FP8 auto-detect support for diffusion checkpoints #2709 (#2913)

Signed-off-by: roG0d <rodgarcas98@gmail.com>
Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com>
Co-authored-by: roG0d <rodgarcas98@gmail.com>

vllm_omni/diffusion/data.py

@@ -208,10 +208,11 @@ class TransformerConfig:
         quant_method: str | None = None
         quant_config: QuantizationConfig | None = None
         disk_qc = params.get("quantization_config")
-        if isinstance(disk_qc, dict) and "quant_method" in disk_qc:
-            quant_method = disk_qc["quant_method"]
-            kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
-            quant_config = build_quant_config(quant_method, **kwargs)
+        if isinstance(disk_qc, dict):
+            raw_quant_method = disk_qc.get("quant_method", disk_qc.get("method"))
+            quant_config = build_quant_config(disk_qc)
+            if quant_config is not None:
+                quant_method = raw_quant_method if raw_quant_method is not None else quant_config.get_name()
 
         return cls(params=params, quant_method=quant_method, quant_config=quant_config)
 
@@ -513,6 +514,10 @@ class OmniDiffusionConfig:
     # str is resolved to {"method": <str>} internally.
     # Per-component: {"transformer": {"method": "fp8"}, "vae": None}
     quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
+    # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
+    # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
+    # has already resolved to vLLM's ModelOpt FP8 linear method.
+    force_cutlass_fp8: bool = False
 
     # Diffusion pipeline Profiling config
     enable_diffusion_pipeline_profiler: bool = False
@@ -630,14 +635,9 @@ class OmniDiffusionConfig:
 
         # Auto-detect quantization from TransformerConfig if not explicitly set.
         # This covers the case where tf_model_config is passed at construction
-        # time.  For late (post-construction) assignment, callers should use
+        # time. For late (post-construction) assignment, callers should use
         # set_tf_model_config() which propagates quant_config automatically.
-        if self.quantization_config is None and self.tf_model_config.quant_config is not None:
-            self.quantization_config = self.tf_model_config.quant_config
-            logger.info(
-                "Auto-detected quantization '%s' from model config",
-                self.tf_model_config.quant_method,
-            )
+        self._propagate_quantization_from_tf_config(self.tf_model_config)
 
         # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
         if self.quantization_config is not None:
@@ -664,6 +664,32 @@ class OmniDiffusionConfig:
                 "valid together with diffusion_load_format=diffusers"
             )
 
+    def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None:
+        if tf_config.quant_config is None:
+            return
+
+        is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False))
+        should_use_checkpoint_config = self.quantization_config is None or (
+            is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config)
+        )
+        if should_use_checkpoint_config:
+            self.quantization_config = tf_config.quant_config
+            logger.info(
+                "Auto-detected quantization '%s' from model config",
+                tf_config.quant_method,
+            )
+
+    @staticmethod
+    def _is_generic_fp8_quant_config(quant_config: object) -> bool:
+        if isinstance(quant_config, str):
+            return quant_config.lower() == "fp8"
+        if isinstance(quant_config, Mapping):
+            method = quant_config.get("method", quant_config.get("quant_method"))
+            return isinstance(method, str) and method.lower() == "fp8"
+        if hasattr(quant_config, "get_name"):
+            return quant_config.get_name() == "fp8"
+        return False
+
     def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
         """Assign `tf_model_config` and propagate quantization if detected.
 
@@ -679,12 +705,7 @@ class OmniDiffusionConfig:
                 `TransformerConfig.from_dict`.
         """
         self.tf_model_config = tf_config
-        if self.quantization_config is None and tf_config.quant_config is not None:
-            self.quantization_config = tf_config.quant_config
-            logger.info(
-                "Auto-detected quantization '%s' from model config",
-                tf_config.quant_method,
-            )
+        self._propagate_quantization_from_tf_config(tf_config)
 
     def update_multimodal_support(self) -> None:
         # Resolve serving-visible multimodal behavior from shared metadata
@@ -716,50 +737,35 @@ class OmniDiffusionConfig:
                 # Skip transformer config loading for diffusers adapter
                 # (non-DiT models don't have a separate transformer folder/config)
                 if self.diffusion_load_format == "diffusers":
-                    self.tf_model_config = TransformerConfig()
-                    try:
-                        diffusers_pipeline_cls_name = config_dict["_class_name"]
-                        self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)
-                    except (KeyError, AttributeError) as exc:
-                        logger.warning(
-                            f"Could not find valid _class_name for diffusers pipeline in model_index.json: {exc}. "
-                            "Without knowing the underlying diffusers pipeline class, the dummy run will input only "
-                            "text prompt, which may cause errors for pipelines that require additional inputs."
-                        )
+                    self.set_tf_model_config(TransformerConfig())
                 else:
                     tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
-                    self.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+                    self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
             else:
                 raise FileNotFoundError("model_index.json not found")
         except (AttributeError, OSError, ValueError, FileNotFoundError):
             # Skip transformer config loading for diffusers adapter
             # (non-DiT models don't have a separate transformer folder/config)
             if self.diffusion_load_format == "diffusers":
-                self.tf_model_config = TransformerConfig()
-                logger.warning(
-                    "Could not find valid model_index.json per diffusers format. "
-                    "This model is suspectedly unsupported by the diffusers backend. "
-                    "Also, without knowing the underlying diffusers pipeline class from model_index.json, "
-                    "the dummy run will input only text prompt, which may cause errors for pipelines "
-                    "that require additional inputs."
-                )
+                self.set_tf_model_config(TransformerConfig())
+                self.update_multimodal_support()
             else:
                 cfg = get_hf_file_to_dict("config.json", self.model)
                 if cfg is None:
                     raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
 
-                self.tf_model_config = TransformerConfig.from_dict(cfg)
+                self.set_tf_model_config(TransformerConfig.from_dict(cfg))
                 model_type = cfg.get("model_type")
                 architectures = cfg.get("architectures") or []
 
                 if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
                     self.model_class_name = "BagelPipeline"
-                    self.tf_model_config = TransformerConfig()
+                    self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()
                 elif model_type == "nextstep":
                     if self.model_class_name is None:
                         self.model_class_name = "NextStep11Pipeline"
-                    self.tf_model_config = TransformerConfig()
+                    self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()
                 elif model_type == "s2v":
                     if self.model_class_name is None:
@@ -794,23 +800,10 @@ class OmniDiffusionConfig:
             cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
             kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
 
-        # Falsy-value check for not-None fields (convert potential None values in YAML config to empty containers)
-        if "diffusers_load_kwargs" in kwargs and kwargs["diffusers_load_kwargs"] is None:
-            kwargs["diffusers_load_kwargs"] = {}
-        if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None:
-            kwargs["diffusers_call_kwargs"] = {}
-
-        # Forward top-level parallel knobs (e.g. --tensor-parallel-size from CLI)
-        # into parallel_config so the diffusion engine sees them.
-        par = kwargs.get("parallel_config", {})
-        if isinstance(par, Mapping):
-            par = dict(par)
-            if par.get("tensor_parallel_size") is None:
-                par.pop("tensor_parallel_size", None)
-            tensor_parallel_size = kwargs.get("tensor_parallel_size")
-            if tensor_parallel_size is not None and "tensor_parallel_size" not in par:
-                par["tensor_parallel_size"] = tensor_parallel_size
-            kwargs["parallel_config"] = par
+        # Convert optional YAML null values to empty containers.
+        for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
+            if key in kwargs and kwargs[key] is None:
+                kwargs[key] = {}
 
         # Filter kwargs to only include valid fields
         valid_fields = {f.name for f in fields(cls)}