[Bugfix] Reject contradictory custom-op directives (#49134)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jul 26, 2026 at 08:42 UTC
0da6e7f3d6ed50fe1262cc2cf44066989a6b4cdd
3 files changed
+59
-11
tests/compile/test_config.py
+26
-2
@@ -83,9 +83,33 @@ def test_copy_pass():
83
def test_custom_op():
84
# proper syntax
85
_ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"])
86
+ _ = CompilationConfig(custom_ops=["none", "+rms_norm"])
87
+ _ = CompilationConfig(custom_ops=["+rms_norm", "+rms_norm"])
88
87
- with pytest.raises(ValueError, match="Invalid syntax '"):
88
- _ = CompilationConfig(custom_ops=["quant_fp8"])
89
+ for custom_ops in (["quant_fp8"], ["+"], ["-"]):
90
+ with pytest.raises(ValueError, match="Invalid syntax '"):
91
+ CompilationConfig(custom_ops=custom_ops)
92
+
93
+
94
+@pytest.mark.parametrize(
95
+ ("custom_ops", "config_kwargs", "match"),
96
+ [
97
+ (["all", "none"], {}, "can contain only one base mode"),
98
+ (
99
+ ["none", "+rms_norm", "-rms_norm"],
100
+ {},
101
+ "cannot both enable and disable.*rms_norm",
102
+ ),
103
+ (
104
+ ["-rotary_embedding"],
105
+ {"pass_config": PassConfig(enable_qk_norm_rope_fusion=True)},
106
+ "cannot both enable and disable.*rotary_embedding",
107
+ ),
108
+ ],
109
+)
110
+def test_reject_contradictory_custom_ops(custom_ops, config_kwargs, match):
111
+ with pytest.raises(ValueError, match=match):
112
+ CompilationConfig(custom_ops=custom_ops, **config_kwargs)
113
114
115
# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
vllm/config/compilation.py
+24
-7
@@ -904,10 +904,6 @@ class CompilationConfig:
904
return handler(value)
905
906
def __post_init__(self) -> None:
907
- count_none = self.custom_ops.count("none")
908
- count_all = self.custom_ops.count("all")
909
- assert count_none + count_all <= 1, "Can only specify 'none' or 'all'"
910
-
907
# TODO(zou3519/luka): There are 2 issues with auto-functionalization V2:
908
# 1. A bug in PyTorch, fixed in 2.7:
909
# https://github.com/pytorch/pytorch/issues/147924
@@ -1000,13 +996,28 @@ class CompilationConfig:
996
)
997
998
for op in self.custom_ops:
1003
- if op[0] not in {"+", "-"} and op not in {"all", "none"}:
999
+ if op not in {"all", "none"} and (len(op) < 2 or op[0] not in {"+", "-"}):
1000
raise ValueError(
1001
f"Invalid syntax '{op}' for custom op, "
1002
"must be 'all', 'none', '+op' or '-op' "
1003
"(where 'op' is the registered op name)"
1004
)
1005
1006
+ base_modes = [op for op in self.custom_ops if op in {"all", "none"}]
1007
+ if len(base_modes) > 1:
1008
+ raise ValueError(
1009
+ "custom_ops can contain only one base mode: 'all' or 'none'"
1010
+ )
1011
+
1012
+ enabled_ops = {op[1:] for op in self.custom_ops if op.startswith("+")}
1013
+ disabled_ops = {op[1:] for op in self.custom_ops if op.startswith("-")}
1014
+ conflicting_ops = sorted(enabled_ops & disabled_ops)
1015
+ if conflicting_ops:
1016
+ raise ValueError(
1017
+ "custom_ops cannot both enable and disable the same operation(s): "
1018
+ f"{', '.join(conflicting_ops)}. Remove either the '+' or '-' directive"
1019
+ )
1020
+
1021
# Currently only eager and inductor backend are supported.
1022
# for piecewise compilation. Custom backends are not supported for
1023
# piecewise compilation. Update when more backends are supported.
@@ -1343,10 +1354,16 @@ class CompilationConfig:
1354
)
1355
1356
def is_custom_op_enabled(self, op: str) -> bool:
1346
- if "all" in self.custom_ops:
1357
+ count_all = self.custom_ops.count("all")
1358
+ count_none = self.custom_ops.count("none")
1359
+ if count_all + count_none != 1:
1360
+ raise ValueError(
1361
+ "custom_ops must contain exactly one base mode: 'all' or 'none'"
1362
+ )
1363
+
1364
+ if count_all:
1365
return f"-{op}" not in self.custom_ops
1366
1349
- assert "none" in self.custom_ops
1367
return f"+{op}" in self.custom_ops
1368
1369
def resolve_cudagraph_mode_and_sizes(
vllm/model_executor/custom_op.py
+9
-2
@@ -284,7 +284,11 @@ class CustomOp(nn.Module):
284
285
enabled = f"+{cls.name}" in custom_ops
286
disabled = f"-{cls.name}" in custom_ops
287
- assert not (enabled and disabled), f"Cannot enable and disable {cls.name}"
287
+ if enabled and disabled:
288
+ raise ValueError(
289
+ "custom_ops cannot both enable and disable the same operation: "
290
+ f"{cls.name}. Remove either the '+' or '-' directive"
291
+ )
292
293
return (CustomOp.default_on() or enabled) and not disabled
294
@@ -299,7 +303,10 @@ class CustomOp(nn.Module):
303
compilation_config = get_cached_compilation_config()
304
count_none = compilation_config.custom_ops.count("none")
305
count_all = compilation_config.custom_ops.count("all")
302
- assert count_none + count_all == 1
306
+ if count_none + count_all != 1:
307
+ raise ValueError(
308
+ "custom_ops must contain exactly one base mode: 'all' or 'none'"
309
+ )
310
311
return not count_none > 0 or count_all > 0
312