From 51cf142bc1dfa372f2c61cfee44535f61afd93de Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Wed, 9 Sep 2026 15:11:31 -0700 Subject: [PATCH] [Bugfix][XPU] Restore N-D output shape for W8A16 FP8 linear (#7301) Signed-off-by: Joshna Medisetty Signed-off-by: Joshna-Medisetty --- .../quantization/test_quantization_fp8.py | 3 +- vllm_omni/platforms/xpu/patch.py | 41 +++++++++++++++++++ vllm_omni/platforms/xpu/platform.py | 5 +++ 3 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 vllm_omni/platforms/xpu/patch.py diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 0da5b6dd3..5793ad95c 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/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" diff --git a/vllm_omni/platforms/xpu/patch.py b/vllm_omni/platforms/xpu/patch.py new file mode 100644 index 000000000..841cf2503 --- /dev/null +++ b/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() diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 3d2cd2d62..df4e7fd55 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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"