Docs: how to set up and run single-node 4P4D (prefill/decode disaggregation) with Mooncake (XGMI/TCP)

[carlushuang]

Update (see the comments below + PR #1495): the original recipe here uses protocol=tcp, but TCP is not the recommended single-node transport — it serializes under concurrency (the "TTFT blow-up" in the perf table below). The correct transport is protocol=hip (XGMI, GPU-direct IPC), which restores scaling (~4× throughput / ~7× TTFT at c4, matching aggregated TP=4). Also, the c16 "hang" was not a transport issue — it's a prefill-engine OOM from a per-rank KV-budget imbalance (fix in PR #1495). protocol=tcp still works out-of-the-box (the stock image's mooncake is Tcp+Rdma only), so the recipe below is valid as the slow fallback; use protocol=hip with a -DUSE_HIP=ON mooncake build for real performance. Full analysis in the comment thread.

---

Summary

atom/kv_transfer/disaggregation/README.md documents the MORI-IO RDMA path for prefill/decode (P/D) disaggregation, but PR #1308 ("Support PD disaggregation on Single node") also enables a single-node setup over the Mooncake connector that works without an RDMA NIC (protocol=tcp). That path is currently undocumented. This issue requests a short "single-node 4P4D" how-to in the README, and records a validated recipe + the gotchas hit while bringing it up.

Validated on one 8×MI355X (gfx950) node, ATOM main, model deepseek-ai/DeepSeek-V4-Pro, split into 4-GPU prefill + 4-GPU decode. Requests returned correct output with kv_transfer_params.do_remote_prefill=true, decode TPOT matching an aggregated TP=8 server.

Proposed how-to (validated recipe)

All three processes run on the same node. The prefill and decode engines each use TP=4; requests go to the proxy.

The recipe below uses protocol=hip (XGMI) — the recommended transport — which needs both engines to see all 8 GPUs in a rotated order so each engine's own 4 GPUs are cuda:0-3 while the peer's are visible for IPC. For the TCP fallback (works on the stock Tcp+Rdma-only mooncake build, but serializes under concurrency), set "protocol":"tcp" and use disjoint visibility instead (HIP_VISIBLE_DEVICES=0,1,2,3 / 4,5,6,7). --max-num-seqs 128 avoids the prefill-engine OOM described in the comments / PR #1495 (until that fix lands).

1. Proxy

python -m atom.kv_transfer.disaggregation.proxy --port 10001

2. Prefill (kv_producer, own GPUs = physical 0-3)

HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
TORCHINDUCTOR_CACHE_DIR=/tmp/ind_p TRITON_CACHE_DIR=/tmp/tri_p HOME=/tmp/home_p \
python -m atom.entrypoints.openai_server \
  --model deepseek-ai/DeepSeek-V4-Pro \
  --kv_cache_dtype bf16 -tp 4 --gpu-memory-utilization 0.85 --max-num-seqs 128 \
  --host 0.0.0.0 --server-port 8003 \
  --kv-transfer-config '{"kv_connector":"mooncake","kv_role":"kv_producer","protocol":"hip","proxy_ip":"127.0.0.1","proxy_ping_port":36367,"http_port":8003,"handshake_port":6301}'

3. Decode (kv_consumer, own GPUs = physical 4-7, rotated to cuda:0-3)

HIP_VISIBLE_DEVICES=4,5,6,7,0,1,2,3 \
TORCHINDUCTOR_CACHE_DIR=/tmp/ind_d TRITON_CACHE_DIR=/tmp/tri_d HOME=/tmp/home_d \
python -m atom.entrypoints.openai_server \
  --model deepseek-ai/DeepSeek-V4-Pro \
  --kv_cache_dtype bf16 -tp 4 --gpu-memory-utilization 0.85 --max-num-seqs 128 \
  --host 0.0.0.0 --server-port 8004 \
  --kv-transfer-config '{"kv_connector":"mooncake","kv_role":"kv_consumer","protocol":"hip","proxy_ip":"127.0.0.1","proxy_ping_port":36367,"http_port":8004,"handshake_port":7401}'

4. Send requests to the proxy

curl -s http://127.0.0.1:10001/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"deepseek-ai/DeepSeek-V4-Pro","prompt":"The capital of France is","max_tokens":16,"temperature":0}'

The response's kv_transfer_params.do_remote_prefill should be true, confirming the KV was computed on the prefill engine and transferred to decode.

Setup notes / gotchas worth documenting

These were not obvious from the README and each cost time:

• Connector selection: kv_connector defaults to moriio (RDMA). Single-node requires "kv_connector":"mooncake". Prefer "protocol":"hip" (XGMI GPU-direct IPC; needs a -DUSE_HIP=ON -DUSE_MNNVL=ON mooncake build + the all-8-GPU rotated visibility above). "protocol":"tcp" is the no-RDMA-NIC fallback that works on the stock build but serializes under concurrency (default protocol is rdma). TransferEngine.initialize(..., "tcp"/"hip", "") returns 0 and ignores the auto-selected rdma{gpu} device.
• Bind --host 0.0.0.0: each engine reports its real IP (get_ip()) to the proxy. With --host 127.0.0.1 the proxy cannot reach the engine and returns 500.
• Distinct handshake_port per engine on the same node (e.g. 6301 / 7401); otherwise the intra-node ZMQ side channels collide.
• Per-engine compile-cache isolation: two same-model engines on one node race on the inductor/triton cache (torch.compile "Guard failed" crash) and on ~/.cache/atom. Set distinct TORCHINDUCTOR_CACHE_DIR / TRITON_CACHE_DIR (and separate HOME) per engine.
• cudagraph works for single-node PD over TCP (no --enforce-eager needed; capture ~7s per engine).
• Restart the proxy after (re)launching engines: a long-lived proxy retains stale registrations for engines that were restarted and then returns empty responses for new requests. Restarting the proxy alone makes both engines re-register via the ping channel (no model reload). A stale-registration eviction / health check on the proxy would remove this footgun.
• DeepSeek-V4 specifics: the Mooncake connector already handles the V4 pools (register_v4_kv_caches, block + slot regions) — registration logs 154 chunks (91 block regions, 62 slot regions), i.e. classical KV + SWA per-request slots + indexer cache are all transferred.

An "8-GPUs-visible reorder" variant (prefill HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7, decode 4,5,6,7,0,1,2,3) also works over TCP and is the prerequisite for an intra-node XGMI/IPC transport; disjoint visibility (above) is the simpler default for TCP.

Observed behavior (for discussion)

Single-node 4P4D over TCP, deepseek-ai/DeepSeek-V4-Pro, ISL/OSL = 1024/1024, bf16 KV, cudagraph. Output throughput / median TTFT / median TPOT vs an aggregated TP=8 server on the same node:

concurrency	4P4D PD (TCP)	Aggregated TP=8
1	70 tok/s / 0.23 s / 13.9 ms	76 tok/s / 139 ms / 13.0 ms
4	91 tok/s / 25.0 s / 13.9 ms	270 tok/s / 127 ms / 14.2 ms
16	140 tok/s / 56.0 s / 14.8 ms	902 tok/s / 143 ms / 17.3 ms

Decode TPOT tracks the aggregated server (~14 ms), but median TTFT grows steeply with concurrency (0.23 s → 25 s → 56 s) and end-to-end throughput stays low. It looks like requests serialize waiting on the single prefill engine + per-request proxy/transfer coordination rather than on prefill compute (a 1k-token prefill is ~150 ms). Is this expected for single-node TCP PD, or is there a recommended configuration (e.g. multiple prefill instances, DP-attention, batching of transfers) to avoid the TTFT blow-up? Documenting the intended operating point would help.

Ask

1. Add a "single-node (no RDMA NIC) via Mooncake TCP" section to atom/kv_transfer/disaggregation/README.md with the recipe above.
2. Confirm the intended concurrency/operating point for single-node TCP PD (the TTFT scaling above).
3. Consider a proxy stale-registration health check so a restarted engine doesn't require a manual proxy restart.

Happy to open a docs PR for item 1 if useful.

[carlushuang]

Update: root-caused the concurrency scaling + the "hang", with fixes

Following up on the two open questions from the original post (the TTFT blow-up with concurrency, and the intended single-node operating point). After more digging, both come down to (a) the KV transport and (b) a memory-sizing issue on the prefill engine — not the proxy. Summary below; all measured on 8×MI355X (gfx950), deepseek-ai/DeepSeek-V4-Pro, TP=4 prefill + TP=4 decode, ISL/OSL=1024/1024, bf16 KV.

1. The TCP TTFT blow-up is the transport, and protocol=hip (XGMI) fixes it

Isolating the prefill path (OSL=16 so decode is negligible), req/s and median TTFT:

concurrency	mooncake tcp	mooncake hip (XGMI)	aggregated TP=4 (ref)
1	2.66 / 207 ms	2.53 / 189 ms	2.76 / 137 ms
4	1.78 / 2797 ms	7.13 / 379 ms	6.66 / 254 ms
8	—	10.56 / 554 ms	—
16	0.82 / 19778 ms	(see §2)	6.90 / 548 ms

protocol=tcp collapses with concurrency (host-staged, per-request serialization). Switching the mooncake connector to protocol=hip (XGMI GPU-direct IPC) restores scaling — ~4× throughput / ~7× TTFT at c4, matching the aggregated TP=4 baseline, and it scales cleanly through c8 (10.6 req/s). cudagraph works with hip (no custom-all-reduce conflict).

Caveat: protocol=hip requires a mooncake built with -DUSE_HIP=ON -DUSE_MNNVL=ON and the HipTransport device-restore fix (kvcache-ai/Mooncake #2566, merged). The stock image's engine.so is Tcp+Rdma only (strings shows Protocol 'nvlink' requires -DUSE_MNNVL=ON), so on this box tcp is the only out-of-the-box option. Suggest documenting the -DUSE_HIP build + protocol=hip for single-node no-RDMA-NIC setups.

2. The "c16 hang" is a producer OOM crash, not a transport hang

With XGMI, c16 still "hangs" — but it's not a hang. The prefill engine OOM-crashes, and because the proxy has no timeout/health-check for a dead prefill engine, the client blocks forever (and /health on the frontend still returns 200 after the ModelRunner workers die). GPUs go idle (0%), which is the tell — it's a host-side wait on a dead engine, not a stuck kernel.

First fatal in the prefill log:

torch.OutOfMemoryError: HIP out of memory. Tried to allocate 1.19 GiB.
GPU 0 has ... 888.00 MiB free ...   (in model_runner.py:2166 forward → deepseek_v4 forward)

At c16 the scheduler batches all 16 prefills into one 16384-token forward; its ~1.19 GiB activation doesn't fit because GPU 0 is pinned at ~99% at startup. c1–c8 (≤8192-token batches) fit in the sliver, which is exactly why it "works up to c8 then hangs".

Two suggestions here regardless of the memory issue below:

• The proxy should time out / health-check the prefill instance so a dead engine surfaces as a 5xx instead of an indefinite client hang.
• Consider a startup log if a rank exceeds the util target by a large margin (see §3).

3. Root cause of the OOM: per-rank memory imbalance + KV autosizing + dummy warmup

Three compounding factors, all on the generic model-runner path (so not DSV4-specific):

(a) get_num_blocks doesn't reconcile per-rank memory. Each rank computes num_kvcache_blocks independently from available_for_kv = util*total - peak_torch - cudagraph - safety; the budget term binds, so every rank picks the same block count sized for the lightest rank. The heavy ranks then overshoot the util target. Post-init actual usage (util 0.85, TP=4): 87.2%, 87.3%, 97.5%, 99.0%. peak_torch is identical across ranks (211 GB) — weights are evenly sharded — so this is not a weight imbalance.

(b) ~34 GB of non-torch memory lands on the 2 lowest ranks. Verified via torch.cuda.mem_get_info: rank0/1 committed 257/251 GB vs 223 GB on the others. It is not weights (peak_torch uniform) and not RCCL (summing NCCL_DEBUG=INFO allocations ≈ 10 GB/rank, uniform). It appears to be runtime/allocator memory concentrated on the coordinator ranks; the exact call-site still needs a torch.cuda.memory._record_memory_history snapshot to pin. This imbalance is ~constant regardless of TP degree, which is why it's invisible at aggregated TP=8 (heavy rank only reaches ~89% with ~31 GB free) but fatal at TP=4/PD (heavy rank at 99% with ~0.9 GB free — half the sharding, plus PD forces util up to fit the 12.8 GB per-request SWA cache).

(c) Warmup under-reserves activation. warmup_model runs a full max_num_batched_tokens forward but with is_dummy_run=True, and DSV4 attention short-circuits the sparse-attention dispatch on dummy runs. So peak_torch misses the real prefill activation, and the KV pool is sized too large.

4. Fixes / workarounds

Validated workaround (no code change) — got c16 working at 13 req/s:

--attn-prefill-chunk-size 8192 --max-num-batched-tokens 8192   # caps per-step prefill activation → OOM-free at any concurrency
--max-num-seqs 128                                             # shrinks the 12.8GB per-request SWA cache

The activation OOM is governed by per-step prefill tokens T_eff = min(max_num_batched_tokens, attn_prefill_chunk_size, concurrency*ISL); capping the chunk decouples OOM from concurrency.

Suggested engine-side fixes:

• get_num_blocks: reconcile available_for_kv across ranks with an all_reduce(MIN) (size the KV pool for the tightest rank) so heavy ranks keep real headroom.
• Warm up V4 without is_dummy_run (or add an explicit sparse-attention activation reserve) so peak_torch reflects the true prefill peak.
• Proxy: dead-prefill-engine timeout (see §2).

Happy to open PRs for the docs (§1) and/or the get_num_blocks reconciliation (§3a) if useful. I can also post the memory._record_memory_history snapshot once I pin the exact rank-0/1 allocation.