fix(pd): route MoRIIO KV-notify to the correct decode DP rank in 2P2D - #181
raviguptaamd wants to merge 1 commit into
Conversation
…uting Captures the uncommitted working-tree edits (from /tmp/router-inspect, the tree that built the validated vllm-router-2p2d-dpfix binary on 2026-06-14) on top of PR vllm-project#181 (11841c0). Without these, the decode->prefill notify targets a different DP rank than the prefill was pinned to -> "remote blocks never arrived" deferred write expiry under 2P2D EP=16 load (reproduced: clean vllm-project#181 wedges, this serves). Adds remote_dp_rank_override flag + remote_dp_size so the MoRIIO connector's rank arithmetic (blake2s(rid)%dp_size) matches the topology on both legs.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6409ac1409
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| "remote_block_ids": serde_json::Value::Null, | ||
| "transfer_id": transfer_id.unwrap_or(""), | ||
| "remote_dp_size": self.intra_node_data_parallel_size, | ||
| "remote_dp_size": self.effective_dp_size(), |
There was a problem hiding this comment.
When --moriio-dp-size is set and MoRI-IO registers READ mode, this sends the prefill leg remote_dp_size as the effective cross-pod size, but the READ decode path still overwrites remote_dp_size with self.intra_node_data_parallel_size in build_decode_kv_transfer_params. In a 2P2D setup such as moriio_dp_size=16 and intra_node_data_parallel_size=8, the two legs of the same transfer disagree about the DP world size, so the decode handshake only targets the local-size world. Use effective_dp_size() in the READ decode path as well.
Useful? React with 👍 / 👎.
| worker_startup_timeout_secs: self.worker_startup_timeout_secs, | ||
| worker_startup_check_interval_secs: self.worker_startup_check_interval, | ||
| intra_node_data_parallel_size: self.intra_node_data_parallel_size, | ||
| moriio_dp_size: 0, |
There was a problem hiding this comment.
For routers started through the PyPI/Python entrypoint, Router.start() builds this RouterConfig, but the new value is hardcoded to 0 and the PyO3 constructor/RouterArgs do not expose a matching field. That means installed vllm-router users cannot configure the new cross-pod MoRI-IO DP size; Wide-EP deployments silently fall back to intra_node_data_parallel_size, so the new round-robin rank assignment and remote_dp_size stamping still use the wrong world size. Thread this option through the Python args/constructor instead of forcing 0.
Useful? React with 👍 / 👎.
…or 2P2D KV-notify (rebased onto upstream tip; upstream already has round-robin)
6409ac1 to
82dc981
Compare
raviguptaamd
changed the title
|
Force-pushed this branch to rebase it — worth flagging what changed, since the diff will look different to anyone who reviewed it earlier. Before: 2 commits, based on June The dropped commit is the round-robin DP-rank fix ( The previous head Nothing about the remaining fix changed — only its base. |
|
Superseded by #223 — closing this. Same fix, but this PR was opened 2026-06-02 against a much older The round-robin commit that used to ride along here ( Old head |
1 participant
Rebased onto current
main(1d10e71) and reduced to a single commit. This PR was previously 2 commits and conflicted; it is now 1 commit, no conflicts.What changed about this PR
The round-robin DP-rank commit that used to ride along here (
11841c0d) is already upstream —prefill_dp_round_robinis invllm_pd_router.rsonmaintoday. Keeping it was the main source of the conflicts, so it's dropped. What remains is the one fix that is still missing upstream.The fix
In 2P2D disaggregated serving over the MoRIIO connector, the decode side's KV-notify is addressed using the intra-node DP size, so with more than one decode pod the notify targets the wrong DP rank. The transfer is then never acknowledged and the request wedges until the deferred write expires:
This adds
moriio_dp_sizeand routes notify througheffective_dp_size(), falling back to the intra-node value when the flag isn't set — so single-pod behaviour is byte-for-byte unchanged.remote_dp_rank_overridecarries the resolved rank across the pod boundary.--moriio-dp-sizeis opt-in. Without it,effective_dp_size()returns the intra-node size and nothing about existing deployments changes.Validation
Built into a GLM-5.1-FP8 disaggregated image and run on MI300X (8 nodes, 2P/2D EP16, MoRI-EP wideEP):
remote blocks never arrived, noReaped deferred sends— the wedge this fixes does not reproduceAlso perf-neutral against the pre-rebase router: 60.0 ms / 440 tok/s here vs 59.5 ms / 444 tok/s on the old branch, i.e. within noise. The router compiles clean (
vllm_router_rs, release).Consumed by ROCm/MAD#206, which currently has to pin a fork of this repo until this merges.