fix(atom): relay KV events from data-parallel ranks

infera/engine/atom/hooks/kv_event_bootstrap.py

@@ -14,8 +14,8 @@
 
 The work is gated on ``INFERA_ATOM_KV_EVENTS_ENDPOINT`` so that unrelated
 Python invocations (pip, debug shells, the launcher itself before it spawns
-ATOM) pay nothing. The infera ATOM launcher sets that env var (to the ZMQ
-bind address) only for the ATOM subprocess it spawns.
+ATOM) pay nothing. The Infera ATOM launcher sets that env var to the local
+relay ingress and enables connect mode only for the ATOM subprocess it spawns.
 """
 
 from __future__ import annotations
@@ -30,7 +30,8 @@ def _activate() -> None:
     try:
         from infera.engine.atom.hooks.kv_events import arm_kv_event_hooks
 
-        arm_kv_event_hooks(endpoint)
+        connect = os.environ.get("INFERA_ATOM_KV_EVENTS_CONNECT") == "1"
+        arm_kv_event_hooks(endpoint, connect=connect)
     except Exception:
         # A sitecustomize/.pth hook must never take down an unrelated
         # interpreter; surface the traceback and continue.

infera/engine/atom/hooks/kv_events.py

@@ -23,11 +23,11 @@
 change is needed — ATOM workers light up KV-aware routing exactly like
 vLLM/SGLang ones.
 
-The hooks are installed inside the ATOM ``EngineCore`` subprocess (which is
-where ``BlockManager`` actually lives) via
-:mod:`infera.engine.atom.hooks.kv_event_bootstrap`. The PUB socket binds
-lazily on the first ``BlockManager`` instantiation, so only the process that
-owns the block manager binds the port.
+The hooks are installed inside every ATOM ``EngineCore`` subprocess (where
+each data-parallel rank owns a ``BlockManager``) via
+:mod:`infera.engine.atom.hooks.kv_event_bootstrap`. Each PUB socket connects
+to the worker-local relay, which multiplexes all ranks onto the single endpoint
+advertised for the logical ATOM worker.
 """
 
 from __future__ import annotations
@@ -101,15 +101,16 @@ def _check_block_manager_compat(BlockManager) -> str | None:
 
 class _Publisher:
     """Thin ZMQ PUB wrapper that encodes batches with the router's msgspec
-    structs (guaranteeing byte-for-byte wire compatibility) and binds lazily.
+    structs (guaranteeing byte-for-byte wire compatibility) and opens lazily.
 
     All publishing happens from the single EngineCore scheduler thread that
     drives ``BlockManager``, so a plain (non-async, non-thread-safe) PUB
     socket is safe.
     """
 
-    def __init__(self, bind_endpoint: str) -> None:
-        self._bind_endpoint = bind_endpoint
+    def __init__(self, endpoint: str, *, connect: bool = False) -> None:
+        self._endpoint = endpoint
+        self._connect = connect
         self._sock = None
         self._encoder = None
         self._events_mod = None
@@ -130,9 +131,16 @@ def ensure_bound(self) -> None:
         # the router rebuilds its view from a fresh subscription if it ever
         # falls behind; we never want a slow subscriber to block the engine.
         sock.setsockopt(zmq.SNDHWM, 100_000)
-        sock.bind(self._bind_endpoint)
+        if self._connect:
+            sock.connect(self._endpoint)
+        else:
+            sock.bind(self._endpoint)
         self._sock = sock
-        logger.info("ATOM kv-events: PUB bound at %s", self._bind_endpoint)
+        logger.info(
+            "ATOM kv-events: PUB %s %s",
+            "connected to" if self._connect else "bound at",
+            self._endpoint,
+        )
 
     def _send_batch(self, events: list) -> None:
         if self._sock is None or not events:
@@ -174,9 +182,9 @@ def block_removed(self, block_hash: int) -> None:
 def _patch_block_manager(BlockManager, publisher: _Publisher) -> None:
     """Wrap ``BlockManager`` methods to publish KV cache events.
 
-    Idempotent. The PUB socket only binds when a ``BlockManager`` is actually
-    instantiated (i.e. inside the EngineCore subprocess), so this is safe to
-    apply in any process that imports the class.
+    Idempotent. The PUB socket is only opened when a ``BlockManager`` is
+    actually instantiated (i.e. inside an EngineCore subprocess), so this is
+    safe to apply in any process that imports the class.
     """
     if getattr(BlockManager, "_infera_kv_patched", False):
         return
