[Bugfix][XPU] Restore N-D output shape for W8A16 FP8 linear (#7301)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com>

tests/diffusion/quantization/test_quantization_fp8.py

@@ -300,7 +300,8 @@ def test_single_stage_zimage_fp8_uses_less_memory():
     assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
 
 
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "xpu": "B60"})
+@pytest.mark.advanced_model
 def test_single_stage_qwen_image_fp8():
     """Qwen-Image (random weights) with FP8 generates valid images."""
     model = "riverclouds/qwen_image_random"

vllm_omni/platforms/xpu/patch.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""XPU-specific patches for upstream vLLM behavior."""
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def _patch_w8a16_fp8_nd_activation() -> None:
+    """Restore the N-D output shape for the XPU W8A16 FP8 linear kernel.
+
+    ``XPUW8A16FP8LinearKernel.apply_weights`` hands its activation straight to
+    ``torch.ops._xpu_C.fp8_gemm_w8a16``, skipping the flatten-then-reshape every
+    other ScaledMM kernel does. The op's registered fake always returns 2-D, so a
+    3-D activation makes inductor bake in a 2-D ``assert_size_stride`` that the
+    rank-3 result trips, killing FLUX.1 ``--quantization fp8`` at startup warmup.
+    Only >2-D inputs are affected, so LLM (token-flattened, 2-D) paths are not.
+    """
+    try:
+        from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
+            XPUW8A16FP8LinearKernel,
+        )
+    except ImportError:
+        logger.debug("XPUW8A16FP8LinearKernel not available; skipping FP8 shape patch.")
+        return
+
+    _original_apply_weights = XPUW8A16FP8LinearKernel.apply_weights
+
+    def _patched_apply_weights(self, layer, x, bias=None):
+        output_shape = (*x.shape[:-1], layer.weight.shape[1])
+        x = x.reshape(-1, x.shape[-1])
+        return _original_apply_weights(self, layer, x, bias).view(*output_shape)
+
+    XPUW8A16FP8LinearKernel.apply_weights = _patched_apply_weights
+
+
+def apply_patches() -> None:
+    """Apply all XPU-specific patches."""
+    _patch_w8a16_fp8_nd_activation()

vllm_omni/platforms/xpu/platform.py

@@ -11,6 +11,7 @@ from vllm.platforms.xpu import XPUPlatform
 
 from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+from vllm_omni.platforms.xpu.patch import apply_patches
 
 logger = init_logger(__name__)
 
@@ -31,6 +32,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
     _omni_enum = OmniPlatformEnum.XPU
 
+    def __init__(self):
+        super().__init__()
+        apply_patches()
+
     @classmethod
     def get_omni_ar_worker_cls(cls) -> str:
         return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker"