[Bugfix] Fix Qwen3.5-FP8 nightly fail. Guard fused_add_rms_norm input/weight dtype mismatch in RMSNorm + quant fusion (#44694)

This commit is contained in:
Vadim Gimpelson
2026-06-06 08:46:14 -04:00
committed by GitHub
parent 00d1fb7747
commit 67d3792d99
@@ -39,6 +39,7 @@ FP4_DTYPE = torch.uint8
_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
_FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default
# 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:
if node.target == _RMS_NORM_OP:
# rms_norm(x, weight, epsilon, variance_size)
x, weight = node.args[0], node.args[1]
if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
return x.meta["val"].dtype == weight.meta["val"].dtype
elif node.target == _FUSED_ADD_RMS_NORM_OP:
# fused_add_rms_norm(x, residual, weight, epsilon, variance_size)
x, weight = node.args[0], node.args[2]
else:
continue
if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
return x.meta["val"].dtype == weight.meta["val"].dtype
return True