[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/distributed/hsdp.py

@@ -140,8 +140,13 @@ def apply_hsdp_to_model(
         fs_rank,
     )
 
+    # When the model contains FP8 parameters (online quantization), let FSDP
+    # keep the original storage dtype on all-gather instead of casting to
+    # hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect
+    # FP8 inputs; an implicit FP8 -> bf16 cast would silently break them.
+    has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters())
     mp_policy = MixedPrecisionPolicy(
-        param_dtype=hsdp_config.param_dtype,
+        param_dtype=None if has_fp8_params else hsdp_config.param_dtype,
         reduce_dtype=hsdp_config.reduce_dtype,
         output_dtype=hsdp_config.output_dtype,
         cast_forward_inputs=False,