Skip to content

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

Draft
cquil11 wants to merge 2 commits into
AMD-AGI:mainfrom
cquil11:fix/atom-dpa-kv-events
Draft

fix(atom): relay KV events from data-parallel ranks#121
cquil11 wants to merge 2 commits into
AMD-AGI:mainfrom
cquil11:fix/atom-dpa-kv-events

Conversation

@cquil11

@cquil11 cquil11 commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Summary\n\n- relay KV events from every ATOM data-parallel EngineCore through one worker-local XSUB/XPUB proxy\n- preserve the one logical KV-event endpoint advertised to Infera for each ATOM worker\n- keep direct-bind behavior available for existing standalone hook callers\n- stop the relay with the worker and retain the existing ATOM process-group cleanup behavior\n\n## Root cause\n\nAn ATOM TP+DP worker creates one BlockManager in every data-parallel EngineCore process. The current hook gives every process the same externally advertised ZMQ bind endpoint. All DP ranks therefore race to bind one port; one wins and the remaining ranks fail during startup with zmq.error.ZMQError: Address already in use.\n\nThis reproduced simultaneously on all three workers of the MI355X DeepSeek-V4 2P/1D DPA validation in InferenceX run 31995111191, Slurm 39443. Every rank completed model load and CUDA-graph capture before failing at BlockManager.init -> publisher.ensure_bound() on the shared endpoint. The Infera health registry consequently remained at 0/2 prefill and 0/1 decode.\n\nThe fix binds one proxy endpoint in the Infera worker process and gives every spawned EngineCore a local IPC ingress to connect its PUB socket. ZMQ forwards all rank streams through the single advertised endpoint, preserving full DP cache visibility rather than publishing only rank zero or silently disabling KV-aware routing.\n\n## Validation\n\n- pytest -q tests/unit/engine tests/unit/router/test_kv_event_client.py: 50 passed\n- focused relay regression uses two independent EngineCore publishers and observes both streams through one advertised endpoint\n- worker wiring test verifies the child environment uses connect mode and the proxy is stopped during teardown\n- Ruff format/check: passed\n- applicable pre-commit hooks: passed\n- git diff --check: passed\n\nHardware revalidation will update the companion InferenceX PR after the current exact-head attempt finishes naturally. No running job was cancelled.

Signed-off-by: Cam Quilici <cjquilici@gmail.com>
Signed-off-by: Cam Quilici <cjquilici@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant