diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index e85149a879..670349a08b 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -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