[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:
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user