[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>
zofia committed
Sep 7, 2026 at 13:44 UTC
ed29dfae6e76235e01fcc0eca21cf0dcbfc40af1
2 files changed
+51
vllm/_xpu_ops.py
+37
@@ -578,6 +578,37 @@ def _xpu_mxfp4_quantize_fake(
578
return x_q, x_s
579
580
581
+def _xpu_fused_input_norm_impl(
582
+ x: torch.Tensor,
583
+ weight: torch.Tensor | None,
584
+ bias: torch.Tensor | None,
585
+ visual_dtype: torch.dtype,
586
+) -> torch.Tensor:
587
+ patches, size = x.shape
588
+ out = torch.empty(
589
+ (patches, size),
590
+ dtype=visual_dtype,
591
+ device=x.device,
592
+ )
593
+ torch.ops._xpu_C.fused_input_norm(out, x.contiguous(), weight, bias)
594
+ return out
595
+
596
+
597
+def _xpu_fused_input_norm_fake(
598
+ x: torch.Tensor,
599
+ weight: torch.Tensor | None,
600
+ bias: torch.Tensor | None,
601
+ visual_dtype: torch.dtype,
602
+) -> torch.Tensor:
603
+ patches, size = x.shape
604
+ out = torch.empty(
605
+ (patches, size),
606
+ dtype=visual_dtype,
607
+ device=x.device,
608
+ )
609
+ return out
610
+
611
+
612
@triton.jit
613
def _softplus(x):
614
return tl.where(x <= 20.0, tl.math.log(tl.math.exp(x) + 1.0), x)
@@ -1281,6 +1312,12 @@ class xpu_ops:
1312
],
1313
)
1314
1315
+ direct_register_custom_op(
1316
+ op_name="xpu_fused_input_norm",
1317
+ op_func=_xpu_fused_input_norm_impl,
1318
+ fake_impl=_xpu_fused_input_norm_fake,
1319
+ )
1320
+
1321
_OPS_REGISTERED = True
1322
1323
vllm/model_executor/models/vision.py
+14
@@ -718,6 +718,20 @@ class FusedInputNorm(nn.Module):
718
patches, size = grid_thw.shape
719
patch_size = size // self.channel
720
721
+ # On XPU, fuse the whole rescale + normalize into a single custom
722
+ # kernel. The eager path below materializes an fp32 intermediate and
723
+ # then casts back, which adds device-side compute that cancels the
724
+ # bandwidth saving of transferring uint8 pixel_values. The fused
725
+ # kernel reads uint8 directly and writes ``visual_dtype`` in one pass.
726
+ if (
727
+ current_platform.is_xpu()
728
+ and grid_thw.dtype == torch.uint8
729
+ and self.weight.dtype == torch.float32
730
+ ):
731
+ return torch.ops.vllm.xpu_fused_input_norm(
732
+ grid_thw, self.weight, self.bias, visual_dtype
733
+ )
734
+
735
# Apply the per-channel affine transform directly instead of via
736
# F.batch_norm. batch_norm dispatches to cuDNN, whose batch-norm
737
# kernels cap the batch dimension near the CUDA grid limit (~65535);