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 1, 18:45 UTC
[Bugfix] Fix Qwen3.5-FP8 nightly fail. Guard fused_add_rms_norm input/weight dtype mismatch in RMSNorm + quant fusion (#44694)
Vadim Gimpelson committed
Jun 6, 2026 at 16:46 UTC
67d3792d99f345015a8bd01a3703be1469b60d7d
1 file changed
+8
-2
vllm/compilation/passes/fusion/rms_quant_fusion.py
+8
-2
@@ -39,6 +39,7 @@ FP4_DTYPE = torch.uint8
39
40
41
_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
42
+_FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default
43
44
45
# TODO: extend rmsnorm quant kernels to support mixed input/weight dtypes,
@@ -49,8 +50,13 @@ def _rms_input_weight_dtype_match(match: pm.Match) -> bool:
50
if node.target == _RMS_NORM_OP:
51
# rms_norm(x, weight, epsilon, variance_size)
52
x, weight = node.args[0], node.args[1]
52
- if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
53
- return x.meta["val"].dtype == weight.meta["val"].dtype
53
+ elif node.target == _FUSED_ADD_RMS_NORM_OP:
54
+ # fused_add_rms_norm(x, residual, weight, epsilon, variance_size)
55
+ x, weight = node.args[0], node.args[2]
56
+ else:
57
+ continue
58
+ if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
59
+ return x.meta["val"].dtype == weight.meta["val"].dtype
60
return True
61
62