fix error

Author: baonudesifeizhai

tests/compile/piecewise/test_simple.py

@@ -143,7 +143,7 @@ def test_simple_piecewise_compile(use_inductor):
 
 @torch.inference_mode()
 @pytest.mark.parametrize("splitting_ops", [["silly::attention"], []])
-def test_simple_inductor_graph_partition(splitting_ops):
+def test_simple_inductor_graph_partition(monkeypatch, splitting_ops):
     if not is_torch_equal_or_newer("2.9.0.dev"):
         pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
 

vllm/compilation/partition_rules.py

@@ -37,29 +37,44 @@ def resolve_defined_ops(op_names: list[str]) -> list[torch._ops.OpOverload]:
     Returns:
         List of successfully resolved operator overloads
     """
-    resolved = []
+    resolved: list[torch._ops.OpOverload] = []
     for op_name in op_names:
-        try:
-            op = lookup_op(op_name)
-            # Handle OpOverloadPacket: get .default if available
-            if hasattr(op, "default"):
-                resolved.append(op.default)
+        overload: torch._ops.OpOverload | None = None
+        candidate_names = [op_name]
+
+        # When the caller omits an explicit overload (e.g. "namespace::op"),
+        # also try the conventional ".default" suffix.
+        if "." not in op_name.split("::")[-1]:
+            candidate_names.append(f"{op_name}.default")
+
+        for candidate in candidate_names:
+            try:
+                op = lookup_op(candidate)
+            except Exception:
+                continue
+
+            # lookup_op may return either an OpOverload (desired) or an
+            # OpOverloadPacket (collection of overloads).
+            if hasattr(op, "overloads"):
+                overloads = list(op.overloads())
+                if "default" in overloads:
+                    overload = op.default
+                elif len(overloads) == 1:
+                    overload = getattr(op, overloads[0])
+                else:
+                    logger.warning(
+                        "Operator '%s' has multiple overloads (%s); please "
+                        "specify the desired overload explicitly.",
+                        candidate,
+                        ", ".join(overloads),
+                    )
             else:
-                resolved.append(op)
-        except Exception:
-            # If lookup fails and no overload specified, try with .default
-            if "." not in op_name.split("::")[-1]:
-                try:
-                    op = lookup_op(f"{op_name}.default")
-                    resolved.append(op)
-                    continue
-                except Exception:
-                    pass
-            # Skip operators that don't exist
-            logger.warning(
-                "Failed to resolve operator for Inductor partition: %s", op_name
-            )
-            continue
+                overload = op  # Already an OpOverload
+
+            if overload is not None:
+                break
+
+        resolved.append(overload)
 
     return resolved