@@ -285,7 +293,7 @@ def patched_allocate_block(self, block_id):
     BlockManager.hash_blocks = patched_hash_blocks
     BlockManager._allocate_block = patched_allocate_block
     BlockManager._infera_kv_patched = True
-    logger.info("ATOM kv-events: BlockManager hooks installed (bind=%s)", publisher._bind_endpoint)
+    logger.info("ATOM kv-events: BlockManager hooks installed (endpoint=%s)", publisher._endpoint)
 
 
 class _DeferredPatchFinder:
@@ -300,8 +308,9 @@ class _DeferredPatchFinder:
     point ``site`` init is long done.
     """
 
-    def __init__(self, endpoint: str) -> None:
+    def __init__(self, endpoint: str, *, connect: bool = False) -> None:
         self._endpoint = endpoint
+        self._connect = connect
         self._done = False
 
     def find_spec(self, fullname, path, target=None):
@@ -314,21 +323,22 @@ def find_spec(self, fullname, path, target=None):
             return None
         self._done = True
         endpoint = self._endpoint
+        connect = self._connect
         loader = spec.loader
         orig_exec = loader.exec_module
 
-        def exec_module(module, _orig=orig_exec, _ep=endpoint):
+        def exec_module(module, _orig=orig_exec, _ep=endpoint, _connect=connect):
             _orig(module)
             try:
-                _patch_block_manager(module.BlockManager, _Publisher(_ep))
+                _patch_block_manager(module.BlockManager, _Publisher(_ep, connect=_connect))
             except Exception:
                 traceback.print_exc()
 
         loader.exec_module = exec_module
         return spec
 
 
-def arm_kv_event_hooks(bind_endpoint: str) -> None:
+def arm_kv_event_hooks(endpoint: str, *, connect: bool = False) -> None:
     """Arrange for ATOM's ``BlockManager`` to publish KV cache events.
 
     If the block-manager module is already imported, patch immediately;
@@ -338,9 +348,13 @@ def arm_kv_event_hooks(bind_endpoint: str) -> None:
     """
     mod = sys.modules.get(_BLOCK_MANAGER_MODULE)
     if mod is not None:
-        _patch_block_manager(mod.BlockManager, _Publisher(bind_endpoint))
+        _patch_block_manager(mod.BlockManager, _Publisher(endpoint, connect=connect))
         return
     if any(isinstance(f, _DeferredPatchFinder) for f in sys.meta_path):
         return
-    sys.meta_path.insert(0, _DeferredPatchFinder(bind_endpoint))
-    logger.info("ATOM kv-events: armed deferred BlockManager patch (bind=%s)", bind_endpoint)
+    sys.meta_path.insert(0, _DeferredPatchFinder(endpoint, connect=connect))
+    logger.info(
+        "ATOM kv-events: armed deferred BlockManager patch (endpoint=%s connect=%s)",
+        endpoint,
+        connect,
+    )

infera/engine/atom/kv_event_proxy.py

@@ -0,0 +1,99 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+###############################################################################
+"""Process-safe relay for ATOM KV-cache events.
+
+ATOM creates one ``EngineCore`` process per data-parallel rank. Each process
+owns an independent ``BlockManager`` and therefore needs to publish its own KV
+events, while Infera advertises one logical KV-event endpoint per ATOM worker.
+The relay gives all EngineCore publishers a local XSUB ingress and forwards
+their streams through one externally advertised XPUB endpoint.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+import uuid
+
+import zmq
+
+logger = logging.getLogger(__name__)
+
+
+class AtomKvEventProxy:
+    """Relay multiple child PUB sockets through one advertised endpoint."""
+
+    def __init__(self, external_bind_endpoint: str) -> None:
+        suffix = f"{os.getpid()}-{uuid.uuid4().hex[:8]}"
+        self.external_bind_endpoint = external_bind_endpoint
+        self.ingress_endpoint = f"ipc:///tmp/infera-atom-kv-{suffix}.sock"
+        self._control_endpoint = f"inproc://infera-atom-kv-control-{suffix}"
+        self._ready = threading.Event()
+        self._error: BaseException | None = None
+        self._thread: threading.Thread | None = None
+        self._controller: zmq.Socket | None = None
+
+    def start(self, timeout: float = 10.0) -> None:
+        if self._thread is not None:
+            return
+        self._thread = threading.Thread(
+            target=self._run,
+            name="atom-kv-event-proxy",
+            daemon=True,
+        )
+        self._thread.start()
+        if not self._ready.wait(timeout):
+            raise TimeoutError("ATOM KV-event proxy did not start")
+        if self._error is not None:
+            raise RuntimeError("ATOM KV-event proxy failed to start") from self._error
+
+        controller = zmq.Context.instance().socket(zmq.PAIR)
+        controller.setsockopt(zmq.LINGER, 0)
+        controller.connect(self._control_endpoint)
+        self._controller = controller
+
+    def stop(self, timeout: float = 5.0) -> None:
+        if self._thread is None:
+            return
+        if self._controller is not None:
+            try:
+                self._controller.send(b"TERMINATE")
+            except zmq.ZMQError:
+                logger.exception("failed to stop ATOM KV-event proxy cleanly")
+            self._controller.close(linger=0)
+            self._controller = None
+        self._thread.join(timeout)
+        if self._thread.is_alive():
+            logger.error("ATOM KV-event proxy did not stop within %.1fs", timeout)
+        self._thread = None
+
+    def _run(self) -> None:
+        context = zmq.Context.instance()
+        ingress = context.socket(zmq.XSUB)
+        egress = context.socket(zmq.XPUB)
+        control = context.socket(zmq.PAIR)
+        for socket in (ingress, egress, control):
+            socket.setsockopt(zmq.LINGER, 0)
+        try:
+            ingress.bind(self.ingress_endpoint)
+            egress.bind(self.external_bind_endpoint)
+            control.bind(self._control_endpoint)
+            logger.info(
+                "ATOM kv-events: proxy ingress=%s advertise-bind=%s",
+                self.ingress_endpoint,
+                self.external_bind_endpoint,
+            )
+            self._ready.set()
+            zmq.proxy_steerable(ingress, egress, None, control)
+        except BaseException as exc:
+            self._error = exc
+            self._ready.set()
+            logger.exception("ATOM KV-event proxy failed")
+        finally:
+            ingress.close(linger=0)
+            egress.close(linger=0)
+            control.close(linger=0)

