diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 940812b6a4..42f605367a 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -21,6 +21,7 @@ __all__ = [ "EpBuffer", "ep_bootstrap", + "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", "ep_combine", @@ -169,6 +170,11 @@ def ep_bootstrap( _ATEXIT_REGISTERED = True +def is_ep_bootstrapped() -> bool: + """Whether EP has been initialized in this process.""" + return _BOOTSTRAPPED + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent.