From 0da6e7f3d6ed50fe1262cc2cf44066989a6b4cdd Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 26 Jul 2026 08:42:14 -0400 Subject: [PATCH] [Bugfix] Reject contradictory custom-op directives (#49134) Signed-off-by: Taneem Ibrahim --- tests/compile/test_config.py | 28 ++++++++++++++++++++++++++-- vllm/config/compilation.py | 31 ++++++++++++++++++++++++------- vllm/model_executor/custom_op.py | 11 +++++++++-- 3 files changed, 59 insertions(+), 11 deletions(-) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index d822b68c50..c45907ee6e 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -83,9 +83,33 @@ def test_copy_pass(): def test_custom_op(): # proper syntax _ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"]) + _ = CompilationConfig(custom_ops=["none", "+rms_norm"]) + _ = CompilationConfig(custom_ops=["+rms_norm", "+rms_norm"]) - with pytest.raises(ValueError, match="Invalid syntax '"): - _ = CompilationConfig(custom_ops=["quant_fp8"]) + for custom_ops in (["quant_fp8"], ["+"], ["-"]): + with pytest.raises(ValueError, match="Invalid syntax '"): + CompilationConfig(custom_ops=custom_ops) + + +@pytest.mark.parametrize( + ("custom_ops", "config_kwargs", "match"), + [ + (["all", "none"], {}, "can contain only one base mode"), + ( + ["none", "+rms_norm", "-rms_norm"], + {}, + "cannot both enable and disable.*rms_norm", + ), + ( + ["-rotary_embedding"], + {"pass_config": PassConfig(enable_qk_norm_rope_fusion=True)}, + "cannot both enable and disable.*rotary_embedding", + ), + ], +) +def test_reject_contradictory_custom_ops(custom_ops, config_kwargs, match): + with pytest.raises(ValueError, match=match): + CompilationConfig(custom_ops=custom_ops, **config_kwargs) # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 8142004748..fbc5dcff40 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -904,10 +904,6 @@ class CompilationConfig: return handler(value) def __post_init__(self) -> None: - count_none = self.custom_ops.count("none") - count_all = self.custom_ops.count("all") - assert count_none + count_all <= 1, "Can only specify 'none' or 'all'" - # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2: # 1. A bug in PyTorch, fixed in 2.7: # https://github.com/pytorch/pytorch/issues/147924 @@ -1000,13 +996,28 @@ class CompilationConfig: ) for op in self.custom_ops: - if op[0] not in {"+", "-"} and op not in {"all", "none"}: + if op not in {"all", "none"} and (len(op) < 2 or op[0] not in {"+", "-"}): raise ValueError( f"Invalid syntax '{op}' for custom op, " "must be 'all', 'none', '+op' or '-op' " "(where 'op' is the registered op name)" ) + base_modes = [op for op in self.custom_ops if op in {"all", "none"}] + if len(base_modes) > 1: + raise ValueError( + "custom_ops can contain only one base mode: 'all' or 'none'" + ) + + enabled_ops = {op[1:] for op in self.custom_ops if op.startswith("+")} + disabled_ops = {op[1:] for op in self.custom_ops if op.startswith("-")} + conflicting_ops = sorted(enabled_ops & disabled_ops) + if conflicting_ops: + raise ValueError( + "custom_ops cannot both enable and disable the same operation(s): " + f"{', '.join(conflicting_ops)}. Remove either the '+' or '-' directive" + ) + # Currently only eager and inductor backend are supported. # for piecewise compilation. Custom backends are not supported for # piecewise compilation. Update when more backends are supported. @@ -1343,10 +1354,16 @@ class CompilationConfig: ) def is_custom_op_enabled(self, op: str) -> bool: - if "all" in self.custom_ops: + count_all = self.custom_ops.count("all") + count_none = self.custom_ops.count("none") + if count_all + count_none != 1: + raise ValueError( + "custom_ops must contain exactly one base mode: 'all' or 'none'" + ) + + if count_all: return f"-{op}" not in self.custom_ops - assert "none" in self.custom_ops return f"+{op}" in self.custom_ops def resolve_cudagraph_mode_and_sizes( diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py index a1514c9206..bf98e46e85 100644 --- a/vllm/model_executor/custom_op.py +++ b/vllm/model_executor/custom_op.py @@ -284,7 +284,11 @@ class CustomOp(nn.Module): enabled = f"+{cls.name}" in custom_ops disabled = f"-{cls.name}" in custom_ops - assert not (enabled and disabled), f"Cannot enable and disable {cls.name}" + if enabled and disabled: + raise ValueError( + "custom_ops cannot both enable and disable the same operation: " + f"{cls.name}. Remove either the '+' or '-' directive" + ) return (CustomOp.default_on() or enabled) and not disabled @@ -299,7 +303,10 @@ class CustomOp(nn.Module): compilation_config = get_cached_compilation_config() count_none = compilation_config.custom_ops.count("none") count_all = compilation_config.custom_ops.count("all") - assert count_none + count_all == 1 + if count_none + count_all != 1: + raise ValueError( + "custom_ops must contain exactly one base mode: 'all' or 'none'" + ) return not count_none > 0 or count_all > 0