[Dev] Drop has_internal_collective from magi_register_custom_op

magi_compiler/_magi_register_custom_op.py

@@ -58,6 +58,8 @@
 import torch
 import torch.utils._pytree as pytree
 
+from magi_compiler.profiling import register_materialize_inputs
+
 from .config import get_compile_config
 from .utils.logger import magi_logger
 
@@ -886,14 +888,6 @@ class _DataclassRuntimeAdapter:
 # ==============================================================================
 
 
-def _maybe_register_op_profiling(op_name: str, has_internal_collective: bool, materialize_inputs: Callable | None) -> None:
-    if not has_internal_collective and materialize_inputs is None:
-        return
-    from magi_compiler.profiling import register_materialize_inputs
-
-    register_materialize_inputs(op_name, materialize_inputs, has_internal_collective=has_internal_collective)
-
-
 def _magi_register_custom_op_impl(
     name: str | None = None,
     mutates_args: tuple[str, ...] = (),
@@ -902,7 +896,6 @@ def _magi_register_custom_op_impl(
     backward_fn: Callable | None = None,
     is_compute_sensitive: bool = False,
     is_subgraph_boundary: bool = False,
-    has_internal_collective: bool = False,
     materialize_inputs: Callable | None = None,
 ):
     def decorator(fn: Callable) -> Callable:
@@ -913,7 +906,8 @@ def decorator(fn: Callable) -> Callable:
             get_compile_config().recompute_config.custom_compute_sensitive_ops.append(op_name)
         if is_subgraph_boundary:
             get_compile_config().splitting_ops.append(op_name)
-        _maybe_register_op_profiling(op_name, has_internal_collective, materialize_inputs)
+
+        register_materialize_inputs(op_name, materialize_inputs)
 
         _validate_op_signature_constraints(fn)
         original_sig, lowered_sig, param_mapping_tree = _lower_op_signature(fn)

magi_compiler/api.py

@@ -202,7 +202,6 @@ def magi_register_custom_op(
     backward_fn: Callable | None = None,
     is_compute_sensitive: bool = False,
     is_subgraph_boundary: bool = False,
-    has_internal_collective: bool = False,
     materialize_inputs: Callable | None = None,
 ):
     """
@@ -235,9 +234,6 @@ def magi_register_custom_op(
             ops are prioritised for saving rather than recomputing.
         is_subgraph_boundary: Split the FX graph at this op during compilation.
             Each sub-graph between boundary ops is compiled independently.
-        has_internal_collective: The op issues NCCL internally (CP all-to-all,
-            EP dispatch, ...).  The profiler then lockstep-replays
-            the whole custom op with a fixed iteration count.
         materialize_inputs: Optional hook with the **same signature as the op**.
             MagiCompiler generic-realizes every argument, then calls
             ``fn(*args, **kwargs)``. Use this to rebuild value-dependent metadata
@@ -323,6 +319,5 @@ def magi_register_custom_op(
         backward_fn=backward_fn,
         is_compute_sensitive=is_compute_sensitive,
         is_subgraph_boundary=is_subgraph_boundary,
-        has_internal_collective=has_internal_collective,
         materialize_inputs=materialize_inputs,
     )

magi_compiler/profiling/materialize_inputs.py

@@ -33,7 +33,7 @@ def _attn_cp_inputs(q, k, v, cp_split_sizes):
         seq = int(q.shape[1] if q.dim() == 4 else q.shape[0])
         return q, k, v, [seq] * len(cp_split_sizes)
 
-    @magi_register_custom_op("mylib::attn_cp", materialize_inputs=_attn_cp_inputs, has_internal_collective=True)
+    @magi_register_custom_op("mylib::attn_cp", materialize_inputs=_attn_cp_inputs)
     def attn_cp(q, k, v, cp_split_sizes):
         ...
 
@@ -51,24 +51,23 @@ def attn_cp(q, k, v, cp_split_sizes):
 # reconstructed from generic size-hinted tensors (value-dependent metadata).
 # Look for a more general approach that does not require model-side hooks.
 
-# op name (OpOverload string, e.g. "mylib::attn_cp") -> hook; plus the set of ops
-# that issue an internal collective (need fixed-iter lockstep replay).
+# op name (OpOverload string, e.g. "mylib::attn_cp") -> optional hook.
+# Every registered name is treated as an internal-collective op (fixed-iter
+# lockstep replay); ``fn`` is only needed when generic realize is not enough.
 _MATERIALIZE_INPUT_HOOKS: dict[str, Callable] = {}
 _INTERNAL_COLLECTIVE_OPS: set[str] = set()
 
 
-def register_materialize_inputs(op_name: str, fn: Callable | None = None, *, has_internal_collective: bool = False) -> None:
-    """Register a same-signature replay-input builder for ``op_name``.
+def register_materialize_inputs(op_name: str, fn: Callable | None = None) -> None:
+    """Mark ``op_name`` for lockstep replay; optionally attach a same-signature hook.
 
-    ``fn`` is optional when only ``has_internal_collective`` is needed (generic
-    realize is already valid).  ``has_internal_collective``: replay with a fixed
-    iteration count under barriers (an adaptive count would desync the internal
-    NCCL op across ranks).
+    ``fn`` rebuilds value-dependent metadata after generic realize.  Omit it when
+    the generic tensors are already valid.  Registration always flags the op as
+    issuing an internal collective (adaptive iter counts would desync NCCL).
     """
     if fn is not None:
         _MATERIALIZE_INPUT_HOOKS[op_name] = fn
-    if has_internal_collective:
-        _INTERNAL_COLLECTIVE_OPS.add(op_name)
+    _INTERNAL_COLLECTIVE_OPS.add(op_name)
 
 
 def get_materialize_inputs_hook(op_name: str) -> Callable | None:

magi_compiler/profiling/runtime_estimator.py

@@ -64,8 +64,8 @@
 
 def snode_issues_collective(snode: BaseSchedulerNode) -> bool:
     """
-    True if replaying / running this snode issues NCCL (collective AG, or a custom
-    op registered with ``has_internal_collective``).
+    True if replaying / running this snode issues NCCL (collective AG, or a
+    ``magi_register_custom_op`` extern).
     """
     if contains_collective(snode):
         return True
@@ -288,9 +288,9 @@ def _op_name(target) -> str:
 
 
 def _extern_has_internal_collective(snode: BaseSchedulerNode) -> bool:
-    """Ops that issue collectives internally (declared via
-    ``register_materialize_inputs(..., has_internal_collective=True)``) must be
-    measured with fixed iterations under a barrier."""
+    """Ops registered via ``register_materialize_inputs`` are treated as
+    issuing an internal collective and must be measured with fixed iterations
+    under a barrier."""
     node = getattr(snode, "node", None)
     origin = node.get_origin_node() if (node is not None and hasattr(node, "get_origin_node")) else None
     target = getattr(origin, "target", None) if origin is not None else None

tests/api_tests/test_register_custom_op.py

@@ -78,10 +78,7 @@ def _materialize(q, k, v, cp_split_sizes):
             return q, k, v, [seq] * len(cp_split_sizes)
 
         @magi_register_custom_op(
-            name="test::materialize_same_sig_op",
-            infer_output_meta_fn=["q"],
-            materialize_inputs=_materialize,
-            has_internal_collective=True,
+            name="test::materialize_same_sig_op", infer_output_meta_fn=["q"], materialize_inputs=_materialize
         )
         def _op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cp_split_sizes: list[int]) -> torch.Tensor:
             return q
