From ed29dfae6e76235e01fcc0eca21cf0dcbfc40af1 Mon Sep 17 00:00:00 2001
From: zofia <110436990+zufangzhu@users.noreply.github.com>
Date: Mon, 7 Sep 2026 13:44:31 +0800
Subject: [PATCH] [XPU] Use fused_input_norm kernel in FusedInputNorm (#52945)

Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
---
 vllm/_xpu_ops.py                     | 37 ++++++++++++++++++++++++++++
 vllm/model_executor/models/vision.py | 14 +++++++++++
 2 files changed, 51 insertions(+)

diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py
index df66e64033..c1f2bb6969 100644
--- a/vllm/_xpu_ops.py
+++ b/vllm/_xpu_ops.py
@@ -578,6 +578,37 @@ def _xpu_mxfp4_quantize_fake(
     return x_q, x_s
 
 
+def _xpu_fused_input_norm_impl(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    visual_dtype: torch.dtype,
+) -> torch.Tensor:
+    patches, size = x.shape
+    out = torch.empty(
+        (patches, size),
+        dtype=visual_dtype,
+        device=x.device,
+    )
+    torch.ops._xpu_C.fused_input_norm(out, x.contiguous(), weight, bias)
+    return out
+
+
+def _xpu_fused_input_norm_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    visual_dtype: torch.dtype,
+) -> torch.Tensor:
+    patches, size = x.shape
+    out = torch.empty(
+        (patches, size),
+        dtype=visual_dtype,
+        device=x.device,
+    )
+    return out
+
+
 @triton.jit
 def _softplus(x):
     return tl.where(x <= 20.0, tl.math.log(tl.math.exp(x) + 1.0), x)
@@ -1281,6 +1312,12 @@ class xpu_ops:
                 ],
             )
 
+            direct_register_custom_op(
+                op_name="xpu_fused_input_norm",
+                op_func=_xpu_fused_input_norm_impl,
+                fake_impl=_xpu_fused_input_norm_fake,
+            )
+
             _OPS_REGISTERED = True
 
 
diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py
index 0ccf410ed5..36008d9996 100644
--- a/vllm/model_executor/models/vision.py
+++ b/vllm/model_executor/models/vision.py
@@ -718,6 +718,20 @@ class FusedInputNorm(nn.Module):
         patches, size = grid_thw.shape
         patch_size = size // self.channel
 
+        # On XPU, fuse the whole rescale + normalize into a single custom
+        # kernel. The eager path below materializes an fp32 intermediate and
+        # then casts back, which adds device-side compute that cancels the
+        # bandwidth saving of transferring uint8 pixel_values. The fused
+        # kernel reads uint8 directly and writes ``visual_dtype`` in one pass.
+        if (
+            current_platform.is_xpu()
+            and grid_thw.dtype == torch.uint8
+            and self.weight.dtype == torch.float32
+        ):
+            return torch.ops.vllm.xpu_fused_input_norm(
+                grid_thw, self.weight, self.bias, visual_dtype
+            )
+
         # Apply the per-channel affine transform directly instead of via
         # F.batch_norm. batch_norm dispatches to cuDNN, whose batch-norm
         # kernels cap the batch dimension near the CUDA grid limit (~65535);