fix(pd): route MoRIIO KV-notify to the correct decode DP rank in 2P2D - #223
raviguptaamd wants to merge 1 commit into
Conversation
…or 2P2D KV-notify (rebased onto upstream tip; upstream already has round-robin)
![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: 82dc9811af
ℹ️ 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".
| if _eff_dp > 1 && prefill_dp_rank.is_none() { | ||
| let rank = self.prefill_dp_round_robin.fetch_add(1, Ordering::Relaxed) | ||
| % self.intra_node_data_parallel_size; | ||
| % _eff_dp; | ||
| prefill_dp_rank = Some(rank); |
There was a problem hiding this comment.
When --moriio-dp-size exceeds the per-pod intra_node_data_parallel_size (the advertised 2P2D example is 16 versus 8), this produces ranks 8–15 and sends them in X-data-parallel-rank to whichever single pod was independently selected by the routing policy. The repository's DP routing treats ranks as replicas of each worker URL (0..intra_node_data_parallel_size), and the counter is not offset or constrained using the selected pod, so requests can address a rank belonging to another pod or one unavailable at that endpoint. Select the pod corresponding to the global rank and translate it to the endpoint's rank, rather than applying a global round-robin rank directly to an independently selected endpoint.
Useful? React with 👍 / 👎.
| "remote_engine_id": serde_json::Value::Null, | ||
| "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.
With MoRI-IO READ mode and moriio_dp_size set for a multi-pod deployment, this prefill leg now advertises the global DP size, but build_decode_kv_transfer_params still overwrites the returned parameters with intra_node_data_parallel_size at line 351. For the documented 16/8 topology, the two legs therefore describe different remote worlds and the decode pull can address only the local eight prefill ranks. Use the same effective size on the decode leg as well, or restrict the new setting explicitly to WRITE mode.
Useful? React with 👍 / 👎.
| intra_node_data_parallel_size: self.intra_node_data_parallel_size, | ||
| moriio_dp_size: 0, |
There was a problem hiding this comment.
When the router is launched through the supported Python Router API rather than the CLI, this hard-coded zero makes effective_dp_size() always fall back to the intra-node size. The constructor has no moriio_dp_size argument or stored field, so Python-based 2P2D deployments cannot enable the fix introduced by this commit at all. Add the option to the Python constructor and propagate it into RouterConfig instead of forcing zero.
Useful? React with 👍 / 👎.
2 participants
Single commit on top of current
main(1d10e71). Supersedes #181, which I'm closing — same fix, but that PR was opened in June against a much older base and its review comments predate a rebase that changed the whole diff. Fresh PR so there's one clean thing to review.Problem
In 2P2D disaggregated serving over the MoRIIO connector, the decode side's KV-notify is addressed using the intra-node DP size. With more than one decode pod that targets the wrong DP rank, so the transfer is never acknowledged and the request wedges until the deferred write expires:
Fix
Adds
moriio_dp_sizeand routes notify througheffective_dp_size(), which falls back to the intra-node value when the flag isn't set.remote_dp_rank_overridecarries the resolved rank across the pod boundary.--moriio-dp-sizeis opt-in: without iteffective_dp_size()returns the intra-node size and existing single-pod deployments are byte-for-byte unchanged.Validation
Built into a GLM-5.1-FP8 disaggregated image, run on MI300X (8 nodes, 2P/2D EP16, MoRI-EP wideEP):
remote blocks never arrivednorReaped deferred sendsreproducesPerf-neutral against the pre-rebase router — 60.0 ms / 440 tok/s here vs 59.5 ms / 444 tok/s before, within noise. Compiles clean (
vllm_router_rs, release).Note on the round-robin commit
The earlier version of this work carried a second commit (
11841c0d, round-robin DP-rank assignment for service discovery). That is already upstream —prefill_dp_round_robinis invllm_pd_router.rsonmaintoday — so it's dropped. Resolving the rebase meant keeping upstream's counter and driving it fromeffective_dp_size()rather than reintroducing a duplicate.Consumed by ROCm/MAD#206, which pins a fork of this repo until this merges.