From 8aa20c569ff1186229c19d5589b31a7539eb7d16 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Thu, 24 Sep 2026 03:46:56 +0000 Subject: [PATCH] refactor(ep/v2): make TokOffExt public for out-of-op callers The tokoff-ext slot allocator is built by EpDispatchCombineOpHip, but aiter's MegaMoE drives EpDispatchPlan directly and needs the same buffer to bind EpArgs.tokOffPeers. Drop the underscore so that path imports it from this module the way it already imports EpDispatchPlan, rather than reaching for a private name. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/mori/ops/dispatch_combine_v2/hip_backend.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/mori/ops/dispatch_combine_v2/hip_backend.py b/python/mori/ops/dispatch_combine_v2/hip_backend.py index 0d3876486a..19d887e594 100644 --- a/python/mori/ops/dispatch_combine_v2/hip_backend.py +++ b/python/mori/ops/dispatch_combine_v2/hip_backend.py @@ -146,7 +146,7 @@ def scale_stride_bytes(scale_bytes: int) -> int: _XDB_FLAG_SLOTS = 256 -class _TokOffExt: +class TokOffExt: """Dispatch's slot allocator word outside the cco window (default-on gfx1250; MORI_EP_TOKOFF_EXT=0 opts out). One int per rank in hipExtMallocWithFlags(hipDeviceMallocUncached) memory, @@ -322,7 +322,7 @@ def _build(self, cfg, comm): self.routing_dest_map = torch.full_like(self.token_dest_map, self._null_flat) self.dest_pe_counter = torch.zeros(cfg.world_size, **i32) # Dispatch's slot allocator word lives in IPC-shared hipExtMallocWithFlags - # memory instead of the cco window (see _TokOffExt): the cco-window slot + # memory instead of the cco window (see TokOffExt): the cco-window slot # atomic serializes on newer fw/KMD stacks. gfx1250 intranode only -- the # only kernel that reads tokOffPeers. Default-on; MORI_EP_TOKOFF_EXT=0 # (or false/no/off) opts back out to the cco-window path (tokOffPeers @@ -332,7 +332,7 @@ def _build(self, cfg, comm): _tokoff_env = os.environ.get("MORI_EP_TOKOFF_EXT", "1").strip().lower() _tokoff_on = _tokoff_env not in ("0", "false", "no", "off") if self._is1250 and _tokoff_on: - self._tokoff_ext = _TokOffExt(cfg.rank, cfg.world_size, dev) + self._tokoff_ext = TokOffExt(cfg.rank, cfg.world_size, dev) self.tok_off_peers = self._tokoff_ext.peers self.total_recv = torch.zeros(1, **i32) self.dispatch_barrier = torch.zeros(1, dtype=torch.uint32, device=dev)