@@ -97,11 +94,11 @@ def _op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cp_split_sizes: list[
             _MATERIALIZE_INPUT_HOOKS.pop("test::materialize_same_sig_op", None)
             _INTERNAL_COLLECTIVE_OPS.discard("test::materialize_same_sig_op")
 
-    def test_has_internal_collective_without_hook(self):
+    def test_decorator_marks_internal_collective_without_hook(self):
         from magi_compiler.profiling import get_materialize_inputs_hook, op_has_internal_collective
         from magi_compiler.profiling.materialize_inputs import _INTERNAL_COLLECTIVE_OPS
 
-        @magi_register_custom_op(name="test::moe_flag_only_op", has_internal_collective=True)
+        @magi_register_custom_op(name="test::moe_flag_only_op")
         def _moe(x: torch.Tensor) -> torch.Tensor:
             return x
 

tests/feature_tests/test_profiling_estimator.py

@@ -259,7 +259,7 @@ def get_estimated_runtime(self):
 
 @pytest.mark.parametrize("sync", [True, False])
 def test_internal_collective_extern_not_measured_in_sync_warmup(monkeypatch, sync):
-    """An extern registered with has_internal_collective=True must NOT be measured
+    """An extern registered via ``register_materialize_inputs`` must NOT be measured
     by __call__ in sync mode (the warm-up runs per-rank without barriers; the
     adaptive benchmarker would issue rank-dependent numbers of the internal NCCL op
     -> hang).  It is seeded analytical + stashed for warm_and_sync.  In non-sync
@@ -274,7 +274,7 @@ def _boom(*a, **k):
         measured["called"] = True
         raise AssertionError("must not be measured in sync warm-up")
 
-    register_materialize_inputs("aten::mm", has_internal_collective=True)
+    register_materialize_inputs("aten::mm")
     monkeypatch.setattr(re_mod, "_measure_extern", _boom)
     try:
         est = ProfilingRuntimeEstimator()

tests/feature_tests/test_profiling_registry.py

@@ -48,18 +48,16 @@ def hook(q, k, v, cp_split_sizes):
 
     register_materialize_inputs("ns::op_a", hook)
     assert get_materialize_inputs_hook("ns::op_a") is hook
-    # not flagged as internal-collective by default
-    assert op_has_internal_collective("ns::op_a") is False
+    assert op_has_internal_collective("ns::op_a") is True
 
 
-def test_internal_collective_flag(clean_registry):
-    register_materialize_inputs("ns::coll_op", has_internal_collective=True)
+def test_register_without_hook_still_marks_collective(clean_registry):
+    register_materialize_inputs("ns::coll_op")
     assert op_has_internal_collective("ns::coll_op") is True
     assert get_materialize_inputs_hook("ns::coll_op") is None
 
-    # a hook registered WITHOUT the flag must not be marked
     register_materialize_inputs("ns::plain_op", lambda *a, **k: None)
-    assert op_has_internal_collective("ns::plain_op") is False
+    assert op_has_internal_collective("ns::plain_op") is True
 
 
 def test_register_overrides_previous(clean_registry):
@@ -77,7 +75,7 @@ def hook2(q):
 
 def test_hook_is_callable_returning_none(clean_registry):
     """A no-op hook (returns None -> keep generic realize) is valid."""
-    register_materialize_inputs("ns::noop", lambda *a, **k: None, has_internal_collective=True)
+    register_materialize_inputs("ns::noop", lambda *a, **k: None)
     hook = get_materialize_inputs_hook("ns::noop")
     assert hook(object()) is None
     assert op_has_internal_collective("ns::noop") is True