From 67d3792d99f345015a8bd01a3703be1469b60d7d Mon Sep 17 00:00:00 2001
From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Date: Sat, 6 Jun 2026 16:46:14 +0400
Subject: [PATCH] [Bugfix] Fix Qwen3.5-FP8 nightly fail. Guard
 fused_add_rms_norm input/weight dtype mismatch in RMSNorm + quant fusion
 (#44694)

---
 vllm/compilation/passes/fusion/rms_quant_fusion.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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