A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 04:45 UTC
[XPU] Enable rms_norm/act quant fusions (#43963)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
liuzhenwei committed
Jun 3, 2026 at 00:14 UTC
3f3e2702c2a620b8c87b69746aa8e7cc400e3cd1
4 files changed
+31
-14
vllm/compilation/passes/fusion/act_quant_fusion.py
+5
-1
@@ -70,7 +70,11 @@ class ActivationQuantPattern(VllmPatternReplacement):
70
self.silu_and_mul_matcher = MatcherSiluAndMul()
71
72
def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
73
- kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}
73
+ kwargs = {
74
+ "dtype": self.quant_dtype,
75
+ "device": current_platform.device_type,
76
+ **kwargs,
77
+ }
78
return torch.empty(*args, **kwargs)
79
80
vllm/compilation/passes/fusion/rms_quant_fusion.py
+12
-4
@@ -54,19 +54,27 @@ def _rms_input_weight_dtype_match(match: pm.Match) -> bool:
54
55
56
def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor:
57
- return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda")
57
+ return torch.empty(
58
+ *args, **kwargs, dtype=torch.bfloat16, device=current_platform.device_type
59
+ )
60
61
62
def empty_fp32(*args: Any, **kwargs: Any) -> torch.Tensor:
61
- return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
63
+ return torch.empty(
64
+ *args, **kwargs, dtype=torch.float32, device=current_platform.device_type
65
+ )
66
67
68
def empty_i32(*args: Any, **kwargs: Any) -> torch.Tensor:
65
- return torch.empty(*args, **kwargs, dtype=torch.int32, device="cuda")
69
+ return torch.empty(
70
+ *args, **kwargs, dtype=torch.int32, device=current_platform.device_type
71
+ )
72
73
74
def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor:
69
- return torch.empty(*args, **kwargs, dtype=torch.int64, device="cuda")
75
+ return torch.empty(
76
+ *args, **kwargs, dtype=torch.int64, device=current_platform.device_type
77
+ )
78
79
80
RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
vllm/compilation/passes/pass_manager.py
+4
@@ -45,6 +45,10 @@ if current_platform.is_cuda():
45
from .fusion.allreduce_rms_fusion import AllReduceFusionPass
46
from .fusion.collective_fusion import AsyncTPPass
47
48
+if current_platform.is_xpu():
49
+ from .fusion.act_quant_fusion import ActivationQuantFusionPass
50
+ from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
51
+
52
from .inductor_pass import (
53
CustomGraphPass,
54
InductorPass,
vllm/platforms/xpu.py
+10
-9
@@ -197,24 +197,25 @@ class XPUPlatform(Platform):
197
)
198
199
# Disable fusion passes not yet supported on XPU.
200
+ from vllm.config.compilation import CompilationMode
201
+
202
pass_config = compilation_config.pass_config
203
fusion_passes_to_disable = {
204
"enable_sp": "Sequence parallelism",
205
"fuse_gemm_comms": "Async TP",
206
"fuse_allreduce_rms": "AllReduce + RMSNorm fusion",
205
- "fuse_norm_quant": "RMSNorm + quant fusion",
206
- "fuse_act_quant": "Activation + quant fusion",
207
"fuse_attn_quant": "Attention + quant fusion",
208
"fuse_act_padding": "Activation + padding fusion",
209
"fuse_rope_kvcache": "RoPE + KV cache fusion",
210
}
211
- for flag, feature_name in fusion_passes_to_disable.items():
212
- if getattr(pass_config, flag):
213
- logger.warning(
214
- "Feature %r is not yet supported on XPU and will be disabled.",
215
- feature_name,
216
- )
217
- setattr(pass_config, flag, False)
211
+ if compilation_config.mode != CompilationMode.NONE:
212
+ for flag, feature_name in fusion_passes_to_disable.items():
213
+ if getattr(pass_config, flag):
214
+ logger.warning(
215
+ "Feature %r is not yet supported on XPU and will be disabled.",
216
+ feature_name,
217
+ )
218
+ setattr(pass_config, flag, False)
219
220
# check and update parallel config
221
parallel_config = vllm_config.parallel_config