[Bugfix] Fix HSDP + FP8 online quantization compatibility (#4494)

Signed-off-by: Kwokho Tsui <kh_tsui@163.com>
Co-authored-by: Kwokho Tsui <kh_tsui@163.com>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -602,10 +602,17 @@ class DiffusersPipelineLoader:
             model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
             with set_current_diffusion_config(self.od_config):
                 model = model_cls(od_config=self.od_config)
-            if not is_hsdp and target_device.type != "cpu":
+            # HSDP normally defers GPU placement to apply_hsdp_to_model to keep peak
+            # load-time memory on CPU. Online quantization (e.g. fp8) runs CUDA-only
+            # kernels inside load_weights via the layerwise loader, so when a quant
+            # config is set we initialize on the accelerator like the non-HSDP path;
+            # apply_hsdp_to_model shards GPU-resident params equally well.
+            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
+            if not hsdp_defer_to_cpu and target_device.type != "cpu":
                 model.to(target_device)
         else:
-            device_ctx = target_device if not is_hsdp else contextlib.nullcontext()
+            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
+            device_ctx = contextlib.nullcontext() if hsdp_defer_to_cpu else target_device
             with device_ctx:
                 if load_format == "default":
                     model = initialize_model(self.od_config)
@@ -646,6 +653,22 @@ class DiffusersPipelineLoader:
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
+        # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
+        # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
+        # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
+        # Rewrite affected layers in-place to row-major contiguous storage and
+        # shift the .t() to GEMM-call time. Layers using other quant methods or
+        # already-contiguous weights are left untouched.
+        if self.quant_config is not None:
+            from vllm_omni.diffusion.quantization.hsdp_fp8 import (
+                prepare_fp8_layers_for_fsdp,
+            )
+
+            for _attr in ("transformer", "transformer_2"):
+                _trans = getattr(model, _attr, None)
+                if _trans is not None:
+                    prepare_fp8_layers_for_fsdp(_trans)
+
         # Collect all transformers to shard (some models have transformer_2 for MoE)
         transformers_to_shard = []
         transformer = getattr(model, "transformer", None)