diff --git a/magi_compiler/_magi_register_custom_op.py b/magi_compiler/_magi_register_custom_op.py index 490c1ee..61344b6 100644 --- a/magi_compiler/_magi_register_custom_op.py +++ b/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) diff --git a/magi_compiler/api.py b/magi_compiler/api.py index 99602ea..6a4cb33 100644 --- a/magi_compiler/api.py +++ b/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, ) diff --git a/magi_compiler/profiling/materialize_inputs.py b/magi_compiler/profiling/materialize_inputs.py index 4214b55..0c657aa 100644 --- a/magi_compiler/profiling/materialize_inputs.py +++ b/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: diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index c6e3f56..490ace1 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/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 diff --git a/tests/api_tests/test_register_custom_op.py b/tests/api_tests/test_register_custom_op.py index 91722c9..d67a7cc 100644 --- a/tests/api_tests/test_register_custom_op.py +++ b/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 diff --git a/tests/feature_tests/test_profiling_estimator.py b/tests/feature_tests/test_profiling_estimator.py index 79317d5..6164539 100644 --- a/tests/feature_tests/test_profiling_estimator.py +++ b/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() diff --git a/tests/feature_tests/test_profiling_registry.py b/tests/feature_tests/test_profiling_registry.py index c88a1f0..a094f03 100644 --- a/tests/feature_tests/test_profiling_registry.py +++ b/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