feat(vllm): DP+EP and PD (prefill/decode) rollout disaggregation via the production vllm-router

[aoshen02]

Summary

Extends the multi-node rollout work (#68 / #85 / #89 lineage) to disaggregated rollout: data-parallel + expert-parallel engines, and prefill/decode (PD) split over NIXL, routed by the real vllm-router with static URLs (SkyRL-style) — not the MiniLB test stub.

Validated on 2-node GB200 (NVL72, arm64/cu13, 4 GPU/node): cross-node PD with tp4-per-engine (prefill on rack1-04, decode on rack1-05), VllmPDRouter starts from static prefill/decode URLs, NIXL handshake engages, rollout generates with 0× 502.

What's in here

• DP+EP (f4ded47): wire tp = gpus // (pp*dp), EP across engines; fix the glm4v_moe fork eager-import that crashed engine construction. Probe + Qwen3-30B-A3B DP=2+EP e2e (a2fea2f).
• PD disaggregation (281364d): prefill/decode engines over the vLLM NIXL KV connector; /sleep+/wake_up honor the tunable weight-transfer timeout instead of a hardcoded 30s (7a9f9e7); health-wait 300s→1200s for large DP/MoE startup (b366228). e2e for PD (f8ee60e, e1a87eb) and non-colocate DP=2 (3d2ea36).
• Production router (tip, 2489dfe): route PD through VllmPDRouter (vllm_pd_disaggregation) with static prefill_urls/decode_urls — engines start first, the orchestrator collects each engine's (url, bootstrap_port)/url, then launches the router with those URLs. Engines pass router_ip=None so they skip /workers self-registration. MiniLB fallback removed. The hand-rolled pd_proxy stays opt-in via SLIME_VLLM_PD_NIXL=1, and now sends kv_transfer_params inside sampling_params.extra_args (where vLLM's disagg /inference/v1/generate actually reads it).
• Hardening / observability: per-engine-TP + strict external-config check (67679bc, fix(vllm): per-engine TP + strict external config check (harden #68) #85), lifecycle reorder of vllm_engine.py (no logic change, 4c00a0c), [vllm-topo] logging (3fee27d, 66dda9d).

Scope / dependencies

• Base is feature/multi_nodes (the multi-node rollout lineage, same base as fix(vllm): unify engine control-plane HTTP timeout as --vllm-engine-request-timeout-secs (use for /sleep+/wake_up) #89), not main — this builds directly on that branch's topology work.
• The colocate cuBLAS grad-accumulation-fusion fix is PR fix(megatron): honor --[no-]gradient-accumulation-fusion on the megatron.bridge provider #101 (model_provider.py), intentionally not included here to keep that fix in its own PR.
• New tests are launch-based GPU e2e (Qwen2.5-0.5B / Qwen3-30B-A3B) plus unit tests for arg parsing and engine config; the e2e require a multi-GPU (and for PD, multi-node) runner.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces multi-node rollout engine topology support for vLLM, including Data Parallelism (DP), Expert Parallelism (EP), and Prefill/Decode (PD) disaggregation via a new NIXL relay proxy. It also resolves an eager import crash in glm4v_moe.py and adds comprehensive end-to-end tests. The review feedback highlights critical improvements: reducing the startup timeout for headless workers to prevent a 20-minute block, moving the weight-sync group guard to cover non-packed paths, defensively handling None values in the proxy's JSON payloads, and registering a cleanup handler to close the aiohttp client session to avoid resource leaks.

[gemini-code-assist]

In multi-node setups, headless workers (node_rank > 0) do not expose an HTTP health endpoint, so they call _wait_worker_process_alive(self.process). Since _wait_worker_process_alive has a default timeout_s of 1200.0 (20 minutes) and loops until the timeout is reached if the process remains alive, this will cause every headless worker to block the init() method for 20 minutes on successful startup.

Since the head node (node_rank == 0) already waits for the entire cluster to be healthy via _wait_server_healthy, the headless workers only need to perform a short sanity check to ensure their subprocesses didn't immediately crash. Passing a short timeout (e.g., 15.0 seconds) to _wait_worker_process_alive will prevent this massive startup delay.

Suggested change

	_wait_worker_process_alive(self.process)
	_wait_worker_process_alive(self.process, timeout_s=15.0)

[gemini-code-assist]

The guard checking if self._model_update_groups is None is currently only implemented in _update_weights_vllm_packed. However, the non-packed path (_update_bucket_weights_from_distributed and _update_expert_bucket_weights_from_distributed) also runs on the PP source rank and calls update_weights_from_distributed using self._model_update_groups. If self._model_update_groups is None there, it will pass None to trainer_send_weights, leading to an opaque hang or crash.

To protect both the packed and non-packed paths, consider moving this guard to the very beginning of _sync_weights_to_rollout_engines so it executes once for the PP source rank before any weight synchronization begins.

[gemini-code-assist]

Using .setdefault() and direct subscripting on sampling_params and extra_args can raise a TypeError if the input JSON explicitly sets either of these fields to None (which is common in client payloads). Ensuring these fields are valid dictionaries before mutating them makes the proxy much more robust.

        pf_body = copy.deepcopy(body)
        sampling_params = pf_body.get("sampling_params")
        if not isinstance(sampling_params, dict):
            sampling_params = {}
            pf_body["sampling_params"] = sampling_params
        sampling_params["max_tokens"] = 1
        extra_args = sampling_params.get("extra_args")
        if not isinstance(extra_args, dict):
            extra_args = {}
            sampling_params["extra_args"] = extra_args
        extra_args["kv_transfer_params"] = {"do_remote_decode": True}

[gemini-code-assist]

Similar to the prefill step, defensively ensure that sampling_params and extra_args are valid dictionaries before setting kv_transfer_params to avoid potential TypeError exceptions when processing client payloads.

Suggested change

	dc_body = copy.deepcopy(body)
	dc_body.setdefault("sampling_params", {}).setdefault("extra_args", {})["kv_transfer_params"] = {
	"do_remote_prefill": True,
	**ktp,
	}
	dc_body = copy.deepcopy(body)
	dc_sampling_params = dc_body.get("sampling_params")
	if not isinstance(dc_sampling_params, dict):
	dc_sampling_params = {}
	dc_body["sampling_params"] = dc_sampling_params
	dc_extra_args = dc_sampling_params.get("extra_args")
	if not isinstance(dc_extra_args, dict):
	dc_extra_args = {}
	dc_sampling_params["extra_args"] = dc_extra_args
	dc_extra_args["kv_transfer_params"] = {
	"do_remote_prefill": True,
	**ktp,
	}

[gemini-code-assist]

The aiohttp.ClientSession created in _ensure_session is never closed, which can lead to resource leaks and unclosed connector warnings when the application shuts down. Registering an on_cleanup signal handler to close the session ensures proper resource cleanup.

Suggested change

	def build_app(self) -> web.Application:
	app = web.Application(client_max_size=512 * 1024 * 1024)
	app.router.add_post(GENERATE_PATH, self.handle_generate)
	app.router.add_get("/health", self.handle_health)
	app.router.add_post("/workers", self.add_worker)
	app.router.add_get("/workers", self.list_workers)
	app.router.add_delete("/workers/{url}", self.remove_worker)
	for ctl in ("/abort_requests", "/resume", "/pause"):
	app.router.add_post(ctl, self.handle_broadcast)
	return app
	async def cleanup(self, app: web.Application) -> None:
	if self._session is not None and not self._session.closed:
	await self._session.close()
	def build_app(self) -> web.Application:
	app = web.Application(client_max_size=512 * 1024 * 1024)
	app.router.add_post(GENERATE_PATH, self.handle_generate)
	app.router.add_get("/health", self.handle_health)
	app.router.add_post("/workers", self.add_worker)
	app.router.add_get("/workers", self.list_workers)
	app.router.add_delete("/workers/{url}", self.remove_worker)
	for ctl in ("/abort_requests", "/resume", "/pause"):
	app.router.add_post(ctl, self.handle_broadcast)
	app.on_cleanup.append(self.cleanup)
	return app

[CalvinXKY]

When you commit a PR, please refer the CONTRIBUTING.md to modify your PR format.

[princepride]

Hi, how is the pr going, do you need my help? 😊