[MIRROR — DO NOT MERGE] feat(dynamo): add native inference backend

docs/training.md

@@ -58,7 +58,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli
 |---|---|
 | `orchestrator.batch_size` | Tasks per trainer step. |
 | `orchestrator.group_size` | Rollouts generated per task. |
-| `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. |
+| `orchestrator.max_off_policy_steps` | On the vLLM admin backend, how many distinct policies may contribute to one rollout before it is discarded (default 8). Dynamo instead drains all live-policy/eval requests before mutating weights, so this setting does not apply to Dynamo runs. |
 | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). |
 | `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). |
 | `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). |

packages/prime-rl-configs/src/prime_rl/configs/inference.py

@@ -24,7 +24,7 @@ class ServerConfig(BaseConfig):
 
 
 class ParallelConfig(BaseConfig):
-    tp: int = 1
+    tp: int = Field(1, ge=1)
     """Tensor parallel size. Forwarded to vLLM as ``--tensor-parallel-size``."""
 
     dp: int = Field(1, ge=1)
@@ -251,7 +251,7 @@ def validate_scorers(self):
 
 
 class BaseInferenceDeploymentConfig(BaseConfig):
-    gpus_per_node: int = 8
+    gpus_per_node: int = Field(8, ge=1)
     """GPUs per node."""
 
     backend_port: int = 8100
@@ -304,10 +304,10 @@ class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig):
     """Extra environment variables exported only on decode nodes."""
 
     prefill_vllm_overrides: dict[str, Any] = {}
-    """Extra vLLM config options merged into --vllm-extra only for prefill ranks (SLURM only)."""
+    """Extra vLLM config options merged into the resolved config only for prefill workers."""
 
     decode_vllm_overrides: dict[str, Any] = {}
-    """Extra vLLM config options merged into --vllm-extra only for decode ranks (SLURM only)."""
+    """Extra vLLM config options merged into the resolved config only for decode workers."""
 
     @property
     def num_prefill_nodes(self) -> int:
@@ -328,7 +328,24 @@ def num_nodes(self) -> int:
 ]
 
 
+class VllmInferenceBackendConfig(BaseConfig):
+    type: Literal["vllm"] = "vllm"
+
+
+class DynamoInferenceBackendConfig(BaseConfig):
+    type: Literal["dynamo"] = "dynamo"
+
+
+InferenceBackendConfig: TypeAlias = Annotated[
+    VllmInferenceBackendConfig | DynamoInferenceBackendConfig,
+    Field(discriminator="type"),
+]
+
+
 class InferenceConfig(BaseConfig):
+    backend: InferenceBackendConfig = VllmInferenceBackendConfig()
+    """Serving backend. Existing configs default to Prime's native vLLM launcher."""
+
     server: ServerConfig = ServerConfig()
 
     model: ModelConfig = Field(default_factory=ModelConfig)
@@ -426,12 +443,81 @@ class InferenceConfig(BaseConfig):
     dry_run: bool = False
     """Only validate and dump resolved configs, then exit early."""
 
+    @property
+    def dynamo_worker_roles(self) -> tuple[Literal["agg", "prefill", "decode"], ...]:
+        """Canonical admin and launch order for Dynamo worker groups."""
+        if self.deployment.type == "disaggregated":
+            return ("prefill",) * self.deployment.num_prefill_replicas + (
+                "decode",
+            ) * self.deployment.num_decode_replicas
+        return ("agg",)
+
+    @property
+    def dynamo_gpus_per_worker(self) -> int:
+        """GPU allocation owned by one independently administered Dynamo worker."""
+        if self.deployment.type == "disaggregated":
+            return self.deployment.gpus_per_node
+        return self.parallel.tp * self.parallel.dp
+
+    @property
+    def dynamo_local_dp(self) -> int:
+        """vLLM data-parallel ranks inside one Dynamo worker."""
+        if self.deployment.type == "disaggregated":
+            return self.deployment.gpus_per_node // self.parallel.tp
+        return self.parallel.dp
+
     @model_validator(mode="after")
     def validate_multi_node_requires_slurm(self):
-        if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None:
+        if self.deployment.type == "multi_node" and self.slurm is None:
+            raise ValueError("Must use SLURM for multi-node deployment.")
+        if self.deployment.type == "disaggregated" and self.slurm is None and self.backend.type != "dynamo":
             raise ValueError("Must use SLURM for multi-node / disaggregated deployment.")
         return self
 
+    @model_validator(mode="after")
+    def validate_disaggregated_topology(self):
+        if self.deployment.type != "disaggregated":
+            return self
+        if self.deployment.gpus_per_node % self.parallel.tp != 0:
+            raise ValueError(
+                "inference.deployment.gpus_per_node must be divisible by inference.parallel.tp "
+                "so every worker contains whole tensor-parallel groups."
+            )
+        if self.backend.type == "dynamo":
+            local_dp = self.deployment.gpus_per_node // self.parallel.tp
+            if "dp" in self.parallel.model_fields_set and self.parallel.dp != local_dp:
+                raise ValueError(
+                    "inference.parallel.dp must equal inference.deployment.gpus_per_node / "
+                    "inference.parallel.tp for a Dynamo disaggregated worker."
+                )
+            if self.data_parallel_size_local is not None and self.data_parallel_size_local != local_dp:
+                raise ValueError(
+                    "inference.data_parallel_size_local must equal inference.deployment.gpus_per_node / "
+                    "inference.parallel.tp for a Dynamo disaggregated worker."
+                )
+        return self
+
+    @model_validator(mode="after")
+    def validate_dynamo_backend(self):
+        if self.backend.type != "dynamo":
+            return self
+        if self.slurm is not None:
+            raise ValueError(
+                "Dynamo is launched locally or through a DynamoGraphDeployment, not Prime's SLURM template."
+            )
+        if self.deployment.type == "multi_node":
+            raise ValueError("Dynamo multi-node inference must use a DynamoGraphDeployment.")
+        if self.enable_lora:
+            raise ValueError("The Dynamo backend does not support LoRA weight updates.")
+        router = getattr(self.deployment, "router", None)
+        if router is not None and router.type == "llm-d":
+            raise ValueError("The Dynamo backend owns request routing and cannot use the llm-d router.")
+        if self.deployment.type == "disaggregated":
+            if self.enable_prefix_caching is False:
+                raise ValueError("Dynamo disaggregated inference requires prefix caching for exact KV-aware routing.")
+            self.enable_prefix_caching = True
+        return self
+
     @model_validator(mode="after")
     def validate_llmd_no_routed_experts(self):
         """Reject routed-expert return with the llm-d router (breaks P/D, unverified for multi-node)."""
@@ -463,9 +549,7 @@ def auto_setup_disaggregated(self):
                 self.enable_expert_parallel = True
             if "enable_eplb" not in self.model_fields_set:
                 self.enable_eplb = False
-            gpus_per_node = self.deployment.gpus_per_node
-            tp = self.parallel.tp
-            dp_per_node = gpus_per_node // tp
+            dp_per_node = self.dynamo_local_dp
             if self.data_parallel_size_local is None:
                 self.data_parallel_size_local = dp_per_node
             if self.parallel.dp == 1:

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -529,7 +529,9 @@ class OrchestratorConfig(BaseConfig):
     """Maximum training steps. If None, runs indefinitely."""
 
     max_off_policy_steps: int = Field(8, ge=0)
-    """Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise."""
+    """Maximum policies allowed to generate one rollout on the vLLM admin backend. Dynamo uses a strict
+    application drain before every weight mutation, so live-policy and eval requests never span versions and
+    this tolerance does not apply there."""
 
     bench: bool = False
     """Benchmark mode. Sets ``max_steps`` to 5 and disables W&B."""
@@ -728,7 +730,6 @@ def resolve_env_config(self):
             if env.algo.sampling.source == "policy":
                 env.sampling.extra_body.setdefault("top_k", -1)
                 env.sampling.extra_body.setdefault("min_p", 0.0)
-                env.sampling.extra_body.setdefault("return_token_ids", True)
             if env.is_legacy:
                 # v0 env: cap per-turn response tokens to the training budget (the legacy
                 # bridge applies extra_env_kwargs via env.set_kwargs).

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -460,6 +460,12 @@ def auto_setup_lora(self):
 
         return self
 
+    @model_validator(mode="after")
+    def validate_auto_setup_does_not_enable_dynamo_lora(self):
+        if self.inference is not None and self.inference.backend.type == "dynamo" and self.inference.enable_lora:
+            raise ValueError("The Dynamo backend does not support LoRA weight updates.")
+        return self
+
     @model_validator(mode="after")
     def auto_setup_router_replay(self):
         if self.trainer.enable_router_replay:
@@ -534,7 +540,23 @@ def auto_setup_deployment(self):
             # fill up inference capacity with dp ranks
             if self.inference is not None:
                 num_infer_gpus = self.deployment.num_infer_gpus
-                if num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp:
+                is_dynamo_disaggregated = (
+                    self.inference.backend.type == "dynamo" and self.inference.deployment.type == "disaggregated"
+                )
+                if is_dynamo_disaggregated:
+                    infer_deploy = self.inference.deployment
+                    expected_infer_gpus = infer_deploy.num_nodes * infer_deploy.gpus_per_node
+                    if num_infer_gpus != expected_infer_gpus:
+                        raise ValueError(
+                            "deployment.num_infer_gpus must equal the Dynamo prefill/decode topology GPU count "
+                            f"({expected_infer_gpus}), got {num_infer_gpus}."
+                        )
+                    if self.weight_broadcast is not None and self.weight_broadcast.type == "nccl":
+                        assert self.trainer.weight_broadcast.type == "nccl"
+                        self.trainer.weight_broadcast.inference_world_size = expected_infer_gpus
+                        assert self.orchestrator.weight_broadcast.type == "nccl"
+                        self.orchestrator.weight_broadcast.inference_world_size = expected_infer_gpus
+                elif num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp:
                     assert num_infer_gpus % self.inference.parallel.tp == 0, (
                         "Number of inference GPUs must be divisible by the tensor parallel size"
                     )
@@ -675,15 +697,38 @@ def auto_setup_inference_client(self):
         if self.inference is None:
             return self
         client = self.orchestrator.model.client
+        client_updates: dict[str, Any] = {}
+        if "admin_api" in client.model_fields_set and client.admin_api != self.inference.backend.type:
+            raise ValueError(
+                "orchestrator.model.client.admin_api conflicts with inference.backend.type; "
+                "configure the backend only under inference."
+            )
+        client_updates["admin_api"] = self.inference.backend.type
         if "dp_rank_count" not in client.model_fields_set:
-            if self.deployment.type == "multi_node":
-                client.dp_rank_count = 1
+            if self.inference.backend.type == "dynamo" or self.deployment.type == "multi_node":
+                client_updates["dp_rank_count"] = 1
             else:
-                client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp
+                client_updates["dp_rank_count"] = self.inference.data_parallel_size_local or self.inference.parallel.dp
+        if self.inference.backend.type == "dynamo":
+            expected_topology = {
+                "dynamo_worker_roles": self.inference.dynamo_worker_roles,
+                "dynamo_gpus_per_worker": self.inference.dynamo_gpus_per_worker,
+            }
+            for field, expected in expected_topology.items():
+                if field in client.model_fields_set and getattr(client, field) != expected:
+                    raise ValueError(
+                        f"orchestrator.model.client.{field} conflicts with the inference topology; "
+                        "configure the topology only under inference."
+                    )
+            client_updates.update(expected_topology)
         if not self.orchestrator.any_policy_sourced and "base_url" not in client.model_fields_set:
             host = self.inference.server.host or "localhost"
             port = self.inference.server.port
-            client.base_url = [f"http://{host}:{port}/v1"]
+            client_updates["base_url"] = [f"http://{host}:{port}/v1"]
+
+        updated_client = client.model_copy(update=client_updates)
+        updated_model = self.orchestrator.model.model_copy(update={"client": updated_client})
+        self.orchestrator = self.orchestrator.model_copy(update={"model": updated_model})
         return self
 
     @model_validator(mode="after")

packages/prime-rl-configs/src/prime_rl/configs/shared.py

@@ -147,6 +147,18 @@ class ClientConfig(BaseConfig):
     admin_base_url: list[str] | None = None
     """Separate base URLs for admin operations (weight updates, health checks). When set, admin clients bypass routers and hit each server directly — used in disaggregated P/D deployments where the router must not handle admin traffic."""
 
+    admin_api: Literal["vllm", "dynamo"] = "vllm"
+    """Admin protocol used for health and weight updates. Auto-derived from the local inference backend."""
+
+    rl_base_url: list[str] | None = None
+    """Dynamo RL worker-discovery URLs. When omitted, they are derived from the frontend URL or ``DYN_RL_DISCOVERY_URL``."""
+
+    dynamo_worker_roles: tuple[Literal["agg", "prefill", "decode"], ...] | None = None
+    """Exact Dynamo worker roles expected during readiness. Auto-derived from the local inference topology."""
+
+    dynamo_gpus_per_worker: int | None = Field(None, ge=1)
+    """GPUs owned by each discovered Dynamo worker. Auto-derived from the local inference topology."""
+
     elastic: ElasticConfig | None = None
     """Elastic inference pool config for DNS-based service discovery. When set, ``base_url`` is ignored and inference servers are discovered dynamically via DNS."""
 
@@ -255,7 +267,8 @@ class MetricsServerConfig(BaseConfig):
 
 
 class BaseTransportConfig(BaseConfig):
-    pass
+    send_timeout_seconds: float = Field(300.0, gt=0, allow_inf_nan=False)
+    """Maximum time to ship one rollout batch before the orchestrator fails closed."""
 
 
 class FileSystemTransportConfig(BaseTransportConfig):

src/prime_rl/entrypoints/inference.py

@@ -148,6 +148,13 @@ def inference_local(config: InferenceConfig):
     logger = setup_logger(config.log.level, json_logging=config.log.json_logging)
 
     if config.dry_run:
+        if config.backend.type == "dynamo":
+            from prime_rl.inference.dynamo import build_dry_run_worker_specs, build_frontend_process
+
+            specs = build_dry_run_worker_specs(config)
+            logger.info(f"Dynamo frontend: {' '.join(build_frontend_process(config).command())}")
+            for spec in specs:
+                logger.info(f"Dynamo {spec.name}: {' '.join(spec.process.command())}")
         logger.success("Dry run complete. To start inference locally, remove --dry-run from your command.")
         return
 
@@ -162,6 +169,12 @@ def inference_local(config: InferenceConfig):
 
     setup_vllm_env(config)
 
+    if config.backend.type == "dynamo":
+        from prime_rl.inference.dynamo import run_dynamo_local
+
+        run_dynamo_local(config)
+        return
+
     from prime_rl.inference.vllm.server import server  # pyright: ignore
 
     server(config, vllm_extra=config.vllm_extra)