feat(dynamo): support prefill decode disaggregation

[jthomson04]

Summary

• add opt-in managed Dynamo prefill/decode worker counts
• launch decode and prefill roles with NIXL KV transfer, prefill KV events, and stable refit ordering
• add a three-GPU GRPO L1 test that reuses the existing Dynamo functional harness

Aggregate mode remains unchanged when dynamo_cfg.disaggregation is null. This PR keeps the current Dynamo 1.3, vLLM 0.23, and NIXL 1.1 runtime pins.

Scope

This PR validates functional correctness. It makes no performance or convergence claim.

Tests

• 59 passed in the affected Dynamo configuration and managed-runtime unit tests
• Ruff check and format check passed for changed Python files
• shell syntax checks passed for the Dynamo functional scripts and ray.sub

The full Dynamo L1 runs the existing two-GPU aggregate test and the new three-GPU prefill/decode test.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jthomson04]

Reviewed by a team of six agents. The source changes came through in good shape — I verified the load-bearing upstream assumptions against pinned ai-dynamo 1.3.0 and vLLM 0.23.0 source, and they all hold: the prefill/decode flag spellings, the decode→backend component mapping that startup gating depends on, prefill workers serving the rl endpoint under --enable-rl, the KVEventsConfig field names, the tcp://*:PORT bind form, the kv_both role choice, and the one-port-per-engine NIXL scheme (correct at any TP — vLLM 0.23.0 uses + data_parallel_index and multiplexes TP ranks over a single socket). The three-GPU config arithmetic, actor-name uniqueness, and the nixl/nixl-cu13 assertions all check out too. Ruff and bash -n are clean.

The findings cluster in two places: test coverage and evidence.

Performance and convergence evidence

This is the main ask. P/D disaggregation is a throughput/latency feature, and the new three-GPU L1 smoke can't demonstrate a win — one TP1 prefill plus one TP1 decode at max_new_tokens: 128 is a smaller configuration than the aggregate baseline it's compared against. It's a good plumbing test, just not evidence of benefit. Right now there's no basis for choosing prefill_workers/decode_workers, and nothing showing NIXL transfer plus the wider refit fan-out costs less than the pipelining gains.

Everything below already exists in the repo:

• Throughput on the existing grpo-nanov3-30ba3b-3n8g-megatron-dynamo-swe1.yaml, holding everything constant except the inference-node split: current aggregate (2× TP4/EP4) vs. two P/D splits on the same 8 GPUs — e.g. decode:1/prefill:1 at TP4 and decode:1/prefill:3 at TP2. Report timing/train/total_step_time and timing/train/valid_tokens_per_sec_per_gpu.
• Latency: dynamo_work_handler_time_to_first_response (TTFT) and vllm:inter_token_latency (ITL).
• Refit cost: refit/* timings — every prefill engine adds an NCCL rank and a cache flush per step.
• Convergence: a short matched run (same seed/model/batch) aggregate vs. P/D on train/token_mult_prob_error, which the repo already treats as red above 1.05. Remote prefill over NIXL should be numerically identical, so this is cheap and is the proof that it is.

If a 3-node sweep is too expensive right now, a single-node 8-GPU run at TP1 (decode:4/prefill:4 vs. 8 aggregate engines) reporting valid_tokens_per_sec_per_gpu and token_mult_prob_error would establish the shape of the tradeoff and let the guide give a concrete starting point instead of just the formula.

One note on feasibility: per-role TTFT/ITL isn't obtainable today. The PR publishes role in worker metadata at dynamo_worker.py:295, but DynamoWorkerEndpoint.from_metadata keeps only instance_id/system_url and the metrics sampler keys by list ordinal, so prefill and decode — which have opposite KV-residency profiles — get averaged into one series. That plumbing lives in metrics.py/refit.py, outside this diff, so it may be a prerequisite for the numbers above or a reasonable follow-up. Your call.

Test coverage

The headline: test_frontend_requires_matching_prefill_and_decode_rl_membership passes with the multi-component readiness gate deleted outright, and again with its identity check deleted. Both reproduced independently, on code the test provably executes. All nine raise statements this PR adds are also unreachable from the suite. Details inline.

---

Two hypotheses were investigated and refuted, noted so they don't resurface: there is no TP>1 NIXL side-channel port collision, and no flush_cache/NIXL-lease race (vLLM's has_finished_requests() keeps the engine non-idle while a lease is held, so the reset callback never fires).

Generated by Claude Code

[jthomson04]

test_dynamo_managed_runtime.py:932

This test feeds a single happy-path payload and then asserts nothing, so neither claim in its name is actually pinned. Two mutations of _wait_for_frontend leave the suite fully green (78 passed, 0 failed), reproduced independently twice:

• Restricting the gate to if component == "backend" — i.e. deleting the entire multi-component behavior this PR adds — passes.
• Deleting observed[component]["generate"] == observed[component]["rl"], the literal "matching ... rl membership" — passes.

A reachability probe confirms this test does execute the gate on every run, so it's not a "code never runs" artifact. The only thing currently pinned is that the happy path doesn't hang, and that's incidental — the 2-element responses iterator would raise StopIteration on a second poll.

Two negative tests close it. The load-bearing assertion in each is that /v1/models is never reached, since readiness is the only thing that unlocks it. The second is the only shape that catches the identity check (counts match, instance-ID sets diverge):

def test_frontend_wait_times_out_while_the_prefill_component_is_missing(monkeypatch):
    health = {"instances": [_instance("backend", "decode-0", e) for e in ("generate", "rl")]}
    runtime, urls = _disaggregated_runtime(monkeypatch, health, iter([0.0, 1.0, 2.0, 100.0]))
    with pytest.raises(RuntimeError) as excinfo:
        runtime._wait_for_frontend(expected_components={"backend": 1, "prefill": 1})
    assert urls == ["http://127.0.0.1:3000/health"] * 2
    assert "'prefill': {'generate': 0, 'rl': 0}" in str(excinfo.value)


def test_frontend_wait_rejects_a_component_whose_generate_and_rl_ids_diverge(monkeypatch):
    health = {"instances": [
        _instance("backend", "decode-0", "generate"), _instance("backend", "decode-0", "rl"),
        _instance("prefill", "prefill-0", "generate"), _instance("prefill", "prefill-1", "rl"),
    ]}
    runtime, urls = _disaggregated_runtime(monkeypatch, health, iter([0.0, 1.0, 2.0, 100.0]))
    with pytest.raises(RuntimeError):
        runtime._wait_for_frontend(expected_components={"backend": 1, "prefill": 1})
    assert urls == ["http://127.0.0.1:3000/health"] * 2

Both were run against this branch and pass on clean source while failing under the respective mutation. Renaming the existing one to something like test_frontend_accepts_a_fully_registered_disaggregated_fleet would also make its scope honest:

Suggested change

	def test_frontend_requires_matching_prefill_and_decode_rl_membership(
	def test_frontend_accepts_a_fully_registered_disaggregated_fleet(

[jthomson04]

Addressed in 81d855f. The happy-path test is renamed, and the focused negative cases cover a missing prefill component and divergent generate/RL instance IDs. Both cases prove that the models endpoint is not reached.

[jthomson04]

arguments.py:29

Nothing currently stops vllm_kwargs: {enable_prefix_caching: false}, and if a user sets it the managed --kv-events-config this PR carefully builds is silently discarded. Dynamo's create_kv_events_config returns None on that flag before it ever looks at the caller's config:

# If prefix caching is not enabled, no events config needed
if not engine_config.enable_prefix_caching:
    logger.info("No kv_events_config required: prefix caching is disabled")
    return None

and args.py:262 then writes that None back over the parsed value. Net effect: the reserved KV-event port is never bound and the router loses prefill-side KV awareness, with only an INFO line. Note the early return sits below the decode short-circuit, so this also applies to aggregated mode under router_mode: kv, which is what the shipped example config uses.

Adding it to the reject-list catches both spellings, since _normalise_flag folds --no- and the vllm_kwargs loop checks the normalised form:

Suggested change

	"--disaggregation-mode",
	"--disaggregation-mode",
	"--enable-prefix-caching",

Worth deciding whether to scope this to P/D only or block it outright — either is defensible, but silently dropping the events config seems worth preventing.

[jthomson04]

Keeping prefix caching optional. When enable_prefix_caching: false, there are no cached prefixes to reuse, so Dynamo intentionally disables KV events. The updated docs state this behavior; the configuration remains valid.

[jthomson04]

config.py:315-330

This guard is nested three levels deep and every level is a silent skip, so the check quietly does nothing for the most common config shape. num_nodes: null is both the shipped default in grpo_math_1B.yaml and explicitly legal — grpo.py asserts inference_nodes is None or inference_nodes == 1 and then coerces it to 1.

Verified by replaying the validator: num_nodes: null, gpus_per_node: null, and an absent resources key all pass validation with a wrong GPU count. So disaggregation: {prefill_workers: 2, decode_workers: 2} with resources: {gpus_per_node: 2, num_nodes: null} validates clean, then fails much later in ManagedDynamoRuntime.__init__ with allocated=2, expected=4 — after Ray init, both clusters, and placement-group creation.

Mirroring grpo.py's coercion and failing loudly on the genuinely-missing cases keeps it consistent with the error-handling guideline. The shipped grpo_math_1B_dynamo_disagg.yaml sets both keys explicitly, so this doesn't break it:

Suggested change

	if self.dynamo_cfg.disaggregation is not None:
	resources = (
	colocated.get("resources") if isinstance(colocated, dict) else None
	)
	if isinstance(resources, dict):
	gpus_per_node = resources.get("gpus_per_node")
	num_nodes = resources.get("num_nodes")
	if gpus_per_node is not None and num_nodes is not None:
	configured_gpus = int(gpus_per_node) * int(num_nodes)
	expected_gpus = self.managed_inference_world_size
	if configured_gpus != expected_gpus:
	raise ValueError(
	"Dynamo disaggregation requires inference GPUs to equal "
	"(prefill_workers + decode_workers) * TP * PP: "
	f"configured={configured_gpus}, expected={expected_gpus}"
	)
	if self.dynamo_cfg.disaggregation is not None:
	resources = (
	colocated.get("resources") if isinstance(colocated, dict) else None
	)
	if not isinstance(resources, dict):
	raise ValueError(
	"Dynamo disaggregation requires "
	"policy.generation.colocated.resources to be set."
	)
	gpus_per_node = resources.get("gpus_per_node")
	if gpus_per_node is None:
	raise ValueError(
	"Dynamo disaggregation requires an explicit "
	"policy.generation.colocated.resources.gpus_per_node."
	)
	# num_nodes may be null on the single-node non-colocated path;
	# grpo.py coerces it to 1 there, so mirror that instead of skipping.
	num_nodes = resources.get("num_nodes")
	configured_gpus = int(gpus_per_node) * int(
	1 if num_nodes is None else num_nodes
	)
	expected_gpus = self.managed_inference_world_size
	if configured_gpus != expected_gpus:
	raise ValueError(
	"Dynamo disaggregation requires inference GPUs to equal "
	"(prefill_workers + decode_workers) * TP * PP: "
	f"configured={configured_gpus}, expected={expected_gpus}"
	)

[jthomson04]

Addressed in 81d855f. P/D now requires inference resources and an explicit gpus_per_node, treats num_nodes: null as one, and rejects an incorrect total during configuration validation.

[jthomson04]

worker_pool.py:183

This line is where node-local node_slot meets fleet-global group_index (via worker_role), and that distinction is currently invisible to the suite. Both existing pool tests use one placement group with one node_ip, so node_slot and group_index are always numerically equal. Setting node_slot = group_index leaves all 24 tests green, including every test this PR adds — verified independently, and a reachability probe confirms the line is executed by two existing tests, so this is a missing assertion rather than dead code. On a real two-node fleet that mutation walks NIXL/KV ports out of their bands.

A two-node case pins it (fails under the mutation, passes clean):

def test_fixed_pool_scopes_ports_per_node_while_roles_stay_fleet_global(monkeypatch):
    launch_kwargs = _run_pool(
        monkeypatch, node_ips=["10.0.0.1", "10.0.0.2"],
        worker_roles=["decode", "prefill"], bundle_counts=[1, 1],
    )
    # Roles are indexed across the whole fleet, ports only within one node.
    assert [k["worker_role"] for k in launch_kwargs] == ["decode", "prefill"]
    assert [k["nixl_port"] for k in launch_kwargs] == [4100, 4100]
    assert [k["kv_event_port"] for k in launch_kwargs] == [None, 4200]

Relatedly, all nine raise statements this PR adds are unreachable from the suite (line-coverage measured): the three role/port ValueErrors in build_dynamo_vllm_argv, both band-overflow raises at worker_pool.py:55 and :66, the three role-count guards at :93/:177/:233, and the new divisibility guard at managed_runtime.py:86 (the existing multinode test uses tp=8, which trips the earlier per-node check first; tp=3 reaches it). Happy to share the full set of tests — they share one _run_pool harness and were all verified passing.

[jthomson04]

Keeping this PR narrow. The current code keeps node-local port slots separate from fleet-global role order, and this thread does not identify a defect in that implementation. This PR does not add a multi-node pool harness or exhaustive tests for defensive raises.

[jthomson04]

grpo_dynamo_disagg.sh:25

Because exec replaces the process, $0 is rewritten to grpo_dynamo.sh — so the DYNAMO_EXP_NAME export above is the only thing keeping EXP_NAME=${DYNAMO_EXP_NAME:-$(basename "$0" .sh)} from evaluating to grpo_dynamo. If it were ever dropped or renamed, rm -rf "${EXP_DIR}" would delete the aggregate run's results directory. It's correct today, but the safety rests on an undocumented interaction:

Suggested change

	exec bash "${script_dir}/grpo_dynamo.sh"
	# DYNAMO_EXP_NAME must be exported: `exec` rewrites $0 to grpo_dynamo.sh, so the
	# EXP_NAME fallback would otherwise collide with the aggregate run's EXP_DIR,
	# which grpo_dynamo.sh removes with `rm -rf`.
	exec bash "${script_dir}/grpo_dynamo.sh"

Two smaller things in the same area, neither blocking:

• The readiness grep grpo_dynamo.sh:57 matches {'backend': 1, 'prefill': 1}, which depends on dict insertion order fixed by two adjacent lines at managed_runtime.py:113-116. Swapping them is a semantically null edit that reddens an L1 test with an unhelpful message — grepping the two keys independently, or sorted() in the print, removes the trap. (The argv greps are fine; list-of-str repr is stable, and all four needles were confirmed to match real output.)
• grpo_dynamo.sh:30 falls back to /opt/dynamo_venv, while both venv.py:26-27 and install.sh:21 fall back to <repo_root>/venvs/dynamo. Masked in CI by Dockerfile:371, but on the guide's own local-checkout path the preflight asserts against a different venv than the run uses. ${PROJECT_ROOT}/venvs/dynamo would align the three.

[jthomson04]

Addressed the L1 reliability items in 81d855f: readiness checks no longer depend on dictionary order, and the test-only Dynamo venv parameterization is reverted. The small wrapper keeps its explicit experiment-name export unchanged.

[jthomson04]

dynamo-generation.md:222

"Use the same command with three GPUs" points back at the two-GPU section, whose sbatch block hardcodes --gres=gpu:2. A reader following the page top-to-bottom submits a 3-GPU config onto a 2-GPU allocation and gets ResourceInsufficientError from placement-group creation after the retry backoff:

Suggested change

	export COMMAND
	export COMMAND
	sbatch \
	--nodes=1 \
	--gres=gpu:3 \
	--exclusive \
	--account=<account> \
	--partition=<partition> \
	ray.sub

Three other doc points, all small:

• dynamo-integration.md:22 says "Dynamo's NixlConnector". The connector is vLLM's; Dynamo only requires it be selected via --kv-transfer-config. Worth fixing since this paragraph is what a contributor reads when the pin moves.
• Under router_mode: kv, KV-aware routing applies to prefill selection only — Dynamo disables KV events on decode workers by design. The PR's argv already gets this right; the new P/D section just reads as though routing is symmetric.
• For the cross-node paragraph, it's worth noting NCCL_* tunes the refit transport only — KV transport is UCX, and vLLM's docs state NCCL vars "are not applicable to NixlConnector". Naming UCX_TLS/UCX_NET_DEVICES would help, and the good news is they pass through vllm_cfg.env_vars cleanly, so it's config-only.

[jthomson04]

Addressed the focused documentation items in 81d855f: the three-GPU sbatch request is complete, NixlConnector is identified as a vLLM connector, prefix-caching behavior is explicit, and the L1 is described as functional rather than a performance test. Broader transport-tuning guidance remains out of scope.

[jthomson04]

worker_pool.py:74

Non-blocking — a few contracts this PR introduces that are currently implicit.

worker_roles is positionally coupled to placement order (self._worker_roles[group_index]), and that same index becomes the engine's NCCL refit rank offset in refit.py:198-205. The rationale lives only as a comment in managed_runtime.py:107-108, so reordering the list or changing how start() iterates would silently re-map refit ranks:

Suggested change

	"""Reserve inference GPUs and launch one worker per model-parallel group."""
	class FixedDynamoWorkerPool:
	"""Reserve inference GPUs and launch one worker per model-parallel group.
	``worker_roles`` is positional: element ``i`` is assigned to the ``i``-th
	engine group in placement-group iteration order, and that same index becomes
	the engine's NCCL refit rank offset (see ``DynamoRefitChannel``). Callers
	must therefore pass decode roles before prefill roles and must not reorder
	the sequence between runs.
	"""

Also worth a line each, if you agree:

• arguments.py:127 — kv_event_port is tri-state against worker_role (required for prefill, forbidden for decode and aggregated, three distinct ValueErrors). Not discoverable from the signature. Noting the coupling is probably enough; the repo disables D417 so a full Args: block isn't expected.
• managed_runtime.py:358 — expected_components keys are Dynamo component names, not role names: decode and aggregated both register as backend, only prefill as prefill (args.py:175-185). That mismatch is the whole reason the two __init__ branches differ.
• worker_pool.py:51 — one port per engine is correct only because vLLM 0.23.0 computes VLLM_NIXL_SIDE_CHANNEL_PORT + data_parallel_index (scheduler.py:64-67) and multiplexes TP ranks over one socket, with dp pinned to 1. Concretely: a reviewer on this PR initially misread it as a TP>1 collision, so a comment naming the version would earn its keep — config.py:28-32 sets the precedent.
• arguments.py:145-149 — kv_role="kv_both" is deprecated in vLLM 0.23.0 but hard-required by Dynamo 1.3.0, whose error text names that exact JSON. Correct as written; a comment would save the next person the round-trip.
• virtual_cluster.py:97 (and the copy at ray.sub:168) — the new [4300, 4999) line reads as a global partition, but _bind_socket_in_range shuffles across the full [3000, 4999). No live collision since one backend runs per job; that invariant is what the comment should state.
• managed_runtime.py:98 — re-derives (prefill + decode) * engine_world_size by hand although this PR adds managed_inference_world_size for exactly that. validated_config is in scope, so expected_world_size = validated_config.managed_inference_world_size works and drops the formula from three places to two.

[jthomson04]

Keeping the internal documentation changes minimal. The existing decode-first ordering comment and user guide cover the refit-rank contract. This PR does not expand comments for every defensive invariant when there is no code defect to fix.

[jthomson04]

/ok to test 81d855f

[jthomson04]

Review update: commit 81d855f6f applies the focused correctness, readiness-test, L1, and documentation fixes. This PR remains functional in scope. It intentionally excludes performance and convergence sweeps, per-role metrics work, a new multi-node pool harness, and exhaustive tests for defensive raises.