infera/engine/atom/worker.py

@@ -11,14 +11,13 @@
 :class:`EngineConfig` back to the launcher for etcd registration.
 
 KV-aware routing is **off by default**. ATOM has no native KV-event stream,
-so when the operator opts in with ``--enable-kv-events`` infera installs a
+so when the operator opts in with ``--enable-kv-events`` Infera installs a
 BlockManager hook (see :mod:`infera.engine.atom.hooks.kv_events`) inside the
-ATOM subprocess that republishes ATOM's prefix-cache index on a ZMQ PUB
-socket in the router's wire format: the launcher passes the bind endpoint to
-the subprocess via ``INFERA_ATOM_KV_EVENTS_ENDPOINT`` and advertises the
-reachable endpoint + block size in :class:`EngineConfig`. When not enabled
-(the default) the config carries no ``kv_events_endpoint`` and the router
-routes the worker round-robin.
+ATOM subprocess that republishes ATOM's prefix-cache index on ZMQ PUB sockets
+in the router's wire format. A worker-local proxy combines every data-parallel
+rank into the one reachable endpoint advertised in :class:`EngineConfig`.
+When not enabled (the default), the config carries no ``kv_events_endpoint``
+and the router routes the worker round-robin.
 """
 
 from __future__ import annotations
@@ -34,6 +33,7 @@
 import httpx
 
 from infera.common.worker_pool import DisaggMode, EngineType
+from infera.engine.atom.kv_event_proxy import AtomKvEventProxy
 from infera.engine.base import BaseEngine, EngineConfig
 
 logger = logging.getLogger(__name__)
@@ -60,13 +60,15 @@ def __init__(
         self.port = port
         self.advertise_host = advertise_host or host
         # ``kv_events_endpoint`` is what the router connects to (advertised);
-        # ``kv_events_bind`` is what the ATOM subprocess binds (``tcp://*:port``).
+        # ``kv_events_bind`` is bound by a local relay. Every data-parallel
+        # EngineCore connects its own publisher to the relay ingress.
         self.kv_events_endpoint = kv_events_endpoint
         self.kv_events_bind = kv_events_bind
         self.kv_block_size = kv_block_size
         self.disagg_mode = disagg_mode
         self.disagg_meta: dict[str, Any] = dict(disagg_meta or {})
         self._proc: subprocess.Popen | None = None
+        self._kv_event_proxy: AtomKvEventProxy | None = None
 
     async def start(self) -> EngineConfig:
         cmd = [
@@ -77,15 +79,19 @@ async def start(self) -> EngineConfig:
         ]
         logger.info("spawning atom subprocess: %s", " ".join(cmd))
         env = os.environ.copy()
-        # Hand the ZMQ bind endpoint to the ATOM subprocess. The site-startup
-        # hook (infera.engine.atom.hooks.kv_event_bootstrap, run via a .pth in
-        # every interpreter — including ATOM's spawned EngineCore) reads this
-        # and installs the BlockManager KV-event publisher.
+        # A TP+DP ATOM worker owns one BlockManager per EngineCore process.
+        # Relay all of their PUB streams through the one endpoint advertised
+        # for this logical worker; otherwise every DP rank races to bind the
+        # same TCP port and all but one EngineCore fail during startup.
         if self.kv_events_bind:
-            env["INFERA_ATOM_KV_EVENTS_ENDPOINT"] = self.kv_events_bind
+            self._kv_event_proxy = AtomKvEventProxy(self.kv_events_bind)
+            self._kv_event_proxy.start()
+            env["INFERA_ATOM_KV_EVENTS_ENDPOINT"] = self._kv_event_proxy.ingress_endpoint
+            env["INFERA_ATOM_KV_EVENTS_CONNECT"] = "1"
             logger.info(
-                "ATOM kv-events enabled: bind=%s advertise=%s block_size=%s",
+                "ATOM kv-events enabled: relay=%s ingress=%s advertise=%s block_size=%s",
                 self.kv_events_bind,
+                self._kv_event_proxy.ingress_endpoint,
                 self.kv_events_endpoint,
                 self.kv_block_size,
             )
@@ -114,36 +120,41 @@ async def start(self) -> EngineConfig:
 
     async def stop(self) -> None:
         logger.info("ATOM engine stopping")
-        if self._proc is None:
-            return
-        # The subprocess is a session leader (start_new_session=True), so its
-        # pid doubles as the process-group id for the whole ATOM tree.
-        pgid = self._proc.pid
-        if self._proc.poll() is None:
+        try:
+            if self._proc is None:
+                return
+            # The subprocess is a session leader (start_new_session=True), so
+            # its pid doubles as the process-group id for the whole ATOM tree.
+            pgid = self._proc.pid
+            if self._proc.poll() is None:
+                try:
+                    os.killpg(pgid, signal.SIGTERM)
+                except ProcessLookupError:
+                    return
+                # Graceful window kept short so we always reach the SIGKILL
+                # sweep within the launcher's own teardown budget (the parent
+                # harness SIGKILLs this process group ~25s after SIGTERM).
+                for _ in range(15):
+                    if self._proc.poll() is not None:
+                        break
+                    await asyncio.sleep(1)
+            # The leader (ATOM's openai_server) exiting does NOT mean the whole
+            # group is gone: ATOM/AITER helpers — e.g. the shared-memory
+            # broadcast worker — can hang and ignore SIGTERM, holding GPU VRAM
+            # until the container dies and starving the next worker (HIP-OOM).
+            # Always sweep the group with SIGKILL so VRAM is released.
             try:
-                os.killpg(pgid, signal.SIGTERM)
+                os.killpg(pgid, signal.SIGKILL)
             except ProcessLookupError:
-                return
-            # Graceful window kept short so we always reach the SIGKILL sweep
-            # within the launcher's own teardown budget (the parent harness
-            # SIGKILLs this process group ~25s after SIGTERM).
-            for _ in range(15):
-                if self._proc.poll() is not None:
-                    break
-                await asyncio.sleep(1)
-        # The leader (ATOM's openai_server) exiting does NOT mean the whole
-        # group is gone: ATOM/AITER helpers — e.g. the shared-memory broadcast
-        # worker — can hang and ignore SIGTERM, holding GPU VRAM until the
-        # container dies and starving the next worker (HIP-OOM). Always sweep
-        # the group with SIGKILL so VRAM is released on teardown.
-        try:
-            os.killpg(pgid, signal.SIGKILL)
-        except ProcessLookupError:
-            pass
-        try:
-            self._proc.wait(timeout=10)
-        except Exception:
-            pass
+                pass
+            try:
+                self._proc.wait(timeout=10)
+            except Exception:
+                pass
+        finally:
+            if self._kv_event_proxy is not None:
+                self._kv_event_proxy.stop()
+                self._kv_event_proxy = None
 
     async def _wait_ready(self, timeout: float | None = None) -> None:
         if timeout is None: