[FP8] enable hunyuan-image-3 diffusion model with fp8 online quant (#1935)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py

@@ -1565,7 +1565,9 @@ class HunYuanAttention(nn.Module):
 
 
 class HunyuanImage3DecoderLayer(nn.Module):
-    def __init__(self, config: HunyuanImage3Config, layer_idx: int, prefix: str = ""):
+    def __init__(
+        self, config: HunyuanImage3Config, quant_config: QuantizationConfig | None, layer_idx: int, prefix: str = ""
+    ):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.layer_idx = layer_idx
@@ -1607,9 +1609,13 @@ class HunyuanImage3DecoderLayer(nn.Module):
             (isinstance(config.num_experts, int) and config.num_experts > 1)
             or (isinstance(config.num_experts, list) and max(config.num_experts) > 1)
         ) and layer_idx >= config.moe_layer_num_skipped:
-            self.mlp = HunYuanSparseMoeBlock(config, layer_id=layer_idx, prefix=f"{prefix}.mlp")
+            self.mlp = HunYuanSparseMoeBlock(
+                config, quant_config=quant_config, layer_id=layer_idx, prefix=f"{prefix}.mlp"
+            )
         else:
-            self.mlp = HunYuanMLP(self.hidden_size, self.intermediate_size, config.hidden_act)
+            self.mlp = HunYuanMLP(
+                self.hidden_size, self.intermediate_size, config.hidden_act, quant_config=quant_config
+            )
 
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1704,14 +1710,15 @@ class HunyuanImage3PreTrainedModel(PreTrainedModel):
 
 
 class HunyuanImage3Model(nn.Module):
-    def __init__(self, config: HunyuanImage3Config, prefix: str = ""):
+    def __init__(self, config: HunyuanImage3Config, quant_config: QuantizationConfig | None, prefix: str = ""):
         super().__init__()
-        quant_config = None
         lora_config = None
         self.num_redundant_experts = 0
         self.config = config
         self.device = get_local_device()
+
         self.quant_config = quant_config
+        logger.debug(f"quant_config: {quant_config}")
         self.padding_idx = config.pad_token_id
         lora_vocab = (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) if lora_config else 0
         self.vocab_size = config.vocab_size + lora_vocab
@@ -1730,6 +1737,7 @@ class HunyuanImage3Model(nn.Module):
             config.num_hidden_layers,
             lambda prefix: HunyuanImage3DecoderLayer(
                 config=config,
+                quant_config=quant_config,
                 layer_idx=int(prefix.split(".")[-1]),
                 prefix=prefix,
             ),

vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py

@@ -20,6 +20,7 @@ from vllm.transformers_utils.config import get_config
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .autoencoder import AutoencoderKLConv3D
@@ -77,7 +78,8 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
                 fall_back_to_pt=True,
             )
         ]
-        self.model = HunyuanImage3Model(self.hf_config)
+        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
+        self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
         self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae)
         self._pipeline = None
         self._tkwrapper = TokenizerWrapper(od_config.model)