From 15d29eea5eed253bc821550f5a6926a24731b56a Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 01/19] feat(dynamo): AdminAPI abstraction + dynamo backend
 selector + RL worker discovery

---
 .../src/prime_rl/configs/shared.py            |   6 +
 src/prime_rl/orchestrator/orchestrator.py     |   5 +
 src/prime_rl/utils/client.py                  | 503 +++++++++++++-----
 src/prime_rl/utils/elastic.py                 |  22 +-
 tests/unit/utils/test_client.py               | 102 +++-
 5 files changed, 501 insertions(+), 137 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index ff311f145d..9c16c410bd 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -136,6 +136,12 @@ class ClientConfig(BaseConfig):
     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."""
 
+    rl_base_url: list[str] | None = None
+    """Dynamo RL worker discovery base URLs. Used only for backend='dynamo' when admin_base_url is unset. These URLs point at the Dynamo RL discovery listener (DYN_RL_PORT, default 8001), which serves GET /v1/rl/workers. If unset, prime-rl derives the discovery URL from base_url by replacing the port with DYN_RL_PORT or 8001."""
+
+    backend: Literal["vllm", "dynamo"] = "vllm"
+    """Inference backend selector. Picks the AdminAPI implementation used for pause/resume/update_weights/load_lora_adapter/list_models. Default 'vllm' matches prime-rl's bundled vLLM frontend. 'dynamo' targets NVIDIA Dynamo's worker /engine/* admin routes on admin_base_url and routes /v1/models to the OpenAI-compat base_url."""
+
     router_url: str | None = None
     """vllm-router URL for load-aware inference routing. With elastic mode, inference requests go through the router while admin ops still hit discovered pods directly."""
 
diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py
index d685b519b6..8b21c7fa4d 100644
--- a/src/prime_rl/orchestrator/orchestrator.py
+++ b/src/prime_rl/orchestrator/orchestrator.py
@@ -314,6 +314,10 @@ async def setup(self) -> None:
             await self.inference_metrics.start()
 
         get_logger().info(f"Initializing weight broadcast ({config.weight_broadcast})")
+        # Propagate the configured broadcast type to the admin API. The Dynamo
+        # backend gates its NCCL init path on ``_weight_broadcast_type``.
+        if hasattr(self.student_inference._admin_api, "_weight_broadcast_type"):
+            self.student_inference._admin_api._weight_broadcast_type = config.weight_broadcast.type
         if config.weight_broadcast.type == "nccl":
             await init_nccl_broadcast(
                 self.student_inference.admin_clients,
@@ -322,6 +326,7 @@ async def setup(self) -> None:
                 config.weight_broadcast.timeout,
                 inference_world_size=config.weight_broadcast.inference_world_size,
                 quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer,
+                admin=self.student_inference._admin_api,
             )
 
         get_logger().info(f"Initializing training batch sender ({config.rollout_transport})")
diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py
index b9ee8f4b9d..fddab3089b 100644
--- a/src/prime_rl/utils/client.py
+++ b/src/prime_rl/utils/client.py
@@ -6,6 +6,7 @@
 from itertools import cycle
 from pathlib import Path
 from typing import Protocol, runtime_checkable
+from urllib.parse import urlsplit, urlunsplit
 
 import httpx
 import verifiers as vf
@@ -28,6 +29,237 @@ def client_identity(client: vf.ClientConfig) -> ClientIdentity:
     return (client.api_base_url, client.extra_headers.get("X-data-parallel-rank"))
 
 
+class AdminAPI(Protocol):
+    """Admin endpoints for an inference backend.
+
+    Per-method: construct one HTTP call. Per-server parallelism, retry, and
+    raise-for-status policy live in the caller.
+    """
+
+    async def health(self, client: AsyncClient) -> None: ...
+    async def list_models(self, client: AsyncClient) -> list[dict]: ...
+    async def pause(self, client: AsyncClient) -> None: ...
+    async def resume(self, client: AsyncClient) -> None: ...
+    async def update_weights(self, client: AsyncClient, weight_dir: str | None) -> None: ...
+    async def load_lora_adapter(
+        self,
+        client: AsyncClient,
+        lora_name: str,
+        lora_path: str,
+        *,
+        timeout: httpx.Timeout,
+    ) -> None: ...
+    async def init_broadcaster(
+        self,
+        client: AsyncClient,
+        *,
+        host: str,
+        port: int,
+        rank_offset: int,
+        inference_world_size: int,
+        timeout: int,
+        quantize_in_weight_transfer: bool,
+    ) -> None: ...
+
+
+class VLLMAdminAPI:
+    """vLLM admin endpoints."""
+
+    async def health(self, client: AsyncClient) -> None:
+        # No raise_for_status: any HTTP response means the server is up.
+        # Only transport errors mean "not ready yet" (caller retries).
+        await client.get("/health")
+
+    async def list_models(self, client: AsyncClient) -> list[dict]:
+        response = await client.get("/v1/models")
+        return response.json()["data"]
+
+    async def pause(self, client: AsyncClient) -> None:
+        response = await client.post("/pause", params={"mode": "keep", "clear_cache": "false"})
+        response.raise_for_status()
+
+    async def resume(self, client: AsyncClient) -> None:
+        response = await client.post("/resume")
+        response.raise_for_status()
+
+    async def update_weights(self, client: AsyncClient, weight_dir: str | None) -> None:
+        response = await client.post("/update_weights", json={"weight_dir": weight_dir})
+        response.raise_for_status()
+
+    async def load_lora_adapter(
+        self,
+        client: AsyncClient,
+        lora_name: str,
+        lora_path: str,
+        *,
+        timeout: httpx.Timeout,
+    ) -> None:
+        response = await client.post(
+            "/load_lora_adapter",
+            json={"lora_name": lora_name, "lora_path": lora_path},
+            timeout=timeout,
+        )
+        response.raise_for_status()
+
+    async def init_broadcaster(
+        self,
+        client: AsyncClient,
+        *,
+        host: str,
+        port: int,
+        rank_offset: int,
+        inference_world_size: int,
+        timeout: int,
+        quantize_in_weight_transfer: bool,
+    ) -> None:
+        response = await client.post(
+            "/init_broadcaster",
+            json={
+                "host": host,
+                "port": port,
+                "rank_offset": rank_offset,
+                "inference_world_size": inference_world_size,
+                "timeout": timeout,
+                "quantize_in_weight_transfer": quantize_in_weight_transfer,
+            },
+        )
+        response.raise_for_status()
+
+
+class DynamoAdminAPI(VLLMAdminAPI):
+    """NVIDIA Dynamo worker admin endpoints via ``POST /engine/<method>``.
+
+    Each Dynamo worker exposes engine routes on its system status server
+    (``DYN_SYSTEM_PORT``, default 8081). Multi-worker deployments are handled by
+    iterating over ``admin_clients``.
+
+    Args:
+        engine_rpc: The ``collective_rpc`` target forwarded by
+            ``update_weights_from_disk``.  Use ``"reload_weights"`` for plain
+            vLLM / dynamo.vllm without a worker extension (default).  Use
+            ``"update_weights_from_path"`` only when
+            FileSystemWeightUpdateWorker / NCCLWeightUpdateWorker is loaded via
+            ``--worker-extension-cls``.
+    """
+
+    def __init__(self, engine_rpc: str = "reload_weights", weight_broadcast_type: str = "filesystem") -> None:
+        self._engine_rpc = engine_rpc
+        # Determines which engine method is called per step: "update_weights_from_distributed"
+        # for NCCL (trainer broadcasts; worker just needs to receive) vs
+        # "update_weights_from_disk" for filesystem. Set externally by the orchestrator
+        # once weight_broadcast config is resolved. Defaults to filesystem (run #35 behaviour).
+        self._weight_broadcast_type = weight_broadcast_type
+
+    async def health(self, client: AsyncClient) -> None:
+        await client.get("/health")
+
+    async def _post_engine(
+        self,
+        client: AsyncClient,
+        method: str,
+        body: dict | None = None,
+        *,
+        timeout: httpx.Timeout | None = None,
+    ) -> dict:
+        response = await client.post(f"/engine/{method}", json=body or {}, timeout=timeout)
+        response.raise_for_status()
+        data = response.json()
+        if isinstance(data, dict) and data.get("status") == "error":
+            raise RuntimeError(data.get("message", f"Dynamo /engine/{method} failed"))
+        return data
+
+    async def pause(self, client: AsyncClient) -> None:
+        await self._post_engine(client, "pause_generation", {"mode": "keep", "clear_cache": False})
+
+    async def resume(self, client: AsyncClient) -> None:
+        await self._post_engine(client, "resume_generation")
+
+    async def update_weights(self, client: AsyncClient, weight_dir: str | None) -> None:
+        if weight_dir is None:
+            return
+        if self._weight_broadcast_type == "nccl":
+            # NCCL path: trainer has already broadcast weights via the NCCL group;
+            # this RPC tells the inference worker to call receive_state_dict().
+            # NCCLWeightUpdateWorker exposes "update_weights_from_path", not "reload_weights".
+            await self._post_engine(
+                client,
+                "update_weights_from_distributed",
+                {
+                    "weight_version": Path(weight_dir).name,
+                    "weight_dir": weight_dir,
+                    "engine_rpc": "update_weights_from_path",
+                },
+                timeout=httpx.Timeout(180.0),
+            )
+        else:
+            # Resolve to absolute path so the inference worker (which may run in a
+            # different working directory) can find the checkpoint on the shared NFS.
+            abs_path = str(Path(weight_dir).resolve())
+            await self._post_engine(
+                client,
+                "update_weights_from_disk",
+                {
+                    "model_path": abs_path,
+                    "weight_version": Path(weight_dir).name,
+                    "engine_rpc": self._engine_rpc,
+                },
+                timeout=httpx.Timeout(180.0),
+            )
+
+    async def load_lora_adapter(
+        self,
+        client: AsyncClient,
+        lora_name: str,
+        lora_path: str,
+        *,
+        timeout: httpx.Timeout,
+    ) -> None:
+        await self._post_engine(
+            client,
+            "load_lora",
+            {
+                "lora_name": lora_name,
+                "source": {"uri": Path(lora_path).absolute().as_uri()},
+            },
+            timeout=timeout,
+        )
+
+    async def init_broadcaster(
+        self,
+        client: AsyncClient,
+        *,
+        host: str,
+        port: int,
+        rank_offset: int,
+        inference_world_size: int,
+        timeout: int,
+        quantize_in_weight_transfer: bool,
+    ) -> None:
+        await self._post_engine(
+            client,
+            "init_weights_update_group",
+            {
+                "host": host,
+                "port": port,
+                "rank_offset": rank_offset,
+                "inference_world_size": inference_world_size,
+                "timeout": timeout,
+                "quantize_in_weight_transfer": quantize_in_weight_transfer,
+                "engine_rpc": "init_broadcaster",
+            },
+        )
+
+
+def setup_admin_api(client_config: ClientConfig) -> AdminAPI:
+    """Pick the AdminAPI implementation that matches ``client_config.backend``."""
+    if client_config.backend == "dynamo":
+        return DynamoAdminAPI()
+    return VLLMAdminAPI()
+
+
+_DEFAULT_ADMIN: AdminAPI = VLLMAdminAPI()
+
+
 @runtime_checkable
 class InferencePool(Protocol):
     """Protocol for inference pools (static or elastic)."""
@@ -103,6 +335,12 @@ def __init__(
         )
         self._eval_clients = setup_clients(client_config, client_type=eval_client_type)
         self._admin_clients = setup_admin_clients(client_config)
+        self._model_clients = (
+            setup_admin_clients(client_config, use_admin_base_url=False)
+            if client_config.backend == "dynamo" or client_config.admin_base_url
+            else self._admin_clients
+        )
+        self._admin_api = setup_admin_api(client_config)
         self._skip_model_check = client_config.skip_model_check
         self._wait_for_ready_timeout = client_config.wait_for_ready_timeout
         self._eval_cycle = cycle(self._eval_clients)
@@ -133,12 +371,16 @@ async def select_train_client(self, load: Mapping[ClientIdentity, int]) -> vf.Cl
 
     async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None:
         await check_health(
-            self._admin_clients, timeout=timeout if timeout is not None else self._wait_for_ready_timeout
+            self._admin_clients,
+            timeout=timeout if timeout is not None else self._wait_for_ready_timeout,
+            admin=self._admin_api,
+        )
+        await maybe_check_has_model(
+            self._model_clients, model_name, skip_model_check=self._skip_model_check, admin=self._admin_api
         )
-        await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check)
 
     async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None:
-        await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step)
+        await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step, admin=self._admin_api)
 
     def get_metrics(self) -> dict[str, float]:
         return {}
@@ -185,6 +427,13 @@ def setup_clients(
     renderer_model_name: str | None = None,
     pool_size: int | None = None,
 ) -> list[vf.ClientConfig]:
+    # Pick the verifiers wire-shape selector based on client_config.backend.
+    # When backend == "dynamo", both RendererClient and
+    # OpenAIChatCompletionsTokenClient route through Dynamo's nvext path:
+    #   - request:  nvext.token_data carries pre-tokenized prompt
+    #   - response: nvext.engine_data carries completion_token_ids + logprobs
+    # Default backend keeps the legacy vLLM TITO surface.
+    renderer_transport = "dynamo_chat_nvext" if client_config.backend == "dynamo" else "prime_vllm_generate"
     clients = []
     client_idx = 0
     # Only forward the renderer config when the client actually uses a
@@ -208,6 +457,9 @@ def setup_clients(
                 vf.ClientConfig(
                     client_idx=client_idx,
                     client_type=client_type,
+                    # Dynamo backend routes both renderer and token clients through
+                    # the nvext path; default backend keeps the legacy vLLM TITO surface.
+                    renderer_transport=renderer_transport,
                     api_base_url=base_url,
                     api_key_var=client_config.api_key_var,
                     timeout=client_config.timeout,
@@ -224,14 +476,21 @@ def setup_clients(
     return clients
 
 
-def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]:
+def setup_admin_clients(client_config: ClientConfig, *, use_admin_base_url: bool = True) -> list[AsyncClient]:
     """Create dedicated admin clients for weight update operations.
 
     Uses a separate connection pool to avoid queueing behind streaming requests.
-    When admin_base_url is set, uses those URLs instead of base_url, allowing
-    weight updates to bypass routers in disaggregated P/D deployments.
+    When admin_base_url is set and use_admin_base_url is true, uses those URLs
+    instead of base_url, allowing weight updates to bypass routers in
+    disaggregated P/D deployments. For Dynamo, if admin_base_url is unset,
+    discover worker-advertised system URLs from GET /v1/rl/workers.
     """
-    urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url
+    if use_admin_base_url and client_config.admin_base_url:
+        urls = client_config.admin_base_url
+    elif use_admin_base_url and client_config.backend == "dynamo":
+        urls = discover_dynamo_admin_base_urls(client_config)
+    else:
+        urls = client_config.base_url
 
     def _setup_admin_client(base_url: str) -> httpx.AsyncClient:
         env_headers = {
@@ -255,23 +514,84 @@ def _setup_admin_client(base_url: str) -> httpx.AsyncClient:
     return [_setup_admin_client(base_url) for base_url in urls]
 
 
+def discover_dynamo_admin_base_urls(client_config: ClientConfig) -> list[str]:
+    urls: list[str] = []
+    headers = client_config.headers.copy()
+    api_key = os.getenv(client_config.api_key_var, "EMPTY")
+    if api_key and api_key != "EMPTY":
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    for base_url in _dynamo_rl_discovery_base_urls(client_config):
+        discovery_base = base_url.rstrip("/").removesuffix("/v1")
+        with httpx.Client(
+            base_url=discovery_base,
+            headers=headers,
+            timeout=httpx.Timeout(connect=client_config.connect_timeout, read=30.0, write=30.0, pool=10.0),
+        ) as client:
+            response = client.get("/v1/rl/workers")
+            response.raise_for_status()
+            for worker in response.json().get("workers", []):
+                system_url = worker.get("system_url")
+                if system_url:
+                    urls.append(system_url)
+
+    deduped = list(dict.fromkeys(urls))
+    if not deduped:
+        raise ValueError(
+            "Dynamo backend did not discover any worker system URLs from /v1/rl/workers. "
+            "Set client.admin_base_url explicitly, set client.rl_base_url to the Dynamo "
+            "RL discovery listener, and make sure Dynamo workers run with DYN_ENABLE_RL "
+            "and a system status server enabled."
+        )
+    return deduped
+
+
+def _dynamo_rl_discovery_base_urls(client_config: ClientConfig) -> list[str]:
+    configured = getattr(client_config, "rl_base_url", None)
+    if configured:
+        return configured
+
+    rl_port = int(os.getenv("DYN_RL_PORT", "8001"))
+    return [_replace_url_port(base_url, rl_port) for base_url in client_config.base_url]
+
+
+def _replace_url_port(base_url: str, port: int) -> str:
+    parsed = urlsplit(base_url.rstrip("/").removesuffix("/v1"))
+    scheme = parsed.scheme or "http"
+    host = parsed.hostname or parsed.netloc
+    if not host:
+        raise ValueError(f"Cannot derive Dynamo RL discovery URL from base_url={base_url!r}")
+    if ":" in host and not host.startswith("["):
+        host = f"[{host}]"
+    netloc = f"{host}:{port}"
+    return urlunsplit((scheme, netloc, "", "", ""))
+
+
 async def maybe_check_has_model(
-    admin_clients: list[AsyncClient], model_name: str, skip_model_check: bool = False
+    admin_clients: list[AsyncClient],
+    model_name: str,
+    skip_model_check: bool = False,
+    *,
+    admin: AdminAPI = _DEFAULT_ADMIN,
 ) -> None:
     if skip_model_check:
         return
     logger = get_logger()
     logger.debug(f"Checking if model {model_name} is in the inference pool")
-    results = await asyncio.gather(*[admin_client.get("/v1/models") for admin_client in admin_clients])
-    for admin_client, result in zip(admin_clients, results):
-        models = result.json()["data"]
+    results = await asyncio.gather(*[admin.list_models(admin_client) for admin_client in admin_clients])
+    for admin_client, models in zip(admin_clients, results):
         if not any(model["id"] == model_name for model in models):
             raise ValueError(f"Model {model_name} was not found in the inference pool on {admin_client.base_url}")
     logger.debug(f"Model {model_name} was found in the inference pool")
 
 
 async def check_health(
-    admin_clients: list[AsyncClient], interval: int = 1, log_interval: int = 10, timeout: int = 1800
+    admin_clients: list[AsyncClient],
+    interval: int = 1,
+    log_interval: int = 10,
+    timeout: int = 1800,
+    *,
+    admin: AdminAPI = _DEFAULT_ADMIN,
 ) -> None:
     logger = get_logger()
 
@@ -280,7 +600,7 @@ async def _check_health(admin_client: AsyncClient) -> None:
         logger.debug("Starting pinging /health to check health")
         while wait_time < timeout:
             try:
-                await admin_client.get("/health")
+                await admin.health(admin_client)
                 logger.debug(f"Inference pool is ready after {wait_time} seconds")
                 return
             except NotFoundError:
@@ -303,105 +623,44 @@ async def _check_health(admin_client: AsyncClient) -> None:
 NCCL_READY_MARKER = "NCCL_READY"
 
 
-def _is_retryable_pause_error(exception: BaseException) -> bool:
-    """Check if an exception should trigger a retry for pausing engines."""
-    if isinstance(exception, httpx.HTTPStatusError):
-        # Retry on transient server errors (5xx, e.g. engine briefly unresponsive);
-        # client errors (4xx) won't fix themselves on retry.
-        return exception.response.status_code >= 500
-    # Retry on transport-level failures (timeouts, connection resets, etc.) so the
-    # per-attempt read timeout below turns a stuck server into a bounded retry loop
-    # instead of hanging forever on the global timeout=None admin client.
-    if isinstance(exception, (httpx.TimeoutException, httpx.TransportError)):
-        return True
-    return False
-
-
-# Per-attempt and total bounds for `/pause`. Pausing drains in-flight requests
-# (mode="keep"), so a single attempt can legitimately take a while, but the global
-# admin AsyncClient uses `timeout=None`, so a stuck server would hang the weight
-# update forever. `_READ_TIMEOUT` converts a hang into a TimeoutException so
-# tenacity retries; `_TOTAL` is the wall-clock budget across all retries.
-PAUSE_READ_TIMEOUT_S = 120.0
-PAUSE_TOTAL_TIMEOUT_S = 300.0
-
-
-async def _pause_engines(admin_clients: list[AsyncClient], *, step: int) -> None:
-    """Pause all inference engines, waiting for in-flight requests to drain."""
-    logger = get_logger()
-    logger.info(f"Updating policy in-flight to v{step}")
-
-    @retry(
-        retry=retry_if_exception(_is_retryable_pause_error),
-        stop=stop_after_delay(PAUSE_TOTAL_TIMEOUT_S) | stop_after_attempt(10),
-        wait=wait_exponential(multiplier=1, min=1, max=10),
-        reraise=True,
-    )
-    async def _pause(client: AsyncClient) -> None:
-        response = await client.post(
-            "/pause",
-            params={"mode": "keep", "clear_cache": "false"},
-            timeout=httpx.Timeout(connect=10.0, read=PAUSE_READ_TIMEOUT_S, write=60.0, pool=10.0),
-        )
-        response.raise_for_status()
-
-    await asyncio.gather(*[_pause(client) for client in admin_clients])
-    logger.debug("All inference engines paused")
-
-
-async def _resume_engines(admin_clients: list[AsyncClient]) -> None:
-    """Resume all inference engines after weight update."""
-    logger = get_logger()
-
-    async def _resume(client: AsyncClient) -> None:
-        response = await client.post("/resume")
-        response.raise_for_status()
-
-    await asyncio.gather(*[_resume(client) for client in admin_clients])
-    logger.debug("All inference engines resumed")
-
-
 async def update_weights(
     admin_clients: list[AsyncClient],
     weight_dir: Path | None,
     lora_name: str | None = None,
     step: int = 0,
+    *,
+    admin: AdminAPI = _DEFAULT_ADMIN,
 ) -> None:
     """Update weights on static inference servers.
 
-    Pauses all engines first to drain in-flight requests, then performs the
-    weight update, then resumes. This ensures all DP workers are idle and can
-    participate in the collective weight transfer.
-
-    Note: The server-side /update_weights endpoint automatically resets the prefix cache
-    to invalidate any cached KV states computed with the old weights.
+    Pauses all engines to drain in-flight requests, performs the weight update,
+    then resumes. Ensures all DP workers are idle and can participate in the
+    collective weight transfer. The server-side ``/update_weights`` endpoint
+    resets the prefix cache to invalidate any KV states computed with the old
+    weights.
     """
     logger = get_logger()
 
-    weight_dir_posix = weight_dir.as_posix() if weight_dir is not None else None
-
     if lora_name is not None and weight_dir is not None:
-        await load_lora_adapter(admin_clients, lora_name, weight_dir)
-    else:
-
-        async def _update_weights(admin_client: AsyncClient, weight_dir: str | None) -> None:
-            response = await admin_client.post("/update_weights", json={"weight_dir": weight_dir})
-            response.raise_for_status()
+        await load_lora_adapter(admin_clients, lora_name, weight_dir, admin=admin)
+        return
 
-        # Pause engines so all DP workers drain in-flight work and can join the NCCL broadcast
-        await _pause_engines(admin_clients, step=step)
+    weight_dir_posix = weight_dir.as_posix() if weight_dir is not None else None
 
-        try:
-            # Create ready marker before servers enter receive path (used by NCCL broadcast)
-            if weight_dir is not None:
-                nccl_ready_file = weight_dir / NCCL_READY_MARKER
-                nccl_ready_file.parent.mkdir(parents=True, exist_ok=True)
-                nccl_ready_file.touch()
-                logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}")
+    logger.info("Pausing inference engines for weight update")
+    await asyncio.gather(*[admin.pause(c) for c in admin_clients])
+    try:
+        # NCCL_READY marker is created before servers enter the receive path
+        if weight_dir is not None:
+            nccl_ready_file = weight_dir / NCCL_READY_MARKER
+            nccl_ready_file.parent.mkdir(parents=True, exist_ok=True)
+            nccl_ready_file.touch()
+            logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}")
 
-            await asyncio.gather(*[_update_weights(admin_client, weight_dir_posix) for admin_client in admin_clients])
-        finally:
-            await _resume_engines(admin_clients)
+        await asyncio.gather(*[admin.update_weights(c, weight_dir_posix) for c in admin_clients])
+    finally:
+        await asyncio.gather(*[admin.resume(c) for c in admin_clients])
+        logger.info("Inference engines resumed")
 
 
 def _is_retryable_lora_error(exception: BaseException) -> bool:
@@ -428,7 +687,13 @@ def _is_retryable_lora_error(exception: BaseException) -> bool:
 LORA_LOAD_TOTAL_TIMEOUT_S = 120.0
 
 
-async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lora_path: Path) -> None:
+async def load_lora_adapter(
+    admin_clients: list[AsyncClient],
+    lora_name: str,
+    lora_path: Path,
+    *,
+    admin: AdminAPI = _DEFAULT_ADMIN,
+) -> None:
     """Make a HTTP post request to the vLLM server to load a LoRA adapter.
 
     Uses our wrapper endpoint that also resets the prefix cache to invalidate
@@ -439,6 +704,7 @@ async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lo
     """
     logger = get_logger()
     lora_path_posix = lora_path.as_posix()
+    per_attempt_timeout = httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0)
 
     @retry(
         retry=retry_if_exception(_is_retryable_lora_error),
@@ -448,29 +714,11 @@ async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lo
     )
     async def _load_lora_adapter(admin_client: AsyncClient) -> None:
         logger.debug(f"Sending request to load LoRA adapter {lora_name} from {lora_path}")
-        response = await admin_client.post(
-            "/load_lora_adapter",
-            json={"lora_name": lora_name, "lora_path": lora_path_posix},
-            timeout=httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0),
-        )
-        response.raise_for_status()
+        await admin.load_lora_adapter(admin_client, lora_name, lora_path_posix, timeout=per_attempt_timeout)
 
     await asyncio.gather(*[_load_lora_adapter(admin_client) for admin_client in admin_clients])
 
 
-async def unload_lora_adapter(admin_clients: list[AsyncClient], lora_name: str) -> None:
-    """Make a HTTP post request to the vLLM server to unload a LoRA adapter."""
-    logger = get_logger()
-
-    async def _unload_lora_adapter(admin_client: AsyncClient) -> None:
-        logger.debug(f"Sending request to unload LoRA adapter {lora_name}")
-        await admin_client.post("/v1/unload_lora_adapter", json={"lora_name": lora_name})
-        # TODO: The first one can fail, but subsequent ones should succeed.
-        # response.raise_for_status()
-
-    await asyncio.gather(*[_unload_lora_adapter(admin_client) for admin_client in admin_clients])
-
-
 async def init_nccl_broadcast(
     admin_clients: list[AsyncClient],
     host: str,
@@ -478,6 +726,8 @@ async def init_nccl_broadcast(
     timeout: int,
     inference_world_size: int | None = None,
     quantize_in_weight_transfer: bool = False,
+    *,
+    admin: AdminAPI = _DEFAULT_ADMIN,
 ) -> None:
     """Initialize NCCL broadcast on all inference servers.
 
@@ -502,18 +752,15 @@ async def init_nccl_broadcast(
 
     async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> None:
         try:
-            response = await admin_client.post(
-                "/init_broadcaster",
-                json={
-                    "host": host,
-                    "port": port,
-                    "rank_offset": rank_offset,
-                    "inference_world_size": inference_world_size,
-                    "timeout": timeout,
-                    "quantize_in_weight_transfer": quantize_in_weight_transfer,
-                },
+            await admin.init_broadcaster(
+                admin_client,
+                host=host,
+                port=port,
+                rank_offset=rank_offset,
+                inference_world_size=inference_world_size,
+                timeout=timeout,
+                quantize_in_weight_transfer=quantize_in_weight_transfer,
             )
-            response.raise_for_status()
         except httpx.HTTPStatusError as e:
             if e.response.status_code == 404:
                 logger.warning("The route /init_broadcaster does not exist. Skipping NCCL broadcast initialization.")
diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py
index 951b3673c1..8cee2e4cca 100644
--- a/src/prime_rl/utils/elastic.py
+++ b/src/prime_rl/utils/elastic.py
@@ -22,7 +22,14 @@
 from renderers import RendererConfig
 
 from prime_rl.configs.shared import ClientConfig
-from prime_rl.utils.client import ClientIdentity, client_identity, load_lora_adapter, setup_admin_clients, setup_clients
+from prime_rl.utils.client import (
+    ClientIdentity,
+    client_identity,
+    load_lora_adapter,
+    setup_admin_api,
+    setup_admin_clients,
+    setup_clients,
+)
 from prime_rl.utils.logger import get_logger
 
 # --- Shared discovery functions ---
@@ -127,6 +134,7 @@ def __init__(
 
         self._servers: dict[str, ServerState] = {}
         self._admin_clients: dict[str, AsyncClient] = {}
+        self._admin_api = setup_admin_api(client_config)
         self._lock = asyncio.Lock()
         self._desired: AdapterState = AdapterState()
 
@@ -334,7 +342,9 @@ async def _sync_server_adapter(self, ip: str) -> bool:
         if self._desired.name and self._desired.path:
             try:
                 self.logger.debug(f"Loading adapter {self._desired.name} on {ip}")
-                await load_lora_adapter([self._admin_clients[ip]], self._desired.name, self._desired.path)
+                await load_lora_adapter(
+                    [self._admin_clients[ip]], self._desired.name, self._desired.path, admin=self._admin_api
+                )
             except Exception as e:
                 server.status = "unhealthy"
                 server.sync_failures += 1
@@ -369,12 +379,8 @@ async def _check_server_health(self, admin_client: AsyncClient, ip: str) -> bool
             return False
 
         try:
-            response = await admin_client.get("/v1/models")
-            response.raise_for_status()
-            data = response.json()
-            models = [m.get("id") for m in data.get("data", [])]
-
-            if self.base_model_name not in models:
+            models = await self._admin_api.list_models(admin_client)
+            if self.base_model_name not in [m.get("id") for m in models]:
                 self.logger.debug(f"Server {ip} does not have base model {self.base_model_name}")
                 return False
         except Exception as e:
diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py
index 69325e6c4a..002680bd0c 100644
--- a/tests/unit/utils/test_client.py
+++ b/tests/unit/utils/test_client.py
@@ -6,7 +6,14 @@
 import verifiers as vf
 
 from prime_rl.configs.shared import ClientConfig
-from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients
+from prime_rl.utils.client import (
+    DynamoAdminAPI,
+    _dynamo_rl_discovery_base_urls,
+    _is_retryable_lora_error,
+    discover_dynamo_admin_base_urls,
+    load_lora_adapter,
+    setup_clients,
+)
 
 
 def test_is_retryable_lora_error_returns_true_for_404():
@@ -49,6 +56,32 @@ def test_load_lora_adapter_succeeds_on_first_attempt():
     )
 
 
+def test_dynamo_load_lora_adapter_uses_existing_lora_engine_route():
+    mock_client = AsyncMock()
+    mock_response = MagicMock()
+    mock_response.raise_for_status = MagicMock()
+    mock_response.json.return_value = {"status": "success"}
+    mock_client.post.return_value = mock_response
+
+    asyncio.run(
+        DynamoAdminAPI().load_lora_adapter(
+            mock_client,
+            "test-lora",
+            "/test/path",
+            timeout=httpx.Timeout(connect=10.0, read=30.0, write=60.0, pool=10.0),
+        )
+    )
+
+    mock_client.post.assert_called_once_with(
+        "/engine/load_lora",
+        json={
+            "lora_name": "test-lora",
+            "source": {"uri": Path("/test/path").absolute().as_uri()},
+        },
+        timeout=httpx.Timeout(connect=10.0, read=30.0, write=60.0, pool=10.0),
+    )
+
+
 def test_setup_clients_assigns_renderer_and_dp_rank_headers():
     from renderers import Qwen3VLRendererConfig
 
@@ -117,3 +150,70 @@ def test_setup_clients_preserves_chat_client_defaults():
             extra_headers_from_state={},
         )
     ]
+
+
+def test_dynamo_rl_discovery_base_urls_derive_from_base_url(monkeypatch):
+    monkeypatch.setenv("DYN_RL_PORT", "18001")
+    client_config = ClientConfig(
+        base_url=["http://frontend.local:8000/v1"],
+        backend="dynamo",
+    )
+
+    assert _dynamo_rl_discovery_base_urls(client_config) == ["http://frontend.local:18001"]
+
+
+def test_dynamo_rl_discovery_base_urls_honor_explicit_config():
+    client_config = ClientConfig(
+        base_url=["http://frontend.local:8000/v1"],
+        backend="dynamo",
+        rl_base_url=["http://frontend.local:8001/v1"],
+    )
+
+    assert _dynamo_rl_discovery_base_urls(client_config) == ["http://frontend.local:8001/v1"]
+
+
+def test_discover_dynamo_admin_base_urls_reads_workers(monkeypatch):
+    calls = []
+
+    class FakeResponse:
+        def raise_for_status(self):
+            pass
+
+        def json(self):
+            return {
+                "workers": [
+                    {"system_url": "http://worker-0:8081"},
+                    {"system_url": "http://worker-0:8081"},
+                    {"system_url": "http://worker-1:8081"},
+                    {},
+                ]
+            }
+
+    class FakeClient:
+        def __init__(self, *, base_url, headers, timeout):
+            calls.append(("init", base_url, headers, timeout))
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+        def get(self, path):
+            calls.append(("get", path))
+            return FakeResponse()
+
+    monkeypatch.setattr("prime_rl.utils.client.httpx.Client", FakeClient)
+
+    client_config = ClientConfig(
+        base_url=["http://frontend.local:8000/v1"],
+        backend="dynamo",
+        rl_base_url=["http://frontend.local:8001/v1"],
+    )
+
+    assert discover_dynamo_admin_base_urls(client_config) == [
+        "http://worker-0:8081",
+        "http://worker-1:8081",
+    ]
+    assert calls[0][0:2] == ("init", "http://frontend.local:8001")
+    assert ("get", "/v1/rl/workers") in calls

From a1e0c5f163d102b06eda1e6139979fbee9ba1658 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 02/19] feat(weight-transfer): NCCL broadcast + FP8/E8M0
 conversion for GB200 (qwen3_moe/glm_moe)

---
 src/prime_rl/inference/vllm/worker/nccl.py    |  79 ++++++++++--
 .../inference/vllm/worker/weight_transfer.py  | 119 +++++++++++++++++-
 .../models/glm4_moe/converting_glm4_moe.py    |  35 ++++--
 .../glm_moe_dsa/converting_glm_moe_dsa.py     |  28 +++--
 .../models/qwen3_moe/converting_qwen3_moe.py  | 103 +++++++++++++++
 .../models/qwen3_moe/modeling_qwen3_moe.py    |   8 ++
 src/prime_rl/trainer/rl/broadcast/nccl.py     |  12 ++
 7 files changed, 360 insertions(+), 24 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/nccl.py b/src/prime_rl/inference/vllm/worker/nccl.py
index 95d35b583c..1acca94dcd 100644
--- a/src/prime_rl/inference/vllm/worker/nccl.py
+++ b/src/prime_rl/inference/vllm/worker/nccl.py
@@ -139,16 +139,81 @@ def init_broadcaster(
             inference_world_size: Total number of inference GPUs across all servers.
         """
         self.quantize_in_weight_transfer = quantize_in_weight_transfer
-        # Use the worker's device index directly as the local rank.
-        # The previous dp_group-based computation broke in vLLM v1 multiprocess
-        # DP mode where each worker is a separate process with a singleton
-        # DP group (rank_in_group is always 0).
+
+        # =====================================================================
+        # PATCHED — DP+TP-aware rank computation (Option B: parallel_config
+        # dispatch). Overrides the image-baked nccl.py via ConfigMap subPath mount.
+        #
+        # The vLLM Worker exposes `self.rank` (executor-wide rank) and
+        # `self.device.index` (local GPU index 0..local_world_size-1). Neither
+        # alone correctly identifies the GPU's position across BOTH TP and DP
+        # multinode topologies:
+        #
+        #   - TP/PP multinode (one executor spans pods):
+        #       self.rank is unique 0..world_size-1
+        #       self.device.index repeats per node (0..3 on each node)
+        #       --> use self.rank
+        #
+        #   - DP-only (TP=1) — e.g. vLLM `--data-parallel-size N --tp 1`:
+        #       self.rank is 0 for every subprocess
+        #       self.device.index is the LOCAL gpu index (0..size_local-1); for a
+        #       MULTINODE-DP follower it does NOT equal the global DP rank
+        #       --> use parallel_config.data_parallel_rank (global DP rank)
+        #
+        # Dispatch on parallel_config (always available on a vLLM Worker via
+        # vllm_config or as a direct attribute, depending on vLLM version).
+        # =====================================================================
+        pc = None
+        for attr_path in ("vllm_config.parallel_config", "parallel_config"):
+            obj = self
+            try:
+                for part in attr_path.split("."):
+                    obj = getattr(obj, part)
+                pc = obj
+                break
+            except AttributeError:
+                continue
+
+        tp_size = getattr(pc, "tensor_parallel_size", 1) if pc is not None else 1
+        pp_size = getattr(pc, "pipeline_parallel_size", 1) if pc is not None else 1
+        worker_rank = getattr(self, "rank", 0)
         local_rank = self.device.index
-        global_rank_inference = rank_offset + local_rank
+        dp_rank = getattr(pc, "data_parallel_rank", None) if pc is not None else None
+
+        if tp_size > 1 or pp_size > 1:
+            # Executor spans multiple ranks (TP or PP); self.rank is canonical.
+            effective_rank = worker_rank
+            rank_source = "self.rank (TP/PP > 1)"
+        elif dp_rank is not None:
+            # Pure DP: use the GLOBAL data-parallel rank. Correct for BOTH single-node
+            # DP (dp_rank == device.index) AND MULTINODE DP, where a follower node's
+            # device.index is the LOCAL gpu index (0..size_local-1) and does NOT equal
+            # the global DP rank (e.g. follower DP ranks 4..7 have device.index 0..3).
+            # Using device.index there collides with the leader's ranks 0..3 and the
+            # NCCL broadcast group never forms (Bootstrap "rank N already checked in").
+            effective_rank = local_rank  # DGD per-pod-offset override (issue #7); proper fix = per-engine offsets in orch
+            rank_source = "self.device.index (DGD per-pod-offset, issue #7)"
+        else:
+            # Fallback (older vLLM lacking data_parallel_rank): single-node DP only,
+            # where device.index uniquely identifies the GPU within the pod (0..N-1).
+            effective_rank = local_rank
+            rank_source = "self.device.index (DP-only fallback)"
+
+        global_rank_inference = rank_offset + effective_rank
 
         logger.info(
-            f"Worker [local_rank={local_rank} rank_offset={rank_offset}] "
-            f"-> [global_rank={global_rank_inference} inference_world_size={inference_world_size}]"
+            "Worker [tp_size=%s pp_size=%s local_rank=%s worker_rank=%s "
+            "effective_rank=%s (%s) rank_offset=%s] "
+            "-> [global_rank=%s inference_world_size=%s] (dp-tp-aware-patch)",
+            tp_size,
+            pp_size,
+            local_rank,
+            worker_rank,
+            effective_rank,
+            rank_source,
+            rank_offset,
+            global_rank_inference,
+            inference_world_size,
         )
 
         self.nccl_broadcast_receiver = NCCLWeightBroadcastReceiver(
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 7b448f021a..efe6f0a7c9 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -84,12 +84,121 @@ def build_expert_map(model: Module) -> dict[str, torch.Tensor]:
     return source_indices_by_module
 
 
+def _try_e8m0_scale_conversion(
+    name: str,
+    received_scale: torch.Tensor,
+    param: torch.Tensor,
+    params: dict[str, torch.Tensor],
+    updated_weights: set[str],
+) -> bool:
+    """Convert a trainer-format blockwise FP8 scale to vLLM's E8M0 TMA layout.
+
+    When DeepGEMM E8M0 is active (GB200 default), vLLM stores weight_scale_inv
+    tensors in a TMA-aligned [N, K/512] packed UE8M0 layout.  The trainer sends
+    [N/128, K/128] float32 blockwise scales.
+
+    This delegates to vLLM's own unified helper
+    ``deepgemm_post_process_fp8_weight_block``, which is the *exact* call vLLM
+    invokes at initial model load (via ``DeepGemmFp8BlockScaledMMKernel`` for
+    dense linears and ``prepare_fp8_moe_layer_for_deepgemm`` for MoE).  The
+    helper:
+
+      1. Re-quantises the FP8 weight to UE8M0 (power-of-two) scales in-place
+         (``requant_weight_ue8m0_inplace``).
+      2. Repacks the scale tensor to the TMA-aligned ``[..., N, K/512]`` layout
+         (``transform_sf_into_required_layout`` with recipe ``(1, 128, 128)``).
+      3. Handles 2D (dense) vs 3D (MoE) dispatch internally.
+
+    Reusing vLLM's wrapper guarantees the broadcast path produces the same
+    in-memory layout as vLLM's own startup load path — drift-free by
+    construction.
+
+    Returns True on success, False if conversion is not applicable or fails.
+    The caller is responsible for raising shape-mismatch errors for False returns.
+
+    ORDERING REQUIREMENT: the corresponding FP8 weight param must have been
+    updated (copied from the received tensor) before this function is called,
+    i.e. the weight must appear before its weight_scale_inv in the state iterator.
+    """
+    # Only handle weight_scale_inv tensors.
+    if not name.endswith("weight_scale_inv"):
+        return False
+
+    # Derive the corresponding weight parameter name.
+    # e.g. "...qkv_proj.weight_scale_inv" -> "...qkv_proj.weight"
+    #      "...w13_weight_scale_inv"      -> "...w13_weight"
+    weight_name = name[: -len("weight_scale_inv")] + "weight"
+    if weight_name not in params:
+        return False
+
+    # The weight must have been updated in this broadcast pass so that
+    # requant_weight_ue8m0_inplace dequantises the *new* FP8 values.
+    if weight_name not in updated_weights:
+        logger.warning(
+            "E8M0 scale conversion for %s: weight %s was not updated in this pass "
+            "(scale arrived before weight). Skipping conversion.",
+            name,
+            weight_name,
+        )
+        return False
+
+    try:
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            deepgemm_post_process_fp8_weight_block,
+        )
+    except ImportError as exc:
+        logger.debug("E8M0 conversion unavailable: %s", exc)
+        return False
+
+    weight_param = params[weight_name]
+
+    try:
+        # Single unified call: re-quantises weight in-place with UE8M0 scales
+        # AND repacks the scale tensor to the TMA-aligned [..., N, K/512] layout
+        # that vLLM stores.  Handles 2D (dense) and 3D (MoE) dispatch internally.
+        # The first return value is the same underlying storage as weight_param
+        # (modified in-place); we only need the new scale tensor.
+        _wq, dg_scale = deepgemm_post_process_fp8_weight_block(
+            wq=weight_param,
+            ws=received_scale,
+            quant_block_shape=(128, 128),
+            use_e8m0=True,
+        )
+        param.copy_(dg_scale.view_as(param))
+        logger.debug("E8M0 scale conversion applied for %s", name)
+        return True
+
+    except Exception as exc:
+        logger.warning("E8M0 scale conversion failed for %s: %s", name, exc)
+        return False
+
+
 @torch.no_grad()
 def load_weights_kernel(model: Module, state_iter: Generator[tuple[str, torch.Tensor], None, None]) -> None:
-    """Load vLLM kernel-format tensors using in-place copy_ updates."""
+    """Load vLLM kernel-format tensors using in-place copy_ updates.
+
+    Handles the GB200 DeepGEMM E8M0 scale mismatch: the trainer broadcasts
+    blockwise [N/128, K/128] float32 scales, while vLLM with E8M0 enabled
+    stores weight_scale_inv in [N, K/512] TMA-aligned packed UE8M0 layout.
+    When is_deep_gemm_e8m0_used() is True, scale tensors with a shape mismatch
+    are automatically converted via requant_weight_ue8m0_inplace +
+    transform_sf_into_required_layout.
+    """
     params = dict(model.named_parameters())
     expert_source_indices = build_expert_map(model)
 
+    # Detect E8M0 once (cached call, cheap).
+    try:
+        from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
+        _use_e8m0 = is_deep_gemm_e8m0_used()
+    except ImportError:
+        _use_e8m0 = False
+
+    # Track which weight params have been updated in this call so we can
+    # safely call requant_weight_ue8m0_inplace on them (it must dequantise
+    # the *new* FP8 values, not the stale ones from the initial model load).
+    updated_weights: set[str] = set()
+
     loaded = 0
     skipped: list[str] = []
     shape_mismatches: list[str] = []
@@ -108,11 +217,19 @@ def load_weights_kernel(model: Module, state_iter: Generator[tuple[str, torch.Te
                 break
 
             if param.shape != tensor.shape:
+                # Attempt E8M0 scale conversion before declaring a mismatch.
+                if _use_e8m0 and _try_e8m0_scale_conversion(
+                    name, tensor, param, params, updated_weights
+                ):
+                    loaded += 1
+                    continue
+
                 shape_mismatches.append(f"{name}: param={list(param.shape)} != received={list(tensor.shape)}")
                 continue
 
         param.copy_(tensor)
         loaded += 1
+        updated_weights.add(name)
 
     if shape_mismatches:
         raise ValueError(f"Kernel weight transfer had {len(shape_mismatches)} shape mismatches: {shape_mismatches}")
diff --git a/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py
index 85025821a1..ea5cf3c6f8 100644
--- a/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py
+++ b/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py
@@ -7,6 +7,18 @@ def get_max_layer_num(state_dict: dict[str, Tensor]) -> int:
     return max(int(i.split(".")[2]) for i in state_dict.keys() if "model.layers." in i) + 1
 
 
+def _expert_indices(state_dict: dict[str, Tensor], layer_idx: int) -> list[int]:
+    prefix = f"model.layers.{layer_idx}.mlp.experts."
+    indices: set[int] = set()
+    for key in state_dict:
+        if not key.startswith(prefix):
+            continue
+        candidate = key[len(prefix) :].split(".", 1)[0]
+        if candidate.isdigit():
+            indices.add(int(candidate))
+    return sorted(indices)
+
+
 def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int):
     i = layer_idx
 
@@ -35,24 +47,29 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int):
         del state_dict[f"model.layers.{i}.mlp.experts.down_proj"]
     else:
         # Old per-expert format
-        num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3
+        expert_indices = _expert_indices(state_dict, i)
+        num_experts = len(expert_indices)
         if num_experts == 0:
             return
 
-        dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape
+        first_expert = expert_indices[0]
+        dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].shape
         w1 = torch.empty(
-            (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype
+            (num_experts, moe_dim, dim),
+            dtype=state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype,
         )  # Gate
         w2 = torch.empty(
-            (num_experts, dim, moe_dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype
+            (num_experts, dim, moe_dim),
+            dtype=state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype,
         )  # Down
         w3 = torch.empty(
-            (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype
+            (num_experts, moe_dim, dim),
+            dtype=state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype,
         )  # Up
-        for j in range(num_experts):
-            w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"])
-            w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"])
-            w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"])
+        for expert_pos, j in enumerate(expert_indices):
+            w1[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"])
+            w2[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"])
+            w3[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"])
 
             del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]
             del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py
index e8e5b311b8..4d31a3cabc 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py
@@ -13,6 +13,18 @@ def _is_moe_layer(state_dict: dict[str, Tensor], layer_idx: int) -> bool:
     return f"model.layers.{layer_idx}.mlp.gate.weight" in state_dict
 
 
+def _expert_indices(state_dict: dict[str, Tensor], layer_idx: int) -> list[int]:
+    prefix = f"model.layers.{layer_idx}.mlp.experts."
+    indices: set[int] = set()
+    for key in state_dict:
+        if not key.startswith(prefix):
+            continue
+        candidate = key[len(prefix) :].split(".", 1)[0]
+        if candidate.isdigit():
+            indices.add(int(candidate))
+    return sorted(indices)
+
+
 def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int):
     i = layer_idx
 
@@ -38,19 +50,21 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int):
         del state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"]
         del state_dict[f"model.layers.{i}.mlp.experts.down_proj"]
     else:
-        num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3
+        expert_indices = _expert_indices(state_dict, i)
+        num_experts = len(expert_indices)
         if num_experts == 0:
             return
 
-        dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape
-        dtype = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype
+        first_expert = expert_indices[0]
+        dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].shape
+        dtype = state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype
         w1 = torch.empty((num_experts, moe_dim, dim), dtype=dtype)
         w2 = torch.empty((num_experts, dim, moe_dim), dtype=dtype)
         w3 = torch.empty((num_experts, moe_dim, dim), dtype=dtype)
-        for j in range(num_experts):
-            w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"])
-            w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"])
-            w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"])
+        for expert_pos, j in enumerate(expert_indices):
+            w1[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"])
+            w2[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"])
+            w3[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"])
 
             del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]
             del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]
diff --git a/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py
index 48ccc5b941..05362de1d3 100644
--- a/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py
+++ b/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py
@@ -1,6 +1,8 @@
 import torch
 from torch import Tensor
 
+from prime_rl.trainer.models.fp8 import quantize_to_fp8_blockwise
+
 
 def get_max_layer_num(state_dict: dict[str, Tensor]) -> int:
     """Get the maximum number of layers in the model."""
@@ -98,3 +100,104 @@ def convert_tt_to_hf_moe(state_dict: dict[str, Tensor]):
     num_layers = get_max_layer_num(state_dict)
     for i in range(num_layers):
         convert_tt_layer_to_hf(state_dict, i)
+
+
+def _emit_weight(out: dict[str, Tensor], name: str, tensor: Tensor, quantize_fp8: bool) -> None:
+    """Emit a 2D projection weight, optionally FP8-quantized with block scales."""
+    if quantize_fp8:
+        fp8_weight, scale = quantize_to_fp8_blockwise(tensor)
+        out[name] = fp8_weight
+        out[name.removesuffix(".weight") + ".weight_scale_inv"] = scale
+    else:
+        out[name] = tensor
+
+
+def _emit_moe_experts(
+    out: dict[str, Tensor],
+    prefix: str,
+    w1: Tensor,
+    w2: Tensor,
+    w3: Tensor,
+    quantize_fp8: bool,
+) -> None:
+    """Emit fused MoE expert weights in vLLM FusedMoE W13 layout.
+
+    vLLM's DEEPGEMM FP8 MoE path keeps weights in `[gate; up]` (W13) order;
+    FLASHINFER paths flip to W31 during `process_weights_after_loading`, so
+    inference must select the DEEPGEMM backend (`VLLM_USE_DEEP_GEMM=1`) for
+    RL weight reloads to stay valid across broadcasts.
+    """
+    w13 = torch.cat([w1, w3], dim=1)
+    if not quantize_fp8:
+        out[f"{prefix}.mlp.experts.w13_weight"] = w13
+        out[f"{prefix}.mlp.experts.w2_weight"] = w2
+        return
+
+    num_experts = w1.shape[0]
+    w13_fp8: list[Tensor] = []
+    w13_scales: list[Tensor] = []
+    w2_fp8: list[Tensor] = []
+    w2_scales: list[Tensor] = []
+    for expert_idx in range(num_experts):
+        q13, s13 = quantize_to_fp8_blockwise(w13[expert_idx])
+        q2, s2 = quantize_to_fp8_blockwise(w2[expert_idx])
+        w13_fp8.append(q13)
+        w13_scales.append(s13)
+        w2_fp8.append(q2)
+        w2_scales.append(s2)
+
+    out[f"{prefix}.mlp.experts.w13_weight"] = torch.stack(w13_fp8)
+    out[f"{prefix}.mlp.experts.w13_weight_scale_inv"] = torch.stack(w13_scales)
+    out[f"{prefix}.mlp.experts.w2_weight"] = torch.stack(w2_fp8)
+    out[f"{prefix}.mlp.experts.w2_weight_scale_inv"] = torch.stack(w2_scales)
+
+
+def _is_dense_layer(state_dict: dict[str, Tensor], layer_idx: int) -> bool:
+    """Dense MLP layers have a fused gate_proj; MoE layers route through mlp.router."""
+    return f"model.layers.{layer_idx}.mlp.gate_proj.weight" in state_dict
+
+
+def convert_tt_layer_to_vllm_kernel(
+    state_dict: dict[str, Tensor],
+    layer_idx: int,
+    quantize_fp8: bool = False,
+) -> dict[str, Tensor]:
+    """Convert a single Qwen3MoE layer from PrimeRL format to vLLM kernel format."""
+    prefix = f"model.layers.{layer_idx}"
+    out: dict[str, Tensor] = {
+        f"{prefix}.input_layernorm.weight": state_dict[f"{prefix}.input_layernorm.weight"],
+        f"{prefix}.post_attention_layernorm.weight": state_dict[f"{prefix}.post_attention_layernorm.weight"],
+        f"{prefix}.self_attn.q_norm.weight": state_dict[f"{prefix}.self_attn.q_norm.weight"],
+        f"{prefix}.self_attn.k_norm.weight": state_dict[f"{prefix}.self_attn.k_norm.weight"],
+    }
+
+    qkv = torch.cat(
+        [
+            state_dict[f"{prefix}.self_attn.q_proj.weight"],
+            state_dict[f"{prefix}.self_attn.k_proj.weight"],
+            state_dict[f"{prefix}.self_attn.v_proj.weight"],
+        ],
+        dim=0,
+    )
+    _emit_weight(out, f"{prefix}.self_attn.qkv_proj.weight", qkv, quantize_fp8)
+    _emit_weight(out, f"{prefix}.self_attn.o_proj.weight", state_dict[f"{prefix}.self_attn.o_proj.weight"], quantize_fp8)
+
+    if _is_dense_layer(state_dict, layer_idx):
+        gate_up = torch.cat(
+            [state_dict[f"{prefix}.mlp.gate_proj.weight"], state_dict[f"{prefix}.mlp.up_proj.weight"]],
+            dim=0,
+        )
+        _emit_weight(out, f"{prefix}.mlp.gate_up_proj.weight", gate_up, quantize_fp8)
+        _emit_weight(out, f"{prefix}.mlp.down_proj.weight", state_dict[f"{prefix}.mlp.down_proj.weight"], quantize_fp8)
+    else:
+        out[f"{prefix}.mlp.gate.weight"] = state_dict[f"{prefix}.mlp.router.gate.weight"]
+        _emit_moe_experts(
+            out,
+            prefix,
+            state_dict[f"{prefix}.mlp.experts.w1"],
+            state_dict[f"{prefix}.mlp.experts.w2"],
+            state_dict[f"{prefix}.mlp.experts.w3"],
+            quantize_fp8,
+        )
+
+    return out
diff --git a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py
index 25d83f9436..59c18efcc6 100644
--- a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py
+++ b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py
@@ -38,6 +38,7 @@
     convert_hf_layer_to_tt,
     convert_hf_to_tt_moe,
     convert_tt_layer_to_hf,
+    convert_tt_layer_to_vllm_kernel,
     convert_tt_to_hf_moe,
 )
 from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids
@@ -170,6 +171,13 @@ def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -
         convert_hf_layer_to_tt(state_dict, layer_idx)
         return state_dict
 
+    @classmethod
+    def convert_layer_to_vllm_kernel(
+        cls, state_dict: dict[str, Tensor], layer_idx: int, quantize_fp8: bool = False
+    ) -> dict[str, Tensor]:
+        """Convert a single layer's weights from PrimeRL format to vLLM FusedMoE/linear kernel format."""
+        return convert_tt_layer_to_vllm_kernel(state_dict, layer_idx, quantize_fp8=quantize_fp8)
+
 
 @auto_docstring
 class Qwen3MoeModel(Qwen3MoePreTrainedModel):
diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py
index 13a887308a..8412e5dd80 100644
--- a/src/prime_rl/trainer/rl/broadcast/nccl.py
+++ b/src/prime_rl/trainer/rl/broadcast/nccl.py
@@ -4,6 +4,7 @@
 from typing import Callable, Generator, cast
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch import Tensor
 from torch.distributed.tensor import DTensor
@@ -156,10 +157,21 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None:
             preprocess_fn = preprocess_layer_checkpoint
 
         for layer_id, layer_state_dict in filter_state_dict_by_layers(state_dict, num_layers, layer_prefix):
+            # _resolve_dtensors triggers DTensor.full_tensor() all_gathers on ALL trainer
+            # ranks (FSDP/EP collectives). Every rank must participate.
             layer_state_dict = self._resolve_dtensors(layer_state_dict)
             layer_state_dict = preprocess_fn(model, layer_state_dict, layer_id)
             if self.world.is_master:
                 broadcast_state_dict(layer_state_dict, self.communicator)
+                # Ensure the inference NCCL sends complete on the master's CUDA stream
+                # before we release non-master ranks to start the next layer's all_gather.
+                torch.cuda.synchronize()
+            # Barrier: prevents ranks 1-N from advancing to the next layer's
+            # _resolve_dtensors (and enqueuing the next FSDP/EP all_gather) before
+            # rank 0 finishes broadcasting the current layer to inference workers.
+            # Without this, ranks 1-N enqueue unmatched all_gathers and the NCCL
+            # watchdog fires after 600 s, killing all trainer processes.
+            dist.barrier()
 
     def _resolve_dtensors(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
         for key, value in list(state_dict.items()):

From 685c13a610dedf86ddc97c4cd9e5ed9b54271c29 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 03/19] feat(broadcast): NFS-safe filesystem broadcast +
 weight_broadcast.keep_recent for LoRA

---
 .../src/prime_rl/configs/trainer.py           | 31 +++++++++++---
 .../trainer/rl/broadcast/filesystem.py        | 40 ++++++++++++++++---
 src/prime_rl/trainer/rl/train.py              |  9 ++++-
 src/prime_rl/trainer/utils.py                 | 24 +++++------
 4 files changed, 80 insertions(+), 24 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 00f4e07deb..1f740c2785 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -457,12 +457,25 @@ class BaseWeightBroadcastConfig(BaseConfig):
 
 class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
-
-    save_sharded: bool = True
-    """Save the weight checkpoint in sharded format."""
-
-    save_format: Literal["safetensors", "torch"] = "safetensors"
-    """Weight checkpoint serialization format."""
+    save_sharded: Annotated[bool, Field(description="Whether to save the weight checkpoint in sharded format.")] = True
+    save_format: Annotated[
+        Literal["safetensors", "torch"], Field(description="The format to save the weight checkpoint in.")
+    ] = "safetensors"
+    keep_recent: Annotated[
+        int | None,
+        Field(
+            ge=0,
+            description=(
+                "Number of recent broadcast step directories to keep on disk. At trainer step N, "
+                "step_{N - keep_recent - 1} is deleted. If None (default), falls back to "
+                "trainer.max_async_level so retention matches the staleness window. Set this above "
+                "max_async_level for LoRA + filesystem broadcast: vLLM lazy-loads the adapter inside "
+                "_prepare_inputs, so a generate request admitted under an older adapter can page that "
+                "dir in *after* the orchestrator has moved on. A +1 or +2 buffer over max_async_level "
+                "is typical."
+            ),
+        ),
+    ] = None
 
 
 class NCCLWeightBroadcastConfig(BaseWeightBroadcastConfig):
@@ -626,6 +639,12 @@ def validate_lora_adapter_saving(self):
                 )
         return self
 
+    # NOTE: HEAD's validate_weight_broadcast_type (nccl => async level 1) and
+    # validate_broadcast_keep_recent (keep_recent >= max_async_level) were dropped
+    # in the rl-sdk-4 merge: main removed trainer.max_async_level (the staleness
+    # window is now orchestrator.max_off_policy_steps). Re-add equivalent guards at
+    # the RLConfig level (which sees both trainer + orchestrator) if needed.
+
     @model_validator(mode="after")
     def validate_opt_and_fsdp_offload(self):
         if self.optim.type == "muon" and self.model.fsdp_cpu_offload:
diff --git a/src/prime_rl/trainer/rl/broadcast/filesystem.py b/src/prime_rl/trainer/rl/broadcast/filesystem.py
index e8e2d68db9..35a8a0c89a 100644
--- a/src/prime_rl/trainer/rl/broadcast/filesystem.py
+++ b/src/prime_rl/trainer/rl/broadcast/filesystem.py
@@ -1,3 +1,4 @@
+import os
 import shutil
 import time
 from pathlib import Path
@@ -104,14 +105,43 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None:
             self.logger.debug(f"Weights broadcasted in {time.perf_counter() - start_time:.2f}s")
 
     def _notify_orchestrator(self, save_dir: Path):
-        """Notify the orchestrator that the weights have been broadcast by writing a 'STABLE' file to a shared filesystem."""
-        stable_file = save_dir / "STABLE"
-        stable_file.touch()
-
-    def maybe_clean(self, interval_to_keep: int | None):
+        """Notify the orchestrator that the weights have been broadcast by writing a 'STABLE' file to a shared filesystem.
+
+        On shared filesystems (NFS / CephFS / PVC-backed volumes), a write on one
+        node is not immediately visible on another. Without explicit flushing, the
+        orchestrator can see the STABLE sentinel before the adapter files
+        (``adapter_model.safetensors``, ``adapter_config.json``) are visible to
+        the inference worker — causing ``LoRAAdapterNotFoundError`` when
+        ``/v1/rl/load_lora_adapter`` is called. We fsync the adapter files and
+        the directory's dentries before touching STABLE, then fsync the
+        directory again so STABLE itself is durable.
+        """
+        for f in save_dir.iterdir():
+            if f.is_file():
+                _fsync_path(f)
+        _fsync_path(save_dir)
+        (save_dir / "STABLE").touch()
+        _fsync_path(save_dir)
+
+    def maybe_clean(self, retention: int, interval_to_keep: int | None):
         for idx in self.multi_run_manager.used_idxs:
             maybe_clean(
                 get_broadcast_dir(self.multi_run_manager.get_run_dir(idx)),
                 self.multi_run_manager.progress[idx].step,
+                retention,
                 interval_to_keep,
             )
+
+
+def _fsync_path(path: Path) -> None:
+    """Best-effort fsync of a file or directory. Silent on OSError because the
+    sync is an availability-of-data hint to the kernel, not a correctness gate
+    on its own — STABLE ordering is what matters."""
+    try:
+        fd = os.open(str(path), os.O_RDONLY)
+        try:
+            os.fsync(fd)
+        finally:
+            os.close(fd)
+    except OSError:
+        pass
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 83afa666dc..dcf2aecb98 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -279,7 +279,14 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                 # Clean up old broadcast directories (unless at ckpt interval if using filesystem weight broadcast)
                 if config.weight_broadcast.type == "filesystem":
                     interval_to_keep = config.ckpt and config.ckpt.interval
-                    weight_broadcast.maybe_clean(interval_to_keep)
+                    # Fallback retention when keep_recent unset: 2 (matches main's
+                    # original step-2 cleanup; trainer no longer carries max_async_level).
+                    retention = (
+                        config.weight_broadcast.keep_recent
+                        if config.weight_broadcast.keep_recent is not None
+                        else 2
+                    )
+                    weight_broadcast.maybe_clean(retention, interval_to_keep)
             else:
                 broadcast_weights_time = 0
                 # Usually the broadcast will set this. If broadcast is skipped, we need to reset this here.
diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py
index 9a8a86a526..10aebd0bce 100644
--- a/src/prime_rl/trainer/utils.py
+++ b/src/prime_rl/trainer/utils.py
@@ -450,18 +450,18 @@ def step(self):
         self.step_num += 1
 
 
-def maybe_clean(path: Path, step: int, interval_to_keep: int | None) -> None:
-    """Delete the broadcast dir from 2 trainer steps ago.
+def maybe_clean(path: Path, step: int, retention: int, interval_to_keep: int | None) -> None:
+    """Delete the broadcast step directory that falls outside the retention window.
 
-    With a 1-step async barrier, the orchestrator at trainer step ``step`` is still consuming the
-    ckpt from ``step - 1``; ``step - 2`` is therefore safe to remove unless it falls on a
-    checkpoint interval that we want to preserve.
+    At trainer step N, deletes step_{N - retention - 1}. Caller picks whether
+    `retention` is the orchestrator's staleness bound (max_async_level) or a
+    wider buffer (weight_broadcast.keep_recent).
     """
     logger = get_logger()
-    candidate_step = max(step - 2, 0)
-    candidate_path = get_step_path(path, candidate_step)
-    if interval_to_keep and candidate_step % interval_to_keep == 0:
-        logger.debug(f"Keeping path {candidate_path} (on ckpt interval)")
-        return
-    logger.debug(f"Removing path {candidate_path}")
-    shutil.rmtree(candidate_path, ignore_errors=True)
+    step = max(step - (retention + 1), 0)
+    candidate_path_to_delete = get_step_path(path, step)
+    keep = bool(interval_to_keep and step % interval_to_keep == 0)
+    logger.debug(f"Considering deleting path {candidate_path_to_delete}")
+    if not keep:
+        logger.debug(f"Removing path {candidate_path_to_delete}")
+        shutil.rmtree(candidate_path_to_delete, ignore_errors=True)

From 4f5ba0b0d96331fb460d8d5489120a05ccdb5e22 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 04/19] feat(orchestrator): dispatch compute_teacher_logprobs
 by renderer_transport (vllm + dynamo nvext)

---
 .../src/prime_rl/configs/orchestrator.py      |   7 +-
 src/prime_rl/orchestrator/utils.py            | 215 ++++++++++++++----
 2 files changed, 175 insertions(+), 47 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
index be5fe249f3..a3c28957b9 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
@@ -907,5 +907,10 @@ def resolve_env_config(self):
             if is_vllm:
                 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)
+                # DRA-TITO-PATCH: Dynamo strict-rejects `return_token_ids`.
+                # Token IDs come back via `nvext.engine_data.completion_token_ids`
+                # (PR #8119 channel) on the rl-sdk-2 backend, so the client does
+                # not need this vLLM-native field set. Stripped here at config
+                # validate time so it never reaches the inference server.
+                # env.sampling.extra_body.setdefault("return_token_ids", True)
         return self
diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py
index 5675ba3f34..592b6e7b76 100644
--- a/src/prime_rl/orchestrator/utils.py
+++ b/src/prime_rl/orchestrator/utils.py
@@ -4,16 +4,20 @@
 from concurrent.futures import ThreadPoolExecutor
 from itertools import cycle
 from pathlib import Path
+from typing import Any
 
 import orjson
+import pandas as pd
 import verifiers as vf
+from rich.console import Console
+from rich.table import Table
 from verifiers.utils.client_utils import setup_openai_client
 from verifiers.utils.save_utils import make_serializable
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
 from prime_rl.transport import TrainingSample
 from prime_rl.utils.client import setup_inference_pool
-from prime_rl.utils.logger import InterceptHandler, get_logger
+from prime_rl.utils.logger import InterceptHandler, format_time, get_logger
 from prime_rl.utils.utils import (
     get_broadcast_dir,
     get_ckpt_dir,
@@ -96,58 +100,177 @@ def set_default_executor(max_workers: int = 64) -> None:
     asyncio.get_event_loop().set_default_executor(ThreadPoolExecutor(max_workers=max_workers))
 
 
+def print_benchmark(history: dict[str, list[Any]]) -> None:
+    """
+    Print benchmark results as rich table. Shows formatted step time values.
+    First N rows show the per-step values, and the last row shows the mean,
+    std, min, and max values.
+    """
+    history.pop("step")
+    assert all(len(v) for v in history.values()), "All metrics must have logged the same number of steps"
+
+    # Turn metric history into pd.DataFrame
+    df = pd.DataFrame(dict(history.items()))
+    columns = {
+        "time/step": "Step Time",
+    }
+    df = df.rename(columns=columns)
+    df = df[list(columns.values())]
+    df = df.iloc[1:]  # Exclude first row
+
+    # Setup console
+    console = Console()
+    table = Table(title="Benchmark")
+
+    # Add columns
+    table.add_column("Step", justify="right")
+    for col in df.columns:
+        table.add_column(col, justify="center", style="magenta")
+
+    # Add formatted rows
+    formatted_df = pd.DataFrame(columns=df.columns)
+    formatted_df["Step Time"] = df["Step Time"].apply(format_time)
+    for step, row in formatted_df.iterrows():
+        table.add_row(*([str(step)] + [str(x) for x in row]))
+
+    # Separator
+    num_table_columns = 1 + len(df.columns)
+    table.add_row(*([""] * num_table_columns))
+
+    # Add row for formatted, aggregated statistics
+    mean_df = df.describe().loc[["mean", "std", "min", "max"], :]
+    formatted_mean_df = pd.DataFrame(columns=mean_df.columns)
+    formatted_mean_df["Step Time"] = mean_df["Step Time"].apply(format_time)
+    mean_row = ["Overall"] + formatted_mean_df.T.apply(
+        lambda row: f"{row['mean']} ± {row['std']} [{row['min']}, {row['max']}]", axis=1
+    ).tolist()
+    table.add_row(*mean_row)
+
+    # Display table
+    console.print(table)
+
+
+def _flatten_prompt_logprobs(raw: list[Any] | None) -> list[float]:
+    """Shared flattener used by both transports.
+
+    ``prompt_logprobs[i]`` is a ``{token_id: Logprob}`` dict for tokens the
+    engine could score, or ``None`` for the leading token which has no
+    preceding context. Flatten to ``list[float]`` with 0.0 in the unscored
+    slot. Accepts both vLLM's typed ``Logprob`` objects and dynamo's
+    ``PromptLogprobEntry`` dict shape (`{logprob, rank?, decoded_token?}`).
+    """
+    flat: list[float] = []
+    for entry in raw or []:
+        if not entry:
+            flat.append(0.0)
+            continue
+        first = next(iter(entry.values()))
+        lp = first.logprob if hasattr(first, "logprob") else first.get("logprob")
+        flat.append(float(lp) if lp is not None else 0.0)
+    return flat
+
+
+async def _compute_teacher_logprobs_vllm(
+    client_config: vf.ClientConfig, model_name: str, sample: TrainingSample
+) -> list[float]:
+    """Legacy path: prime-rl's vLLM sidecar ``/inference/v1/generate``."""
+    import httpx
+    from vllm.entrypoints.serve.disagg.protocol import GenerateResponse
+
+    client = setup_openai_client(client_config)
+    # Two escape hatches from ``AsyncOpenAI.post``:
+    #   1. URL — ``/inference/v1/generate`` is mounted at server root, not
+    #      under ``/v1``. Pass an absolute URL so the SDK's ``_prepare_url``
+    #      skips the base-url merge.
+    #   2. Parse — vLLM's ``GenerateResponse`` isn't an ``openai.BaseModel``.
+    #      Use ``cast_to=httpx.Response`` and validate the body ourselves.
+    base = str(client.base_url).rstrip("/").removesuffix("/v1")
+    http_response = await client.post(
+        f"{base}/inference/v1/generate",
+        cast_to=httpx.Response,
+        body={
+            "model": model_name,
+            "token_ids": list(sample.prompt_ids) + list(sample.completion_ids),
+            "sampling_params": {
+                "max_tokens": 1,
+                "temperature": 1.0,
+                "top_p": 1.0,
+                "prompt_logprobs": 1,
+            },
+        },
+    )
+    response = GenerateResponse.model_validate_json(http_response.content)
+    return _flatten_prompt_logprobs(response.prompt_logprobs)
+
+
+async def _compute_teacher_logprobs_dynamo(
+    client_config: vf.ClientConfig, model_name: str, sample: TrainingSample
+) -> list[float]:
+    """rl-sdk-2 path: dynamo via ``/v1/chat/completions`` with nvext envelope.
+
+    Wire shape (per plan.md A1+A3 — already implemented on rl-sdk-2):
+      - top-level ``prompt_logprobs: 1`` (CommonExt sampling param)
+      - ``nvext.token_data`` carries pre-tokenized prompt
+      - ``nvext.extra_fields = ["prompt_logprobs"]`` opts into the response
+        field; dynamo emits ``response.nvext.prompt_logprobs`` shaped as
+        ``[None | {token_id: {logprob, rank?, decoded_token?}}]``, which the
+        shared flattener consumes unchanged.
+
+    Required engine adapter support: vLLM worker must populate
+    ``LLMEngineOutput.prompt_logprobs`` when ``SamplingParams.prompt_logprobs``
+    is set. Without that (A10), the response payload is None.
+    """
+    client = setup_openai_client(client_config)
+    token_ids = list(sample.prompt_ids) + list(sample.completion_ids)
+    body = {
+        "model": model_name,
+        "messages": [],
+        "max_completion_tokens": 1,
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "prompt_logprobs": 1,
+        "nvext": {
+            "token_data": token_ids,
+            "extra_fields": ["prompt_logprobs"],
+        },
+    }
+    # Dynamo's response is just a standard chat-completion JSON with an extra
+    # ``nvext`` field. Use ``cast_to=httpx.Response`` so we can read the raw
+    # body and pluck ``nvext.prompt_logprobs`` ourselves — the OpenAI SDK
+    # response models drop unknown fields.
+    import httpx as _httpx
+
+    http_response = await client.post(
+        "/chat/completions",
+        cast_to=_httpx.Response,
+        body=body,
+    )
+    payload = http_response.json()
+    nvext_resp = (payload or {}).get("nvext") or {}
+    raw = nvext_resp.get("prompt_logprobs")
+    return _flatten_prompt_logprobs(raw)
+
+
 async def compute_teacher_logprobs(
     clients: list[vf.ClientConfig],
     model_name: str,
     samples: list[TrainingSample],
 ) -> list[list[float]]:
-    """Compute teacher model logprobs for a batch of training samples via prefill."""
-    import httpx
-    from vllm.entrypoints.serve.disagg.protocol import GenerateResponse
+    """Compute teacher model logprobs for a batch of training samples via prefill.
+
+    Dispatches to the vLLM-sidecar or dynamo-nvext path based on the
+    per-client ``renderer_transport``:
+
+      - ``prime_vllm_generate``  (default): POST ``/inference/v1/generate``
+      - ``dynamo_chat_nvext``               : POST ``/v1/chat/completions`` with nvext
+
+    Both flatten to ``list[float]`` via the shared helper.
+    """
 
     async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]:
-        client = setup_openai_client(client_config)
-
-        # Two escape hatches from ``AsyncOpenAI.post``:
-        #   1. URL — ``/inference/v1/generate`` is mounted at server root, not
-        #      under ``/v1``. Pass an absolute URL so the SDK's
-        #      ``_prepare_url`` skips the base-url merge (it short-circuits
-        #      when the path passes ``httpx.URL.is_relative_url`` as False).
-        #   2. Parse — vLLM's ``GenerateResponse`` is a plain
-        #      ``pydantic.BaseModel`` and the SDK's parse layer rejects any
-        #      ``cast_to`` that doesn't subclass ``openai.BaseModel``. Use
-        #      ``cast_to=httpx.Response`` so the SDK still builds the request
-        #      (preserving ``auth_headers``, retries, timeouts, idempotency
-        #      keys) and just hands us the raw response to validate ourselves.
-        base = str(client.base_url).rstrip("/").removesuffix("/v1")
-        http_response = await client.post(
-            f"{base}/inference/v1/generate",
-            cast_to=httpx.Response,
-            body={
-                "model": model_name,
-                "token_ids": list(sample.prompt_ids) + list(sample.completion_ids),
-                "sampling_params": {
-                    "max_tokens": 1,
-                    "temperature": 1.0,
-                    "top_p": 1.0,
-                    "prompt_logprobs": 1,
-                },
-            },
-        )
-        response = GenerateResponse.model_validate_json(http_response.content)
-        # ``prompt_logprobs[i]`` is a ``{token_id: Logprob}`` dict for tokens
-        # the engine could score, or ``None`` for the leading token which has
-        # no preceding context. Flatten to ``list[float]`` with 0.0 in the
-        # unscored slot.
-        flat: list[float] = []
-        for entry in response.prompt_logprobs or []:
-            if not entry:
-                flat.append(0.0)
-                continue
-            first = next(iter(entry.values()))
-            lp = first.logprob if hasattr(first, "logprob") else first.get("logprob")
-            flat.append(float(lp) if lp is not None else 0.0)
-        return flat
+        if client_config.renderer_transport == "dynamo_chat_nvext":
+            return await _compute_teacher_logprobs_dynamo(client_config, model_name, sample)
+        return await _compute_teacher_logprobs_vllm(client_config, model_name, sample)
 
     return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)])
 

From 7e4bc48cb0f68e095f35f281e3eaa12430dadd74 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 05/19] feat(inference): vLLM 0.22 patches (fp32 lm-head, int64
 silu_mul_quant, padded scrub)

---
 src/prime_rl/_compat.py           |   4 +-
 src/prime_rl/inference/patches.py | 209 ++++++++++++++++++++++++++++++
 2 files changed, 212 insertions(+), 1 deletion(-)

diff --git a/src/prime_rl/_compat.py b/src/prime_rl/_compat.py
index 268c38376b..1ab251a777 100644
--- a/src/prime_rl/_compat.py
+++ b/src/prime_rl/_compat.py
@@ -35,7 +35,9 @@
 # ---------------------------------------------------------------------------
 try:
     import transformers.integrations.hub_kernels as _hub_kernels
-except ImportError:
+    # AttributeError: module exists in some builds but lazy_load_kernel was added in 5.5
+    _ = _hub_kernels.lazy_load_kernel
+except (ImportError, AttributeError):
     _hub_kernels = None  # transformers < 5.5, no patch needed
 
 if _hub_kernels is not None:
diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py
index 8190f210b0..7d85153c47 100644
--- a/src/prime_rl/inference/patches.py
+++ b/src/prime_rl/inference/patches.py
@@ -18,6 +18,9 @@ def transformers_v5_compat():
     _patch_qwen35_lora()
     _patch_lora_key_prefix()
     monkey_patch_deep_gemm_silu_mul_quant_int64()
+    monkey_patch_deep_gemm_silu_mul_quant_packed_int64()
+    # monkey_patch_dp_engine_core_pause_resume_deadlock()  # DISABLED: use vLLM PR #39366-native two-phase pause (avoid double pause/resume fix)
+    monkey_patch_fp32_lm_head()
     monkey_patch_vllm_padded_input_scrub()
     monkey_patch_return_routed_experts_with_nixl_connector()
 
@@ -207,6 +210,207 @@ def monkey_patch_deep_gemm_silu_mul_quant_int64():
     logger.warning("Enabled int64-addressing Triton patch for vLLM DeepGEMM SiLU/mul FP8 quant.")
 
 
+@triton.jit
+def _silu_mul_quant_fp8_packed_kernel_int64(
+    input_ptr,
+    output_q_ptr,
+    output_scale_ptr,
+    M: tl.int64,
+    input_stride_m: tl.int64,
+    output_q_stride_m: tl.int64,
+    output_scale_stride_k: tl.int64,
+    clamp_limit,
+    N: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    fp8_min: tl.constexpr,
+    fp8_max: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    HAS_CLAMP: tl.constexpr,
+):
+    """int64-addressing variant of vLLM's _silu_mul_quant_fp8_packed_kernel.
+
+    Mirrors `vllm/model_executor/layers/quantization/utils/fp8_utils.py:152`
+    but casts row/column offsets to tl.int64 so the address arithmetic in
+    `base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m` does not
+    overflow when M * input_stride_m exceeds 2**31 (e.g. Qwen3-235B profile_run
+    on EP=2 × DP=4 layouts, 128 experts × 6144 fused MoE dim).
+    """
+    N_2: tl.constexpr = N // 2
+
+    # Grid layout: dim0 = M-blocks (large, up to 2**31-1 on Blackwell),
+    # dim1 = packed groups (small, ~3-6). The upstream kernel placed the
+    # large M dim on grid Y, which is capped at 65,535 and produced
+    # `Triton Error [CUDA]: invalid argument` on Qwen3-235B profile_run.
+    pid_m = tl.program_id(0)
+    pid_pack = tl.program_id(1)
+    m_offset = (pid_m * BLOCK_M).to(tl.int64)
+
+    if m_offset >= M:
+        return
+
+    offs_m = tl.arange(0, BLOCK_M).to(tl.int64)
+    offs_n = tl.arange(0, GROUP_SIZE).to(tl.int64)
+    row_mask = (m_offset + offs_m) < M
+
+    base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m
+    base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m
+
+    packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32)
+
+    for pack_idx in tl.static_range(4):
+        group_id = pid_pack * 4 + pack_idx
+
+        if group_id < NUM_GROUPS:
+            n_offset = (group_id * GROUP_SIZE).to(tl.int64)
+
+            act_ptrs = input_ptr + base_row_offset + n_offset + offs_n[None, :]
+            act_in = tl.load(act_ptrs, mask=row_mask[:, None], other=0.0)
+
+            mul_ptrs = act_ptrs + N_2
+            mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0)
+
+            act_f32 = act_in.to(tl.float32)
+            mul_f32 = mul_in.to(tl.float32)
+
+            if HAS_CLAMP:
+                act_f32 = tl.minimum(act_f32, clamp_limit)
+                mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit)
+
+            y = (act_f32 / (1.0 + tl.exp(-act_f32))) * mul_f32
+            # Round through bf16 to match unfused precision path
+            y = y.to(tl.bfloat16).to(tl.float32)
+
+            absmax = tl.max(tl.abs(y), axis=1)
+
+            scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
+            exponent = tl.ceil(tl.log2(scale_raw))
+            scale = tl.math.exp2(exponent)
+
+            y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max)
+
+            out_q_ptrs = output_q_ptr + base_out_offset + n_offset + offs_n[None, :]
+            tl.store(
+                out_q_ptrs,
+                y_q.to(output_q_ptr.dtype.element_ty),
+                mask=row_mask[:, None],
+            )
+
+            exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
+            packed_scale = packed_scale | (exponent_biased << (pack_idx * 8))
+
+    scale_ptrs = output_scale_ptr + pid_pack.to(tl.int64) * output_scale_stride_k + m_offset + offs_m
+    tl.store(scale_ptrs, packed_scale, mask=row_mask)
+
+
+def silu_mul_quant_fp8_packed_triton_int64(
+    input: torch.Tensor,
+    group_size: int = 128,
+    output_q: torch.Tensor | None = None,
+    clamp_limit: float | None = None,
+):
+    """int64-addressing variant of vLLM's silu_mul_quant_fp8_packed_triton.
+
+    Same semantics and return shape as the upstream function, but launches the
+    int64 kernel above. Required for Qwen3-235B FP8 inference with DeepGEMM
+    enabled — the upstream packed kernel overflows on row offsets at profile_run
+    shapes (see issues.md Issue 7 in work/bis-dev/may-26/01-qwen-235b-whiteffiber).
+    """
+    assert input.dim() == 2
+    assert input.is_contiguous()
+
+    M, N = input.shape
+    N_2 = N // 2
+
+    assert N_2 % group_size == 0
+
+    fp8_dtype = torch.float8_e4m3fn
+    finfo = torch.finfo(fp8_dtype)
+    fp8_min, fp8_max = finfo.min, finfo.max
+
+    num_groups_per_row = N_2 // group_size
+    num_packed_groups = (num_groups_per_row + 3) // 4
+    tma_aligned_M = ((M + 3) // 4) * 4
+
+    if output_q is None:
+        output_q = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device)
+
+    output_scale_packed = torch.zeros(
+        (num_packed_groups, tma_aligned_M),
+        dtype=torch.int32,
+        device=input.device,
+    ).T[:M, :]
+
+    BLOCK_M = 8
+    # gridX gets the large M dim (Blackwell cap is 2**31-1); gridY gets
+    # the small packed-group dim (Blackwell cap is 65,535). Swapping vs.
+    # upstream avoids "invalid argument" when (M+7)//8 > 65,535
+    # (Qwen3-235B with 128 experts on EP=2 trips this in profile_run).
+    grid = ((M + BLOCK_M - 1) // BLOCK_M, num_packed_groups)
+
+    num_warps = max(4, group_size // 32)
+    num_stages = 2
+
+    has_clamp = clamp_limit is not None
+    _silu_mul_quant_fp8_packed_kernel_int64[grid](
+        input,
+        output_q,
+        output_scale_packed,
+        M,
+        input.stride(0),
+        output_q.stride(0),
+        output_scale_packed.stride(1),
+        clamp_limit if has_clamp else 0.0,
+        N=N,
+        NUM_GROUPS=num_groups_per_row,
+        fp8_min=fp8_min,
+        fp8_max=fp8_max,
+        GROUP_SIZE=group_size,
+        BLOCK_M=BLOCK_M,
+        HAS_CLAMP=has_clamp,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+
+    return output_q, output_scale_packed
+
+
+def monkey_patch_deep_gemm_silu_mul_quant_packed_int64():
+    """Replace vLLM's silu_mul_quant_fp8_packed_triton with our int64 variant.
+
+    The upstream packed kernel in vLLM 0.21.0 uses int32 arithmetic for row
+    offsets (`(m_offset + offs_m[:, None]) * input_stride_m`). For Qwen3-235B
+    FP8 inference with `VLLM_USE_DEEP_GEMM=1`, profile_run hits M*stride > 2**31
+    on the 128-expert layout and Triton emits CUDA "invalid argument".
+
+    We patch:
+      1. fp8_utils.silu_mul_quant_fp8_packed_triton (the public wrapper)
+      2. deep_gemm_moe.fused_silu_mul_fp8_quant_packed (the alias inside the
+         DeepGEMM MoE experts module — imported by name at module load)
+    """
+    import sys
+
+    from vllm.logger import init_logger
+    from vllm.model_executor.layers.quantization.utils import fp8_utils
+
+    logger = init_logger(__name__)
+
+    fp8_utils.silu_mul_quant_fp8_packed_triton = silu_mul_quant_fp8_packed_triton_int64
+
+    deep_gemm_moe_module = sys.modules.get("vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe")
+    if deep_gemm_moe_module is not None:
+        # deep_gemm_moe.py re-exports the wrapper under the name
+        # `fused_silu_mul_fp8_quant_packed` (see line 208 in the upstream module).
+        if hasattr(deep_gemm_moe_module, "fused_silu_mul_fp8_quant_packed"):
+            deep_gemm_moe_module.fused_silu_mul_fp8_quant_packed = silu_mul_quant_fp8_packed_triton_int64
+        if hasattr(deep_gemm_moe_module, "silu_mul_quant_fp8_packed_triton"):
+            deep_gemm_moe_module.silu_mul_quant_fp8_packed_triton = silu_mul_quant_fp8_packed_triton_int64
+
+    logger.warning(
+        "Enabled int64-addressing Triton patch for vLLM DeepGEMM packed SiLU/mul FP8 quant."
+    )
+
+
 def _patch_qwen35_lora():
     """Fix Qwen3.5 LoRA: align packed_modules_mapping with output_sizes.
 
@@ -803,6 +1007,10 @@ def monkey_patch_fp32_lm_head():
 
     logger = init_logger(__name__)
 
+    if getattr(LogitsProcessor, "_prime_rl_fp32_lm_head_patch_installed", False):
+        logger.debug("fp32 lm_head patch already installed; skipping.")
+        return
+
     _original_init = LogitsProcessor.__init__
     _original_get_logits = LogitsProcessor._get_logits
 
@@ -835,4 +1043,5 @@ def _patched_get_logits(self, hidden_states, lm_head, embedding_bias):
 
     LogitsProcessor.__init__ = _patched_init
     LogitsProcessor._get_logits = _patched_get_logits
+    LogitsProcessor._prime_rl_fp32_lm_head_patch_installed = True
     logger.info("Installed fp32 lm_head patch (native out_dtype=fp32 mm).")

From 9f2d880b57c36330a953aff6646dfb1a2cbadb66 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 06/19] build(deps): point verifiers/renderers submodules at
 biswapanda rl-sdk-4 forks (TITO)

---
 .gitmodules                | 4 ++--
 deps/pydantic-config       | 2 +-
 deps/renderers             | 2 +-
 deps/research-environments | 2 +-
 deps/verifiers             | 2 +-
 5 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.gitmodules b/.gitmodules
index 2041f460ee..d139f9810f 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,9 +1,9 @@
 [submodule "verifiers"]
 	path = deps/verifiers
-	url = git@github.com:PrimeIntellect-ai/verifiers.git
+	url = git@github.com:biswapanda/verifiers.git
 [submodule "renderers"]
 	path = deps/renderers
-	url = git@github.com:PrimeIntellect-ai/renderers.git
+	url = git@github.com:biswapanda/renderers.git
 [submodule "research-environments"]
 	path = deps/research-environments
 	url = git@github.com:PrimeIntellect-ai/research-environments.git
diff --git a/deps/pydantic-config b/deps/pydantic-config
index 896ade4e69..8fc28bf218 160000
--- a/deps/pydantic-config
+++ b/deps/pydantic-config
@@ -1 +1 @@
-Subproject commit 896ade4e69d8d8dff2d4b0a431b7e1c7c12d638f
+Subproject commit 8fc28bf21828b20ddd6812d27b240daec484dc86
diff --git a/deps/renderers b/deps/renderers
index e6dba5ad6c..6a215742b0 160000
--- a/deps/renderers
+++ b/deps/renderers
@@ -1 +1 @@
-Subproject commit e6dba5ad6c50ca83d4ffa462145037082542e52a
+Subproject commit 6a215742b0d89a1d4e79e8d217efa4909ed8b0a5
diff --git a/deps/research-environments b/deps/research-environments
index c752781984..3c58236993 160000
--- a/deps/research-environments
+++ b/deps/research-environments
@@ -1 +1 @@
-Subproject commit c752781984c1b4fbb0a3d7f4aac1e7ed67cc749e
+Subproject commit 3c58236993eef79897ec9c5552d36aa591d178cc
diff --git a/deps/verifiers b/deps/verifiers
index 05c66c2358..ee3482aebf 160000
--- a/deps/verifiers
+++ b/deps/verifiers
@@ -1 +1 @@
-Subproject commit 05c66c235875d785754f2b7078db0e7deeddbeae
+Subproject commit ee3482aebfaf35e47ec73a55db9276364d63e1cd

From a03a55934ab402eaca93bdce889dbb6ceafefe2b Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 07/19] build(image): Dockerfile.cuda.runtime (vLLM 0.22, COPY
 deps, DeepGEMM) + Dockerfile.dynamo

---
 Dockerfile.cuda.runtime              | 189 ++++++++++++++++++++++++
 Dockerfile.dynamo                    |  79 ++++++++++
 scripts/docker-arm64-post-install.sh |   4 +
 scripts/vllm-pr-39366.patch          | 206 +++++++++++++++++++++++++++
 4 files changed, 478 insertions(+)
 create mode 100644 Dockerfile.cuda.runtime
 create mode 100644 Dockerfile.dynamo
 create mode 100644 scripts/vllm-pr-39366.patch

diff --git a/Dockerfile.cuda.runtime b/Dockerfile.cuda.runtime
new file mode 100644
index 0000000000..38499fd90a
--- /dev/null
+++ b/Dockerfile.cuda.runtime
@@ -0,0 +1,189 @@
+# Multi-stage Dockerfile for prime-rl with NVRTC support on GB200 (sm_100a).
+#
+# WHY THIS EXISTS (separate from Dockerfile.cuda):
+#   The original Dockerfile.cuda uses `python:3.12-slim` as the runtime base.
+#   That image has no CUDA toolkit, so tilelang's JIT path (which compiles the
+#   sparse-MLA kernels at runtime via NVRTC) fails with:
+#
+#       atomic.h(7): catastrophic error: cannot open source file "cuda/atomic"
+#
+#   `cuda/atomic` is a libcudacxx (CCCL) header. It is shipped only by the
+#   CUDA dev/devel toolkit, not by the `nvidia-cuda-*` pip wheels. Without it
+#   NVRTC cannot compile any kernel that pulls in tilelang's atomic.h.
+#
+#   This Dockerfile uses NVIDIA's cuda-dl-base (devel) image for both stages,
+#   so the runtime image carries the libcudacxx / CCCL headers tilelang needs
+#   and the wheels in /app/.venv keep their CUDA 12.x ABI.
+#
+# BASE IMAGE:
+#   nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04
+#     - matches dynamo/container/context.yaml `prime-rl.cuda12.9` entry
+#     - ships Python 3.12 by default (Ubuntu 24.04)
+#     - includes libcudacxx, CCCL, cuDNN, nvcc, and full CUDA dev headers
+#     - forward-compatible with the CUDA 12.8 wheels pinned in uv.lock
+#
+# USAGE:
+#   docker buildx build --platform linux/arm64 \
+#       --build-arg TARGETARCH=arm64 \
+#       -f Dockerfile.cuda.runtime -t <tag> .
+
+ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base
+ARG BASE_IMAGE_TAG=25.06-cuda12.9-devel-ubuntu24.04
+
+############################
+##### Build stage ##########
+############################
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS builder
+LABEL maintainer="prime intellect"
+LABEL repository="prime-rl"
+
+# Set en_US.UTF-8 locale by default
+RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment
+
+# CUDA_HOME / PATH from base image are already correct (/usr/local/cuda), but
+# pin them explicitly so downstream tooling (tilelang, flash-attn) sees them.
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=$PATH:/usr/local/cuda/bin
+
+# Install build tooling.
+ARG DEBIAN_FRONTEND=noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
+ENV TZ=Etc/UTC
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    curl \
+    sudo \
+    git \
+    ninja-build \
+    && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+# Install uv.
+ADD https://astral.sh/uv/install.sh /uv-installer.sh
+RUN INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR="/usr/local/bin" sh /uv-installer.sh && rm /uv-installer.sh
+ENV PATH="/usr/local/bin:$PATH"
+ENV UV_PYTHON_INSTALL_DIR="/usr/local/share/uv/python"
+ENV UV_CACHE_DIR="/usr/local/share/uv/cache"
+
+# Install Python dependencies (gradual copies help with caching).
+WORKDIR /app
+
+ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy
+
+COPY pyproject.toml /app/pyproject.toml
+COPY uv.lock /app/uv.lock
+COPY README.md /app/README.md
+COPY src/ /app/src/
+COPY packages/ /app/packages/
+COPY deps/ /app/deps/
+COPY configs /app/configs
+COPY examples /app/examples
+COPY benchmarks/scripts /app/benchmarks/scripts
+
+RUN --mount=type=cache,target=/app/.cache/uv \
+    uv sync --extra flash-attn --extra flash-attn-3 --extra flash-attn-cute --extra envs --extra gpt-oss --group mamba-ssm --locked --no-dev
+
+# arm64: build flash-attn + DeepGEMM from source.
+ARG TARGETARCH
+COPY scripts/docker-arm64-post-install.sh /app/scripts/docker-arm64-post-install.sh
+COPY scripts/install_deep_gemm.sh /app/scripts/install_deep_gemm.sh
+RUN if [ "$TARGETARCH" = "arm64" ]; then /app/scripts/docker-arm64-post-install.sh; fi
+
+# vLLM PR #39366 (two-phase DP pause) is native in vLLM 0.22 — no patch needed
+# (the rl-sdk-4 merge bumped vLLM 0.20.2 -> 0.22 and dropped this patch).
+
+############################
+##### Runtime stage ########
+############################
+# Same image so libcudacxx, CCCL headers, and the system Python 3.12 are all
+# present at runtime — the original Dockerfile.cuda switched to python:3.12-slim
+# here and lost the CUDA dev headers, which broke tilelang's NVRTC backend.
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG}
+
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=/usr/local/cuda/bin:$PATH
+
+ARG DEBIAN_FRONTEND=noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
+ENV TZ=Etc/UTC
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    wget \
+    clang \
+    tmux \
+    iperf \
+    openssh-server \
+    git \
+    git-lfs \
+    gpg \
+    sudo \
+    iputils-ping \
+    net-tools \
+    curl \
+    vim \
+    libibverbs1 \
+    ibverbs-providers \
+    python3.12 \
+    python3.12-venv \
+    && apt-get clean autoclean \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+# Ensure `python` / `python3` point at 3.12 (ubuntu24.04 ships 3.12 already,
+# but the symlinks aren't created by default).
+RUN ln -sf /usr/bin/python3.12 /usr/local/bin/python \
+ && ln -sf /usr/bin/python3.12 /usr/local/bin/python3 \
+ && ln -sf /usr/bin/python3.12 /usr/local/bin/python3.12
+
+ARG USER_ID=1000
+ARG GROUP_ID=1000
+# Ubuntu 24.04 ships a default `ubuntu` user at uid 1000; remove it so the
+# explicit appuser keeps uid/gid 1000 (matches Dockerfile.cuda + k8s manifests).
+RUN userdel -r ubuntu 2>/dev/null || true \
+ && groupadd --gid $GROUP_ID appuser \
+ && useradd --uid $USER_ID --gid appuser --create-home --shell /bin/bash appuser \
+ && usermod -aG sudo appuser \
+ && echo "appuser ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers
+
+# Install uv for development use.
+ADD https://astral.sh/uv/install.sh /uv-installer.sh
+RUN INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR="/usr/local/bin" sh /uv-installer.sh && rm /uv-installer.sh
+
+USER appuser
+ENV PATH="/usr/local/bin:$PATH"
+WORKDIR /app
+# Copy the application + venv from the builder.
+COPY --from=builder --chown=appuser:appuser /app /app
+
+# Copy and set up entrypoint script.
+COPY --chown=appuser:appuser scripts/docker-entrypoint.sh /app/docker-entrypoint.sh
+RUN chmod +x /app/docker-entrypoint.sh
+
+# Repoint venv Python symlinks at the runtime-stage interpreter (the builder
+# used a uv-managed Python that does not exist here).
+RUN rm /app/.venv/bin/python && ln -s /usr/bin/python3.12 /app/.venv/bin/python
+RUN rm /app/.venv/bin/python3 && ln -s /usr/bin/python3.12 /app/.venv/bin/python3
+RUN rm /app/.venv/bin/python3.12 && ln -s /usr/bin/python3.12 /app/.venv/bin/python3.12
+
+# python3.12-dev: Python.h headers for vLLM's Triton CudaUtils JIT-compile at runtime.
+# Required by vLLM serving (dynamo.vllm / inference); absent from the cuda-dl-base
+# runtime, which crashed inference with `fatal error: Python.h: No such file or directory`.
+USER root
+RUN apt-get update && apt-get install -y --no-install-recommends python3.12-dev \
+ && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+USER appuser
+
+# Place executables in the environment at the front of the path.
+ENV PATH="/app/.venv/bin:$PATH"
+
+# HuggingFace Hub timeouts (defaults are 10s which causes issues on slow networks).
+ENV HF_HUB_ETAG_TIMEOUT=500
+ENV HF_HUB_DOWNLOAD_TIMEOUT=300
+
+# Enable FP8 grouped-GEMM kernels in vLLM MoE layers (requires DeepGEMM, built
+# during the arm64 post-install step above).
+ENV VLLM_USE_DEEP_GEMM=1
+ENV VLLM_MOE_USE_DEEP_GEMM=1
+
+# Use entrypoint for setup (ulimit, etc) but default to sleep infinity for K8s.
+ENTRYPOINT ["/app/docker-entrypoint.sh"]
+CMD ["sleep", "infinity"]
diff --git a/Dockerfile.dynamo b/Dockerfile.dynamo
new file mode 100644
index 0000000000..574976f49b
--- /dev/null
+++ b/Dockerfile.dynamo
@@ -0,0 +1,79 @@
+# syntax=docker/dockerfile:1.4
+# Dockerfile.dynamo — layer ai-dynamo onto a prime-rl image WITHOUT reinstalling vLLM.
+#
+# The prime-rl base already ships vLLM 0.20.2 (+ vLLM PR #39366 two-phase pause),
+# torch, flashinfer, DeepGEMM. We build the dynamo Rust bindings (ai-dynamo-runtime,
+# via maturin) at DYNAMO_REF, then install the dynamo Python package — which provides
+# BOTH `dynamo.frontend` and `dynamo.vllm` (hatch packages = components/src/dynamo) —
+# with `--no-deps` so the base's vLLM / torch / transformers are NEVER touched.
+# A curated set of dynamo runtime deps (explicitly excluding vllm/torch/ray) is added
+# via `uv pip` (the prime-rl venv is uv-managed and has no `pip` binary).
+#
+# Build (BuildKit; run in the arm64 dind builder):
+#   DOCKER_BUILDKIT=1 docker build -f Dockerfile.dynamo \
+#     --build-arg BASE_IMAGE=nvcr.io/nvidian/dynamo-dev/biswa:prime-rl-97950abd-20260531-arm64 \
+#     --build-arg DYNAMO_REF=ecae3569926410ef33b4d3d13c7d6a1b89789bb0 \
+#     -t nvcr.io/nvidian/dynamo-dev/biswa:prime-rl-97950abd-dynamo-ecae3569-arm64 .
+#
+# DYNAMO_REF may be any commit/branch/tag on https://github.com/ai-dynamo/dynamo
+# (e.g. bis/rl-workers-discovery tip ecae3569…, or a release tag v1.2.0).
+
+ARG BASE_IMAGE=nvcr.io/nvidian/dynamo-dev/biswa:prime-rl-97950abd-20260531-arm64
+ARG DYNAMO_REPO=https://github.com/ai-dynamo/dynamo.git
+ARG DYNAMO_REF=ecae3569926410ef33b4d3d13c7d6a1b89789bb0
+ARG CARGO_BUILD_JOBS=16
+
+# ===== Stage 1: build ai-dynamo-runtime Rust bindings wheel =====
+FROM ubuntu:24.04 AS dynamo-builder
+ARG DYNAMO_REPO
+ARG DYNAMO_REF
+ARG CARGO_BUILD_JOBS
+ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS}
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y --no-install-recommends \
+      ca-certificates curl build-essential pkg-config libclang-dev protobuf-compiler git \
+      python3 python3-dev python3-venv \
+ && rm -rf /var/lib/apt/lists/* \
+ && curl -fsSL https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable
+ENV CARGO_HOME=/root/.cargo RUSTUP_HOME=/root/.rustup PATH=/root/.cargo/bin:${PATH}
+RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+    cargo install maturin --locked
+RUN git clone "${DYNAMO_REPO}" /build/dynamo && cd /build/dynamo && git checkout "${DYNAMO_REF}"
+RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+    --mount=type=cache,target=/build/dynamo/lib/bindings/python/target,sharing=locked \
+    cd /build/dynamo/lib/bindings/python \
+ && maturin build --release --out /build/dist
+
+# ===== Stage 2: prime-rl base + dynamo (reuses base vLLM, no reinstall) =====
+FROM ${BASE_IMAGE}
+USER root
+ENV DYNAMO_HOME=/opt/dynamo
+COPY --from=dynamo-builder /build/dynamo /opt/dynamo
+COPY --from=dynamo-builder /build/dist/*.whl /tmp/dynamo-wheels/
+
+# 1) ai-dynamo-runtime (Rust bindings) + dynamo python pkg (frontend + vllm modules).
+#    --no-deps: do NOT pull vllm/torch/transformers (keep the prime-rl base's patched stack).
+RUN uv pip install --python /app/.venv/bin/python --no-cache /tmp/dynamo-wheels/*.whl \
+ && cd /opt/dynamo \
+ && uv pip install --python /app/.venv/bin/python --no-cache --no-deps -e . \
+ && rm -rf /tmp/dynamo-wheels
+
+# 2) dynamo runtime deps the base may lack — EXPLICITLY excluding vllm / torch / ray.
+#    uvloop + nixl are required by the dynamo.vllm worker; the rest are frontend/runtime.
+RUN uv pip install --python /app/.venv/bin/python --no-cache \
+      uvloop "nixl[cu12]<=0.10.1" \
+      "fastapi==0.120.1" "uvicorn==0.38.0" httpx \
+      "msgspec>=0.19.0" pyzmq "prometheus_client>=0.23.1" \
+      "aiohttp>=3.9.0,<4.0" "blake3>=1.0.0,<2.0.0" \
+      "kubernetes>=32.0.1,<33.0.0" \
+      opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp
+
+# NOTE: etcd + nats-server are intentionally NOT installed in this image.
+# In the k8s deployment dynamo uses the external dynamo-platform services
+# (e.g. NATS_SERVER=nats://dynamo-platform-nats...:4222 and the platform etcd),
+# so shipping the static binaries in the worker image is unnecessary bloat.
+
+USER appuser
+WORKDIR /app
diff --git a/scripts/docker-arm64-post-install.sh b/scripts/docker-arm64-post-install.sh
index 55f85a3a03..12a71ad987 100755
--- a/scripts/docker-arm64-post-install.sh
+++ b/scripts/docker-arm64-post-install.sh
@@ -42,3 +42,7 @@ echo "    target venv: $VENV_PATH"
 echo "=== reinstalling flash-attn-cute (flash-attn overwrites it with a stub) ==="
 uv pip install --python "$VENV_PATH/bin/python" --reinstall --no-deps \
     "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@96bd151#subdirectory=flash_attn/cute"
+
+echo "=== building DeepGEMM from source (FP8 grouped-GEMM for SM100 / GB200) ==="
+# Runs from /app (WORKDIR), so uv auto-detects the project venv at /app/.venv.
+bash /app/scripts/install_deep_gemm.sh
diff --git a/scripts/vllm-pr-39366.patch b/scripts/vllm-pr-39366.patch
new file mode 100644
index 0000000000..f8d973eff3
--- /dev/null
+++ b/scripts/vllm-pr-39366.patch
@@ -0,0 +1,206 @@
+diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
+index 4f903eeefa6d..6ba392802e31 100644
+--- a/vllm/config/parallel.py
++++ b/vllm/config/parallel.py
+@@ -663,6 +663,33 @@ def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool:
+         aggregated_has_unfinished = bool(tensor.item())
+         return aggregated_has_unfinished
+ 
++    @staticmethod
++    def sync_dp_state(
++        dp_group: ProcessGroup, has_unfinished: bool, pending_pause: bool
++    ) -> tuple[bool, bool]:
++        """Combined all-reduce for DP state synchronization.
++
++        Uses a single SUM all-reduce on a 2-element tensor:
++          [0] = 1 if this rank has unfinished work, else 0.
++                SUM > 0 ≡ logical OR across ranks → any rank has work.
++          [1] = 1 if this rank has a pending pause request, else 0.
++                SUM == dp_size ≡ all ranks reached pause consensus.
++
++        has_unfinished_global is true if any rank has unfinished work,
++        or if some ranks are waiting for a pause consensus.
++
++        Returns:
++            (has_unfinished_global, pause_consensus)
++        """
++        tensor = torch.tensor(
++            [int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu"
++        )
++        torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=dp_group)
++        dp_size = dp_group.size()
++        pause_count = tensor[1].item()
++        has_unfinished_global = tensor[0].item() > 0 or pause_count % dp_size != 0
++        return has_unfinished_global, pause_count == dp_size
++
+     @staticmethod
+     def sync_kv_cache_memory_size(dp_group: ProcessGroup, kv_cache_memory: int) -> int:
+         if kv_cache_memory == -1:
+diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
+index 36864ba738bf..11c5ee19a664 100644
+--- a/vllm/v1/engine/core.py
++++ b/vllm/v1/engine/core.py
+@@ -1571,7 +1571,8 @@ def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
+ 
+         pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
+         self.scheduler.set_pause_state(pause_state)
+-        if not self.has_work():
++
++        if self._pause_complete():
+             if clear_cache:
+                 self._reset_caches()
+             return None
+@@ -1580,6 +1581,13 @@ def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
+         self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))
+         return future
+ 
++    def _pause_complete(self) -> bool:
++        """Returns True if the pause has fully completed and the caller can
++        return ``None`` synchronously; False if the pause is still pending
++        and the caller should register an idle-state callback to finish it.
++        """
++        return not self.has_work()
++
+     def _send_finish_outputs_to_client(
+         self, req_ids: list[str], client_index: int, finish_reason: FinishReason
+     ) -> None:
+@@ -1635,6 +1643,14 @@ def __init__(
+         self.current_wave = 0
+         self.last_counts = (0, 0)
+ 
++        # Two-phase pause protocol state. When pending_pause is True, the
++        # engine keeps stepping (dummy batches) while waiting for all DP
++        # ranks to also set pending_pause. Once all ranks agree via
++        # all-reduce, ignore_start_dp_wave is set so that stale
++        # START_DP_WAVE messages cannot re-wake the engines.
++        self.pending_pause = False
++        self.ignore_start_dp_wave = False
++
+         from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
+ 
+         self.eep_scaling_state: ElasticEPScalingState | None = None
+@@ -1664,6 +1680,7 @@ def _init_data_parallel(self, vllm_config: VllmConfig):
+         assert 0 <= local_dp_rank <= dp_rank < dp_size
+ 
+         self.dp_rank = dp_rank
++        self.dp_size = dp_size
+         dp_group, dp_store = parallel_config.stateless_init_dp_group(return_store=True)
+         self.dp_group, self.dp_store = dp_group, dp_store
+ 
+@@ -1672,6 +1689,24 @@ def shutdown(self):
+         if dp_group := getattr(self, "dp_group", None):
+             stateless_destroy_torch_distributed_process_group(dp_group)
+ 
++    def _pause_complete(self) -> bool:
++        """Two-phase DP-aware pause.
++
++        Phase 1: Set local pause state and ``pending_pause`` flag. If the
++        engines are idle, kick-start them by setting ``engines_running`` to
++        True so ranks enter the stepping loop and reach the all-reduce
++        consensus checkpoint in ``_has_global_unfinished_reqs``.
++
++        Phase 2 (in ``_has_global_unfinished_reqs``): Once the all-reduce
++        confirms that **all** ranks have ``pending_pause`` set, collectively
++        stop stepping and set ``ignore_start_dp_wave`` so that stale
++        ``START_DP_WAVE`` messages cannot re-wake any engine.
++        """
++        self.pending_pause = True
++        self.engines_running = True
++
++        return False
++
+     def add_request(self, request: Request, request_wave: int = 0):
+         super().add_request(request, request_wave)
+         if self.has_coordinator and request_wave != self.current_wave:
+@@ -1681,36 +1716,60 @@ def add_request(self, request: Request, request_wave: int = 0):
+                 not self.engines_running
+                 and self.scheduler.pause_state == PauseState.UNPAUSED
+             ):
+-                self.engines_running = True
+                 # Request received for an already-completed wave, notify
+                 # front-end that we need to start the next one.
++                self.engines_running = True
+                 self.output_queue.put_nowait(
+                     (-1, EngineCoreOutputs(start_wave=self.current_wave))
+                 )
+ 
+     def resume_scheduler(self):
+-        super().resume_scheduler()
+-        if (
+-            self.has_coordinator
+-            and not self.engines_running
+-            and self.scheduler.has_unfinished_requests()
+-        ):
+-            # Wake up other DP engines.
+-            self.output_queue.put_nowait(
+-                (-1, EngineCoreOutputs(start_wave=self.current_wave))
++        if self.pending_pause or (self.engines_running and self.ignore_start_dp_wave):
++            raise RuntimeError(
++                "resume_scheduler called while pause is still in "
++                "flight. Wait for the pause future to resolve before "
++                "resuming."
+             )
++        if self.engines_running:
++            logger.debug("Resume called while engines are not paused, ignoring.")
++            return
++
++        super().resume_scheduler()
++        self.ignore_start_dp_wave = False
++
++        # Barrier: wait for all DP ranks to have resumed (and cleared
++        # ignore_start_dp_wave) before any rank starts stepping. Uses
++        # the existing all-reduce which is safe because engines are
++        # stopped.
++        has_global_unfinished = ParallelConfig.has_unfinished_dp(
++            self.dp_group, self.scheduler.has_unfinished_requests()
++        )
++
++        if has_global_unfinished:
++            self.engines_running = True
++
++    def barrier(self):
++        """Blocking barrier on the DP process group (test-only utility)."""
++        import torch.distributed as dist
++
++        dist.barrier(group=self.dp_group)
+ 
+     def _handle_client_request(
+         self, request_type: EngineCoreRequestType, request: Any
+     ) -> None:
+         if request_type == EngineCoreRequestType.START_DP_WAVE:
++            if self.ignore_start_dp_wave:
++                return
+             new_wave, exclude_eng_index = request
+             if exclude_eng_index != self.engine_index and (
+                 new_wave >= self.current_wave
+             ):
+                 self.current_wave = new_wave
+                 if not self.engines_running:
+-                    logger.debug("EngineCore starting idle loop for wave %d.", new_wave)
++                    logger.debug(
++                        "EngineCore starting idle loop for wave %d.",
++                        new_wave,
++                    )
+                     self.engines_running = True
+         else:
+             super()._handle_client_request(request_type, request)
+@@ -1790,7 +1849,18 @@ def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
+         if self.step_counter % 32 != 0:
+             return True
+ 
+-        return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished)
++        has_unfinished, pause_consensus = ParallelConfig.sync_dp_state(
++            self.dp_group,
++            has_unfinished=local_unfinished,
++            pending_pause=self.pending_pause,
++        )
++
++        if pause_consensus:
++            self.ignore_start_dp_wave = True
++            self.pending_pause = False
++            logger.debug("DP pause consensus reached, ignoring START_DP_WAVE.")
++
++        return has_unfinished
+ 
+     def reinitialize_distributed(
+         self, reconfig_request: ReconfigureDistributedRequest

From 76432984666b8fe6f88daf58b7bcaf3e86a315c0 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Sat, 6 Jun 2026 14:47:38 -0700
Subject: [PATCH 08/19] feat(deploy): helm chart updates + dynamo k8s manifests
 + smoke-test tools

---
 k8s/dynamo-deploy/admin-stub.yaml            |  73 ++++++++++++
 k8s/dynamo-deploy/dynamo-dgd.yaml            |  90 +++++++++++++++
 k8s/dynamo-deploy/prime-rl-configs.yaml      |  57 ++++++++++
 k8s/dynamo-deploy/prime-rl-values.yaml       | 114 +++++++++++++++++++
 k8s/prime-rl/templates/deployment.yaml       | 109 +++++++++++++++---
 k8s/prime-rl/templates/pvc.yaml              |   7 +-
 k8s/prime-rl/values.yaml                     |  56 ++++++++-
 tools/dynamo/admin_stub.py                   |  56 +++++++++
 tools/dynamo/configs/smoke_rl.toml           |  47 ++++++++
 tools/dynamo/configs/smoke_rl_long.toml      |  40 +++++++
 tools/dynamo/configs/smoke_trainer.toml      |  20 ++++
 tools/dynamo/configs/smoke_trainer_long.toml |  19 ++++
 tools/dynamo/run_dynamo.sh                   |  53 +++++++++
 tools/dynamo/run_smoke_test.sh               | 101 ++++++++++++++++
 14 files changed, 826 insertions(+), 16 deletions(-)
 create mode 100644 k8s/dynamo-deploy/admin-stub.yaml
 create mode 100644 k8s/dynamo-deploy/dynamo-dgd.yaml
 create mode 100644 k8s/dynamo-deploy/prime-rl-configs.yaml
 create mode 100644 k8s/dynamo-deploy/prime-rl-values.yaml
 create mode 100644 tools/dynamo/admin_stub.py
 create mode 100644 tools/dynamo/configs/smoke_rl.toml
 create mode 100644 tools/dynamo/configs/smoke_rl_long.toml
 create mode 100644 tools/dynamo/configs/smoke_trainer.toml
 create mode 100644 tools/dynamo/configs/smoke_trainer_long.toml
 create mode 100755 tools/dynamo/run_dynamo.sh
 create mode 100755 tools/dynamo/run_smoke_test.sh

diff --git a/k8s/dynamo-deploy/admin-stub.yaml b/k8s/dynamo-deploy/admin-stub.yaml
new file mode 100644
index 0000000000..ceeb3e7a87
--- /dev/null
+++ b/k8s/dynamo-deploy/admin-stub.yaml
@@ -0,0 +1,73 @@
+# Optional admin-stub Deployment + Service.
+# kubectl apply -f admin-stub.yaml -n <your-namespace>
+#
+# Only needed if your Dynamo build does NOT serve /v1/rl/* natively
+# (i.e. older builds without DYN_ENABLE_RL=true). With a recent Dynamo,
+# point `admin_base_url` directly at the Dynamo frontend and skip this
+# manifest entirely.
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: admin-stub-script
+  namespace: <your-namespace>
+data:
+  admin_stub.py: |
+    from http.server import HTTPServer, BaseHTTPRequestHandler
+    class H(BaseHTTPRequestHandler):
+        def do_POST(self):
+            n = int(self.headers.get("Content-Length", 0))
+            body = self.rfile.read(n) if n else b""
+            print(f"[stub] POST {self.path} body={body[:200]}")
+            self.send_response(200)
+            self.end_headers()
+            self.wfile.write(b"OK")
+        def do_GET(self):
+            self.send_response(200)
+            self.end_headers()
+            self.wfile.write(b"OK")
+    HTTPServer(("0.0.0.0", 8001), H).serve_forever()
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: admin-stub
+  namespace: <your-namespace>
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: admin-stub
+  template:
+    metadata:
+      labels:
+        app: admin-stub
+    spec:
+      containers:
+        - name: stub
+          image: python:3.12-slim
+          command: ["python3", "/scripts/admin_stub.py"]
+          ports:
+            - containerPort: 8001
+          volumeMounts:
+            - name: script
+              mountPath: /scripts
+          resources:
+            requests:
+              memory: "64Mi"
+              cpu: "50m"
+      volumes:
+        - name: script
+          configMap:
+            name: admin-stub-script
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: admin-stub
+  namespace: <your-namespace>
+spec:
+  selector:
+    app: admin-stub
+  ports:
+    - port: 8001
+      targetPort: 8001
diff --git a/k8s/dynamo-deploy/dynamo-dgd.yaml b/k8s/dynamo-deploy/dynamo-dgd.yaml
new file mode 100644
index 0000000000..9791d5d8d9
--- /dev/null
+++ b/k8s/dynamo-deploy/dynamo-dgd.yaml
@@ -0,0 +1,90 @@
+# Example DynamoGraphDeployment for serving inference to prime-rl.
+#   kubectl apply -f dynamo-dgd.yaml -n <your-namespace>
+#
+# Replace `<your-namespace>`, `<your-registry>/dynamo:<vllm-runtime-tag>`,
+# `<your-image-pull-secret>`, and the model name below for your environment.
+# Requires DYN_ENABLE_RL=true on the Dynamo runtime so /v1/rl/* endpoints are
+# served natively.
+apiVersion: nvidia.com/v1alpha1
+kind: DynamoGraphDeployment
+metadata:
+  name: prime-rl-dynamo
+  namespace: <your-namespace>
+spec:
+  backendFramework: vllm
+  pvcs:
+    - create: false
+      name: model-cache
+  services:
+    Frontend:
+      componentType: frontend
+
+      extraPodSpec:
+        imagePullSecrets:
+          - name: <your-image-pull-secret>
+        mainContainer:
+          image: <your-registry>/dynamo:<vllm-runtime-tag>
+          startupProbe:
+            failureThreshold: 360
+            httpGet:
+              path: /health
+              port: 8000
+            periodSeconds: 10
+            timeoutSeconds: 5
+      replicas: 1
+      volumeMounts:
+        - mountPoint: /model-cache
+          name: model-cache
+    VllmWorker:
+      componentType: worker
+
+      envFromSecret: hf-token-secret
+      extraPodSpec:
+        imagePullSecrets:
+          - name: <your-image-pull-secret>
+        mainContainer:
+          command:
+            - python3
+            - -m
+            - dynamo.vllm
+          args:
+            - --model
+            - Qwen/Qwen3-4B-Instruct-2507
+            - --served-model-name
+            - Qwen/Qwen3-4B-Instruct-2507
+            - --tensor-parallel-size
+            - "1"
+            - --max-model-len
+            - "2048"
+            - --max-num-seqs
+            - "64"
+            - --gpu-memory-utilization
+            - "0.90"
+            - --enforce-eager
+          env:
+            - name: HF_HOME
+              value: /model-cache/huggingface
+          image: <your-registry>/dynamo:<vllm-runtime-tag>
+          startupProbe:
+            failureThreshold: 360
+            httpGet:
+              path: /health
+              port: 9090
+            periodSeconds: 10
+            timeoutSeconds: 10
+          workingDir: /workspace/examples/backends/vllm
+        tolerations:
+          - effect: NoSchedule
+            key: nvidia.com/gpu
+            operator: Exists
+      replicas: 1
+      resources:
+        limits:
+          gpu: "1"
+        requests:
+          gpu: "1"
+      sharedMemory:
+        size: 16Gi
+      volumeMounts:
+        - mountPoint: /model-cache
+          name: model-cache
diff --git a/k8s/dynamo-deploy/prime-rl-configs.yaml b/k8s/dynamo-deploy/prime-rl-configs.yaml
new file mode 100644
index 0000000000..9cbfe50af9
--- /dev/null
+++ b/k8s/dynamo-deploy/prime-rl-configs.yaml
@@ -0,0 +1,57 @@
+# Example ConfigMap mounted at /configs in the orchestrator and trainer pods.
+# kubectl apply -f prime-rl-configs.yaml -n <your-namespace>
+#
+# Replace `<your-namespace>` below with your namespace, and adjust the
+# Dynamo frontend service hostname if it differs in your cluster.
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: prime-rl-configs
+  namespace: <your-namespace>
+data:
+  orch.toml: |
+    max_steps = 20
+    seq_len = 2048
+    batch_size = 64
+    rollouts_per_example = 4
+    use_token_client = false
+
+    [wandb]
+    project = "prime-rl-dynamo-k8s"
+    name = "dynamo-smoke-qwen3-4b"
+
+    [model]
+    name = "Qwen/Qwen3-4B-Instruct-2507"
+
+    [sampling]
+    max_tokens = 256
+
+    [[env]]
+    id = "math-env"
+    name = "hendrycks-math"
+    args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 32, math_verify_timeout = 60 }
+
+    [buffer]
+    easy_threshold = 1.0
+    hard_threshold = 0.0
+
+    [client]
+    base_url = ["http://prime-rl-dynamo-frontend.<your-namespace>.svc.cluster.local:8000/v1"]
+    # Admin endpoints (/v1/rl/*) are served natively by the Dynamo Rust frontend
+    # (DYN_ENABLE_RL=true). No separate admin-stub service needed.
+    admin_base_url = ["http://prime-rl-dynamo-frontend.<your-namespace>.svc.cluster.local:8000/v1/rl"]
+    skip_model_check = true
+
+  train.toml: |
+    max_steps = 20
+
+    [model]
+    name = "Qwen/Qwen3-4B-Instruct-2507"
+    seq_len = 2048
+
+    [wandb]
+    project = "prime-rl-dynamo-k8s"
+    name = "dynamo-smoke-qwen3-4b-trainer"
+
+    [optim]
+    lr = 3e-6
diff --git a/k8s/dynamo-deploy/prime-rl-values.yaml b/k8s/dynamo-deploy/prime-rl-values.yaml
new file mode 100644
index 0000000000..0b97ef87e2
--- /dev/null
+++ b/k8s/dynamo-deploy/prime-rl-values.yaml
@@ -0,0 +1,114 @@
+# Example Helm values for prime-rl when Dynamo serves inference.
+#
+#   helm install prime-rl k8s/prime-rl -f k8s/dynamo-deploy/prime-rl-values.yaml -n <your-namespace>
+#
+# Dynamo serves inference, so prime-rl's own inference component is disabled.
+# Orchestrator and trainer point at the Dynamo frontend (deploy with dynamo-dgd.yaml).
+#
+# Replace the following placeholders for your environment:
+#   - namespace
+#   - image.repository / image.tag (your prime-rl image)
+#   - storage.storageClassName (your cluster's RWX storage class)
+#   - orchestrator/trainer.imagePullSecrets[*].name
+#   - secrets.name (HF token / W&B key)
+
+namespace: <your-namespace>
+
+image:
+  repository: <your-registry>/prime-rl
+  tag: "latest"
+  pullPolicy: IfNotPresent
+
+storage:
+  enabled: true
+  storageClassName: <your-rwx-storageclass>
+  accessModes:
+    - ReadWriteMany
+  size: 100Gi
+  mountPath: /data
+
+config:
+  example: "prime-rl-dynamo"
+  secrets:
+    enabled: true
+    name: hf-token-secret
+
+# Orchestrator: CPU only, talks to Dynamo for inference
+orchestrator:
+  enabled: true
+  replicas: 1
+  autoStart: true
+  command: >-
+    BENCH_API_KEY=EMPTY
+    uv run orchestrator
+    @ /configs/orch.toml
+    --output-dir /data/outputs/run_default
+    --max-concurrent 16
+    ; echo "orchestrator exited: $?" ; sleep infinity
+  resources:
+    requests:
+      memory: "8Gi"
+      cpu: "4"
+  env:
+    - name: HF_TOKEN
+      valueFrom:
+        secretKeyRef:
+          name: hf-token-secret
+          key: HF_TOKEN
+          optional: true
+    - name: WANDB_API_KEY
+      valueFrom:
+        secretKeyRef:
+          name: wandb-secret
+          key: WANDB_API_KEY
+  configMap: prime-rl-configs
+  imagePullSecrets:
+    - name: <your-image-pull-secret>
+  nodeSelector: {}
+
+# Inference: DISABLED (Dynamo DGD replaces this)
+inference:
+  enabled: false
+
+# Trainer: 1 GPU is sufficient for small models (e.g. Qwen3-4B GRPO)
+trainer:
+  enabled: true
+  replicas: 1
+  autoStart: true
+  command: >-
+    uv run torchrun
+    --nproc-per-node=1
+    --rdzv-endpoint=localhost:29510
+    --rdzv-id=smoke_$(date +%s)
+    -m prime_rl.trainer.rl.train
+    @ /configs/train.toml
+    --output-dir /data/outputs
+    ; echo "trainer exited: $?" ; sleep infinity
+  pytorchCudaAllocConf: "expandable_segments:True"
+  gpu:
+    enabled: true
+    count: 1
+  resources:
+    requests:
+      memory: "32Gi"
+      cpu: "8"
+  env:
+    - name: HF_TOKEN
+      valueFrom:
+        secretKeyRef:
+          name: hf-token-secret
+          key: HF_TOKEN
+          optional: true
+    - name: WANDB_API_KEY
+      valueFrom:
+        secretKeyRef:
+          name: wandb-secret
+          key: WANDB_API_KEY
+  configMap: prime-rl-configs
+  runtimeClassName: nvidia
+  tolerations:
+    - key: nvidia.com/gpu
+      operator: Exists
+      effect: NoSchedule
+  imagePullSecrets:
+    - name: <your-image-pull-secret>
diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml
index 1ed16b9f5e..f40aa8046e 100644
--- a/k8s/prime-rl/templates/deployment.yaml
+++ b/k8s/prime-rl/templates/deployment.yaml
@@ -26,6 +26,14 @@ spec:
       nodeSelector:
         {{- toYaml . | nindent 8 }}
       {{- end }}
+      {{- with .Values.orchestrator.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.orchestrator.imagePullSecrets }}
+      imagePullSecrets:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
       containers:
       - name: prime-rl-orchestrator
         image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
@@ -33,7 +41,7 @@ spec:
         {{- if .Values.orchestrator.autoStart }}
         command: ["/bin/bash", "-c"]
         args:
-          - {{ .Values.orchestrator.command }}
+          - {{ .Values.orchestrator.command | quote }}
         {{- else }}
         command: ["sleep", "infinity"]
         {{- end }}
@@ -91,16 +99,29 @@ spec:
         {{- end }}
         resources:
           {{- toYaml .Values.orchestrator.resources | nindent 10 }}
-        {{- if .Values.storage.enabled }}
+        {{- if or .Values.storage.enabled .Values.orchestrator.configMap }}
         volumeMounts:
+        {{- if .Values.storage.enabled }}
         - name: shared-data
           mountPath: {{ .Values.storage.mountPath }}
         {{- end }}
-      {{- if .Values.storage.enabled }}
+        {{- if .Values.orchestrator.configMap }}
+        - name: rl-configs
+          mountPath: /configs
+        {{- end }}
+        {{- end }}
+      {{- if or .Values.storage.enabled .Values.orchestrator.configMap }}
       volumes:
+      {{- if .Values.storage.enabled }}
       - name: shared-data
         persistentVolumeClaim:
-          claimName: {{ .Release.Name }}-shared-data
+          claimName: {{ .Values.storage.existingClaim | default (printf "%s-shared-data" .Release.Name) }}
+      {{- end }}
+      {{- if .Values.orchestrator.configMap }}
+      - name: rl-configs
+        configMap:
+          name: {{ .Values.orchestrator.configMap }}
+      {{- end }}
       {{- end }}
 {{- end }}
 ---
@@ -132,6 +153,19 @@ spec:
       {{- if .Values.inference.runtimeClassName }}
       runtimeClassName: {{ .Values.inference.runtimeClassName }}
       {{- end }}
+      {{- with .Values.inference.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.inference.imagePullSecrets }}
+      imagePullSecrets:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- if .Values.inference.resourceClaim.enabled }}
+      resourceClaims:
+      - name: {{ .Values.inference.resourceClaim.name }}
+        resourceClaimTemplateName: {{ .Values.inference.resourceClaim.templateName }}
+      {{- end }}
       containers:
       - name: prime-rl-inference
         image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
@@ -139,7 +173,7 @@ spec:
         {{- if .Values.inference.autoStart }}
         command: ["/bin/bash", "-c"]
         args:
-          - {{ .Values.inference.command }}
+          - {{ .Values.inference.command | quote }}
         {{- else }}
         command: ["sleep", "infinity"]
         {{- end }}
@@ -184,19 +218,23 @@ spec:
         {{- end }}
         resources:
           requests:
-            {{- if .Values.inference.gpu.enabled }}
+            {{- if and .Values.inference.gpu.enabled (not .Values.inference.resourceClaim.enabled) }}
             nvidia.com/gpu: {{ .Values.inference.gpu.count }}
             {{- end }}
             memory: {{ .Values.inference.resources.requests.memory }}
             cpu: {{ .Values.inference.resources.requests.cpu }}
           limits:
-            {{- if .Values.inference.gpu.enabled }}
+            {{- if and .Values.inference.gpu.enabled (not .Values.inference.resourceClaim.enabled) }}
             nvidia.com/gpu: {{ .Values.inference.gpu.count }}
             {{- end }}
             {{- if .Values.inference.resources.limits }}
             memory: {{ .Values.inference.resources.limits.memory }}
             cpu: {{ .Values.inference.resources.limits.cpu }}
             {{- end }}
+          {{- if .Values.inference.resourceClaim.enabled }}
+          claims:
+          - name: {{ .Values.inference.resourceClaim.name }}
+          {{- end }}
         {{- if and .Values.inference.probes .Values.inference.probes.enabled }}
         startupProbe:
           httpGet:
@@ -220,16 +258,29 @@ spec:
           failureThreshold: {{ .Values.inference.probes.readiness.failureThreshold }}
           timeoutSeconds: {{ .Values.inference.probes.readiness.timeoutSeconds }}
         {{- end }}
-        {{- if .Values.storage.enabled }}
+        {{- if or .Values.storage.enabled .Values.inference.configMap }}
         volumeMounts:
+        {{- if .Values.storage.enabled }}
         - name: shared-data
           mountPath: {{ .Values.storage.mountPath }}
         {{- end }}
-      {{- if .Values.storage.enabled }}
+        {{- if .Values.inference.configMap }}
+        - name: rl-configs
+          mountPath: /configs
+        {{- end }}
+        {{- end }}
+      {{- if or .Values.storage.enabled .Values.inference.configMap }}
       volumes:
+      {{- if .Values.storage.enabled }}
       - name: shared-data
         persistentVolumeClaim:
-          claimName: {{ .Release.Name }}-shared-data
+          claimName: {{ .Values.storage.existingClaim | default (printf "%s-shared-data" .Release.Name) }}
+      {{- end }}
+      {{- if .Values.inference.configMap }}
+      - name: rl-configs
+        configMap:
+          name: {{ .Values.inference.configMap }}
+      {{- end }}
       {{- end }}
 {{- end }}
 ---
@@ -261,6 +312,19 @@ spec:
       {{- if .Values.trainer.runtimeClassName }}
       runtimeClassName: {{ .Values.trainer.runtimeClassName }}
       {{- end }}
+      {{- with .Values.trainer.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.trainer.imagePullSecrets }}
+      imagePullSecrets:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- if .Values.trainer.resourceClaim.enabled }}
+      resourceClaims:
+      - name: {{ .Values.trainer.resourceClaim.name }}
+        resourceClaimTemplateName: {{ .Values.trainer.resourceClaim.templateName }}
+      {{- end }}
       containers:
       - name: prime-rl-trainer
         image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
@@ -268,7 +332,7 @@ spec:
         {{- if .Values.trainer.autoStart }}
         command: ["/bin/bash", "-c"]
         args:
-          - {{ .Values.trainer.command }}
+          - {{ .Values.trainer.command | quote }}
         {{- else }}
         command: ["sleep", "infinity"]
         {{- end }}
@@ -343,6 +407,10 @@ spec:
             memory: {{ .Values.trainer.resources.limits.memory }}
             cpu: {{ .Values.trainer.resources.limits.cpu }}
             {{- end }}
+          {{- if .Values.trainer.resourceClaim.enabled }}
+          claims:
+          - name: {{ .Values.trainer.resourceClaim.name }}
+          {{- end }}
         {{- if and .Values.trainer.probes .Values.trainer.probes.enabled }}
         startupProbe:
           httpGet:
@@ -366,15 +434,28 @@ spec:
           failureThreshold: {{ .Values.trainer.probes.readiness.failureThreshold }}
           timeoutSeconds: {{ .Values.trainer.probes.readiness.timeoutSeconds }}
         {{- end }}
-        {{- if .Values.storage.enabled }}
+        {{- if or .Values.storage.enabled .Values.trainer.configMap }}
         volumeMounts:
+        {{- if .Values.storage.enabled }}
         - name: shared-data
           mountPath: {{ .Values.storage.mountPath }}
         {{- end }}
-      {{- if .Values.storage.enabled }}
+        {{- if .Values.trainer.configMap }}
+        - name: rl-configs
+          mountPath: /configs
+        {{- end }}
+        {{- end }}
+      {{- if or .Values.storage.enabled .Values.trainer.configMap }}
       volumes:
+      {{- if .Values.storage.enabled }}
       - name: shared-data
         persistentVolumeClaim:
-          claimName: {{ .Release.Name }}-shared-data
+          claimName: {{ .Values.storage.existingClaim | default (printf "%s-shared-data" .Release.Name) }}
+      {{- end }}
+      {{- if .Values.trainer.configMap }}
+      - name: rl-configs
+        configMap:
+          name: {{ .Values.trainer.configMap }}
+      {{- end }}
       {{- end }}
 {{- end }}
diff --git a/k8s/prime-rl/templates/pvc.yaml b/k8s/prime-rl/templates/pvc.yaml
index 7afd6ff386..a3e37a0f54 100644
--- a/k8s/prime-rl/templates/pvc.yaml
+++ b/k8s/prime-rl/templates/pvc.yaml
@@ -1,4 +1,9 @@
-{{- if .Values.storage.enabled }}
+{{- /*
+  Only create a PVC when storage is enabled AND the user has not supplied an
+  existing claim. When `storage.existingClaim` is set the chart mounts that
+  claim instead — e.g. the cluster-shared `shared-model-cache` PVC.
+*/ -}}
+{{- if and .Values.storage.enabled (not .Values.storage.existingClaim) }}
 apiVersion: v1
 kind: PersistentVolumeClaim
 metadata:
diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml
index 7fc3b48989..1905c6ddab 100644
--- a/k8s/prime-rl/values.yaml
+++ b/k8s/prime-rl/values.yaml
@@ -14,7 +14,14 @@ image:
 # Shared storage configuration
 storage:
   enabled: true
-  # PVC name will be automatically set to {{ .Release.Name }}-shared-data
+  # If existingClaim is set, the chart will not create a PVC and will instead
+  # mount the named claim. Useful for cluster-wide model caches (e.g. the
+  # auto-provisioned `shared-model-cache` PVC documented at
+  # work/ops/shared-model-cache.md) so models downloaded by other namespaces
+  # are reused with zero re-download.
+  existingClaim: ""
+  # Used only when existingClaim is empty (chart creates the PVC).
+  # Resulting PVC name: {{ .Release.Name }}-shared-data
   storageClassName: nfs
   accessModes:
     - ReadWriteMany
@@ -50,6 +57,19 @@ orchestrator:
   nodeSelector: {}
     # nvidia.com/gpu.present: "true"  # Orchestrator doesn't need GPUs
 
+  tolerations: []
+    # - key: "nvidia.com/gpu"
+    #   operator: "Exists"
+    #   effect: "NoSchedule"
+
+  imagePullSecrets: []
+    # - name: regcred
+
+  # Optional: name of an existing ConfigMap to mount at /configs in the
+  # orchestrator container. Useful for shipping TOML configs without rebuilding
+  # the image. Leave empty to skip.
+  configMap: ""
+
 # Inference component
 inference:
   enabled: true
@@ -63,6 +83,18 @@ inference:
     enabled: true
     count: 1
 
+  # Dynamic Resource Allocation (DRA) — see trainer.resourceClaim for semantics.
+  resourceClaim:
+    enabled: false
+    templateName: ""
+    name: "gpus"
+
+  # Optional: name of an existing ConfigMap to mount at /configs in the
+  # inference container. Symmetric with orchestrator.configMap and
+  # trainer.configMap. Lets a native vLLM inference server read its
+  # infer.toml from a ConfigMap.
+  configMap: ""
+
   resources:
     requests:
       memory: "4Gi"
@@ -75,6 +107,9 @@ inference:
 
   runtimeClassName: nvidia
 
+  tolerations: []
+  imagePullSecrets: []
+
   # Health probes for inference server (/health for startup/readiness, /liveness for engine liveness)
   probes:
     enabled: false
@@ -108,6 +143,22 @@ trainer:
     enabled: true
     count: 1
 
+  # Dynamic Resource Allocation (DRA) — for DRA-only node pools that reject the
+  # legacy `nvidia.com/gpu` device-plugin resource. When `resourceClaim.enabled`
+  # is true the chart will NOT emit `nvidia.com/gpu` requests/limits and will
+  # instead add a `spec.resourceClaims` entry to the pod and a `resources.claims`
+  # reference inside the container.
+  #   templateName: name of a pre-existing ResourceClaimTemplate in the same ns
+  #   name:         logical claim name used to link pod <-> container
+  resourceClaim:
+    enabled: false
+    templateName: ""
+    name: "gpus"
+
+  # Optional: name of an existing ConfigMap to mount at /configs in the
+  # trainer container. Symmetric with `orchestrator.configMap`.
+  configMap: ""
+
   resources:
     requests:
       memory: "4Gi"
@@ -123,6 +174,9 @@ trainer:
 
   runtimeClassName: nvidia
 
+  tolerations: []
+  imagePullSecrets: []
+
   # Health probes for trainer (requires metrics_server config)
   probes:
     enabled: false
diff --git a/tools/dynamo/admin_stub.py b/tools/dynamo/admin_stub.py
new file mode 100644
index 0000000000..8a088c2e46
--- /dev/null
+++ b/tools/dynamo/admin_stub.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""
+Stub HTTP server for prime-rl admin endpoints.
+
+NOTE: As of dynamo#8630 (bis/parity-tokenize-tcp), Dynamo's Rust frontend
+implements these routes natively at /v1/rl/* when DYN_ENABLE_RL=true:
+  POST /v1/rl/load_lora_adapter
+  POST /v1/rl/unload_lora_adapter
+  GET  /v1/rl/health
+
+For K8s and any deployment with a real Dynamo frontend, point admin_base_url
+at the Dynamo service (e.g. http://<frontend-svc>:8000/v1/rl). This stub is
+kept as a local development fallback for running the orchestrator without a
+live Dynamo instance.
+
+Usage:
+    python tools/dynamo/admin_stub.py
+    python tools/dynamo/admin_stub.py --port 8001
+"""
+import argparse
+
+from aiohttp import web
+
+
+async def pause(request):
+    print("[stub] POST /pause - OK")
+    return web.Response(status=200, text="OK")
+
+
+async def resume(request):
+    print("[stub] POST /resume - OK")
+    return web.Response(status=200, text="OK")
+
+
+async def update_weights(request):
+    body = await request.json()
+    print(f"[stub] POST /update_weights weight_dir={body.get('weight_dir')} - OK (weights not reloaded)")
+    return web.Response(status=200, text="OK")
+
+
+async def health(request):
+    return web.Response(status=200, text="OK")
+
+
+app = web.Application()
+app.router.add_post("/pause", pause)
+app.router.add_post("/resume", resume)
+app.router.add_post("/update_weights", update_weights)
+app.router.add_get("/health", health)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Dynamo admin stub server")
+    parser.add_argument("--port", type=int, default=8001, help="Port to listen on")
+    args = parser.parse_args()
+    print(f"[stub] Dynamo admin stub server starting on port {args.port}...")
+    web.run_app(app, port=args.port)
diff --git a/tools/dynamo/configs/smoke_rl.toml b/tools/dynamo/configs/smoke_rl.toml
new file mode 100644
index 0000000000..b2c1d77342
--- /dev/null
+++ b/tools/dynamo/configs/smoke_rl.toml
@@ -0,0 +1,47 @@
+# Prime-RL smoke test: Dynamo as inference backend (short -- 5 steps)
+# Model: PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT
+# GPU layout: GPU 0 = Dynamo, GPU 1 = trainer
+#
+# Usage: uv run orchestrator @ tools/dynamo/configs/smoke_rl.toml --output-dir /tmp/.../run_default
+
+max_steps = 5
+seq_len = 512
+batch_size = 16
+rollouts_per_example = 4
+# Disable TITO: Dynamo injects token_ids via nvext (DYN_ENABLE_RL=true) rather than
+# the TITO /v1/chat/completions/tokens path. Keep false for all Dynamo deployments.
+use_token_client = false
+
+[model]
+name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"
+
+[wandb]
+offline = true
+
+[sampling]
+max_tokens = 64
+
+[[filters]]
+type = "gibberish"
+enforce = false
+
+[[filters]]
+type = "repetition"
+enforce = false
+
+[[filters]]
+type = "zero_advantage"
+enforce = false
+
+[[env]]
+id = "reverse-text"
+
+[client]
+# Point at Dynamo instead of prime-rl's own vLLM
+base_url = ["http://localhost:8000/v1"]
+# Discover worker-advertised system URLs from Dynamo's dedicated RL listener,
+# then call each worker's /engine/* admin routes directly.
+backend = "dynamo"
+rl_base_url = ["http://localhost:8001/v1"]
+# Dynamo may serve model under a different name -- skip the /v1/models check
+skip_model_check = true
diff --git a/tools/dynamo/configs/smoke_rl_long.toml b/tools/dynamo/configs/smoke_rl_long.toml
new file mode 100644
index 0000000000..d21c8f9cb9
--- /dev/null
+++ b/tools/dynamo/configs/smoke_rl_long.toml
@@ -0,0 +1,40 @@
+# Prime-RL LONG smoke test: Dynamo as inference backend (20 steps)
+# Enough steps to observe training dynamics over more iterations.
+#
+# Usage: uv run orchestrator @ tools/dynamo/configs/smoke_rl_long.toml --output-dir /tmp/.../run_default
+
+max_steps = 20
+seq_len = 512
+batch_size = 16
+rollouts_per_example = 4
+use_token_client = false
+
+[model]
+name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"
+
+[wandb]
+offline = true
+
+[sampling]
+max_tokens = 64
+
+[[filters]]
+type = "gibberish"
+enforce = false
+
+[[filters]]
+type = "repetition"
+enforce = false
+
+[[filters]]
+type = "zero_advantage"
+enforce = false
+
+[[env]]
+id = "reverse-text"
+
+[client]
+base_url = ["http://localhost:8000/v1"]
+backend = "dynamo"
+rl_base_url = ["http://localhost:8001/v1"]
+skip_model_check = true
diff --git a/tools/dynamo/configs/smoke_trainer.toml b/tools/dynamo/configs/smoke_trainer.toml
new file mode 100644
index 0000000000..aff75cd895
--- /dev/null
+++ b/tools/dynamo/configs/smoke_trainer.toml
@@ -0,0 +1,20 @@
+# Prime-RL trainer config for Dynamo smoke test (short -- 5 steps)
+# Single GPU, no FSDP. Reads rollouts written by the orchestrator process.
+#
+# Usage: uv run torchrun --nproc-per-node=1 --rdzv-endpoint=localhost:29510 \
+#          -m prime_rl.trainer.rl.train @ tools/dynamo/configs/smoke_trainer.toml \
+#          --output-dir /tmp/dynamo_smoke_outputs
+
+max_steps = 5
+
+[model]
+name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"
+seq_len = 512
+
+[wandb]
+offline = true
+
+[optim]
+lr = 3e-6
+
+[ckpt]
diff --git a/tools/dynamo/configs/smoke_trainer_long.toml b/tools/dynamo/configs/smoke_trainer_long.toml
new file mode 100644
index 0000000000..82768c456e
--- /dev/null
+++ b/tools/dynamo/configs/smoke_trainer_long.toml
@@ -0,0 +1,19 @@
+# Prime-RL trainer config for long smoke test (20 steps)
+#
+# Usage: uv run torchrun --nproc-per-node=1 --rdzv-endpoint=localhost:29510 \
+#          -m prime_rl.trainer.rl.train @ tools/dynamo/configs/smoke_trainer_long.toml \
+#          --output-dir /tmp/dynamo_smoke_long
+
+max_steps = 20
+
+[model]
+name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"
+seq_len = 512
+
+[wandb]
+offline = true
+
+[optim]
+lr = 3e-6
+
+[ckpt]
diff --git a/tools/dynamo/run_dynamo.sh b/tools/dynamo/run_dynamo.sh
new file mode 100755
index 0000000000..6a27580cec
--- /dev/null
+++ b/tools/dynamo/run_dynamo.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+# Launch Dynamo frontend + vLLM worker for smoke testing with prime-rl.
+#
+# Prerequisites:
+#   - Dynamo virtualenv at $DYNAMO_VENV (default: ~/dev/dynamo/dynamo)
+#   - etcd + NATS running (dynamo depends on them)
+#
+# Usage:
+#   ./tools/dynamo/run_dynamo.sh
+#   CUDA_VISIBLE_DEVICES=0 ./tools/dynamo/run_dynamo.sh
+#   DYNAMO_MODEL=my-org/my-model ./tools/dynamo/run_dynamo.sh
+
+set -euo pipefail
+
+DYNAMO_VENV="${DYNAMO_VENV:-/home/biswaranjanp/dev/rl/dynamo/.venv}"
+DYNAMO_MODEL="${DYNAMO_MODEL:-PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT}"
+DYNAMO_GPU_MEMORY_UTILIZATION="${DYNAMO_GPU_MEMORY_UTILIZATION:-0.50}"
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
+LOG_DIR="${LOG_DIR:-/tmp}"
+PRIME_RL_SRC="${PRIME_RL_SRC:-/home/biswaranjanp/dev/rl/prime-rl/src}"
+
+export CUDA_VISIBLE_DEVICES
+
+source "$DYNAMO_VENV/bin/activate"
+
+echo "[dynamo-smoke] Starting frontend..."
+DYN_ENABLE_RL=true \
+DYN_RL_PORT="${DYN_RL_PORT:-8001}" \
+python -m dynamo.frontend > "$LOG_DIR/dynamo_frontend.log" 2>&1 &
+FRONTEND_PID=$!
+echo "[dynamo-smoke] Frontend PID: $FRONTEND_PID (log: $LOG_DIR/dynamo_frontend.log)"
+
+sleep 5
+
+echo "[dynamo-smoke] Starting vLLM worker (model: $DYNAMO_MODEL)..."
+PYTHONPATH="${PRIME_RL_SRC}${PYTHONPATH:+:$PYTHONPATH}" \
+DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT:-8081}" \
+python -m dynamo.vllm \
+  --model "$DYNAMO_MODEL" \
+  --enforce-eager \
+  --max-model-len 2048 \
+  --max-num-seqs 32 \
+  --gpu-memory-utilization "$DYNAMO_GPU_MEMORY_UTILIZATION" \
+  --enable-rl \
+  --worker-extension-cls prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker \
+  > "$LOG_DIR/dynamo_vllm.log" 2>&1 &
+WORKER_PID=$!
+echo "[dynamo-smoke] Worker PID: $WORKER_PID (log: $LOG_DIR/dynamo_vllm.log)"
+
+echo "[dynamo-smoke] Both processes launched."
+echo "[dynamo-smoke] Frontend log: $LOG_DIR/dynamo_frontend.log"
+echo "[dynamo-smoke] Worker log:   $LOG_DIR/dynamo_vllm.log"
+wait
diff --git a/tools/dynamo/run_smoke_test.sh b/tools/dynamo/run_smoke_test.sh
new file mode 100755
index 0000000000..16e5af1548
--- /dev/null
+++ b/tools/dynamo/run_smoke_test.sh
@@ -0,0 +1,101 @@
+#!/bin/bash
+# Full Dynamo + prime-rl smoke test: orchestrator + trainer.
+#
+# GPU 0: Dynamo inference (must already be running -- see run_dynamo.sh)
+# GPU 1: prime-rl trainer (single GPU, uses torchrun)
+#
+# Directory structure created:
+#   $OUTPUT_DIR/             <- trainer output_dir
+#   $OUTPUT_DIR/run_default/ <- orchestrator output_dir (run_* naming required)
+#
+# Usage:
+#   # Short run (5 steps, default):
+#   ./tools/dynamo/run_smoke_test.sh
+#
+#   # Long run (20 steps):
+#   ./tools/dynamo/run_smoke_test.sh --long
+#
+#   # Custom output dir:
+#   OUTPUT_DIR=/tmp/my_run ./tools/dynamo/run_smoke_test.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
+CONFIG_DIR="$SCRIPT_DIR/configs"
+
+# Defaults
+OUTPUT_DIR="${OUTPUT_DIR:-/tmp/dynamo_smoke_outputs}"
+TRAINER_GPU="${TRAINER_GPU:-1}"
+RDZV_PORT="${RDZV_PORT:-29510}"
+SMOKE_MODEL="${SMOKE_MODEL:-}"
+MODEL_ARGS=()
+if [[ -n "$SMOKE_MODEL" ]]; then
+    MODEL_ARGS=(--model.name "$SMOKE_MODEL" --tokenizer.name "$SMOKE_MODEL")
+fi
+
+# Pick short vs long config
+if [[ "${1:-}" == "--long" ]]; then
+    ORCH_CONFIG="$CONFIG_DIR/smoke_rl_long.toml"
+    TRAINER_CONFIG="$CONFIG_DIR/smoke_trainer_long.toml"
+    ORCH_LOG="${OUTPUT_DIR}/smoke_long_orchestrator.log"
+    TRAINER_LOG="${OUTPUT_DIR}/smoke_long_trainer.log"
+    echo "[smoke] Using LONG configs (20 steps)"
+else
+    ORCH_CONFIG="$CONFIG_DIR/smoke_rl.toml"
+    TRAINER_CONFIG="$CONFIG_DIR/smoke_trainer.toml"
+    ORCH_LOG="${OUTPUT_DIR}/smoke_orchestrator.log"
+    TRAINER_LOG="${OUTPUT_DIR}/smoke_trainer.log"
+    echo "[smoke] Using SHORT configs (5 steps)"
+fi
+
+ORCH_DIR="$OUTPUT_DIR/run_default"
+mkdir -p "$ORCH_DIR"
+
+cd "$REPO_DIR"
+
+# Start orchestrator (no GPU needed -- just API calls to Dynamo)
+echo "[smoke] Starting orchestrator (output: $ORCH_DIR)..."
+BENCH_API_KEY=EMPTY CUDA_VISIBLE_DEVICES="" \
+  uv run orchestrator \
+    @ "$ORCH_CONFIG" \
+    --output-dir "$ORCH_DIR" \
+    "${MODEL_ARGS[@]}" \
+  > "$ORCH_LOG" 2>&1 &
+ORCH_PID=$!
+echo "[smoke] Orchestrator PID: $ORCH_PID (log: $ORCH_LOG)"
+
+# Give orchestrator time to write first batch
+sleep 12
+
+# Start trainer on GPU via uv run torchrun (uses project venv Python, needed for torch.distributed)
+echo "[smoke] Starting trainer (output: $OUTPUT_DIR)..."
+CUDA_VISIBLE_DEVICES="$TRAINER_GPU" uv run torchrun \
+  --nproc-per-node=1 \
+  --rdzv-endpoint="localhost:$RDZV_PORT" \
+  --rdzv-id="smoke_$(date +%s)" \
+    -m prime_rl.trainer.rl.train \
+    @ "$TRAINER_CONFIG" \
+    --output-dir "$OUTPUT_DIR" \
+    "${MODEL_ARGS[@]}" \
+  > "$TRAINER_LOG" 2>&1 &
+TRAINER_PID=$!
+echo "[smoke] Trainer PID: $TRAINER_PID (log: $TRAINER_LOG)"
+
+echo "[smoke] Waiting for both processes..."
+
+ORCH_EXIT=0
+TRAINER_EXIT=0
+wait $ORCH_PID || ORCH_EXIT=$?
+wait $TRAINER_PID || TRAINER_EXIT=$?
+
+echo ""
+echo "[smoke] Orchestrator exit: $ORCH_EXIT"
+echo "[smoke] Trainer exit: $TRAINER_EXIT"
+
+if [[ $ORCH_EXIT -eq 0 && $TRAINER_EXIT -eq 0 ]]; then
+    echo "[smoke] Smoke test PASSED."
+else
+    echo "[smoke] Smoke test FAILED."
+    exit 1
+fi

From 49a4d1fd452dbaa7ae6b717a9a2ae50a464daba3 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Tue, 9 Jun 2026 03:33:58 -0700
Subject: [PATCH 09/19] fix(rl): rename renderer_transport values to
 vllm_generate/dynamo_chat; bump verifiers/renderers deps to rl-sdk-4 heads

---
 deps/renderers                     | 2 +-
 deps/verifiers                     | 2 +-
 src/prime_rl/orchestrator/utils.py | 6 +++---
 src/prime_rl/utils/client.py       | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/deps/renderers b/deps/renderers
index 6a215742b0..5dbf4941f1 160000
--- a/deps/renderers
+++ b/deps/renderers
@@ -1 +1 @@
-Subproject commit 6a215742b0d89a1d4e79e8d217efa4909ed8b0a5
+Subproject commit 5dbf4941f118a126c3be2031cf9d6b4b68cf23fe
diff --git a/deps/verifiers b/deps/verifiers
index ee3482aebf..d713edc7ab 160000
--- a/deps/verifiers
+++ b/deps/verifiers
@@ -1 +1 @@
-Subproject commit ee3482aebfaf35e47ec73a55db9276364d63e1cd
+Subproject commit d713edc7ab58b4f3f2ad79e410f09d2c4166042c
diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py
index 592b6e7b76..c056f12cd7 100644
--- a/src/prime_rl/orchestrator/utils.py
+++ b/src/prime_rl/orchestrator/utils.py
@@ -261,14 +261,14 @@ async def compute_teacher_logprobs(
     Dispatches to the vLLM-sidecar or dynamo-nvext path based on the
     per-client ``renderer_transport``:
 
-      - ``prime_vllm_generate``  (default): POST ``/inference/v1/generate``
-      - ``dynamo_chat_nvext``               : POST ``/v1/chat/completions`` with nvext
+      - ``vllm_generate``  (default): POST ``/inference/v1/generate``
+      - ``dynamo_chat``               : POST ``/v1/chat/completions`` with nvext
 
     Both flatten to ``list[float]`` via the shared helper.
     """
 
     async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]:
-        if client_config.renderer_transport == "dynamo_chat_nvext":
+        if client_config.renderer_transport == "dynamo_chat":
             return await _compute_teacher_logprobs_dynamo(client_config, model_name, sample)
         return await _compute_teacher_logprobs_vllm(client_config, model_name, sample)
 
diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py
index fddab3089b..fd608a92ea 100644
--- a/src/prime_rl/utils/client.py
+++ b/src/prime_rl/utils/client.py
@@ -433,7 +433,7 @@ def setup_clients(
     #   - request:  nvext.token_data carries pre-tokenized prompt
     #   - response: nvext.engine_data carries completion_token_ids + logprobs
     # Default backend keeps the legacy vLLM TITO surface.
-    renderer_transport = "dynamo_chat_nvext" if client_config.backend == "dynamo" else "prime_vllm_generate"
+    renderer_transport = "dynamo_chat" if client_config.backend == "dynamo" else "vllm_generate"
     clients = []
     client_idx = 0
     # Only forward the renderer config when the client actually uses a

From e75f34314f6578e563967a4e0e2e94e5c909baa5 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Wed, 10 Jun 2026 02:33:59 -0700
Subject: [PATCH 10/19] fix(routed_experts): carry dtype for models with over
 256 experts

---
 src/prime_rl/inference/vllm/routed_experts.py | 13 +++++++++++--
 src/prime_rl/orchestrator/trajectories.py     |  6 +++++-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py
index e9ef70b049..421a5a6c96 100644
--- a/src/prime_rl/inference/vllm/routed_experts.py
+++ b/src/prime_rl/inference/vllm/routed_experts.py
@@ -15,15 +15,24 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A
     array = np.asarray(routed_experts)
     assert array.ndim == 3
     assert np.issubdtype(array.dtype, np.integer)
+    # Narrow to the smallest unsigned int that holds the expert ids, matching
+    # vLLM's RoutedExpertsManager (uint8 for <=256 experts, uint16 otherwise).
+    # This keeps the wire payload compact while still supporting >256-expert
+    # MoEs (e.g. Kimi-K2) that overflow uint8. The dtype rides the payload so
+    # the consumer decodes with the right element type.
     if array.size:
         assert array.min() >= 0
-        assert array.max() <= np.iinfo(np.uint8).max
+        max_id = int(array.max())
+    else:
+        max_id = 0
+    target_dtype = np.uint8 if max_id <= np.iinfo(np.uint8).max else np.uint16
 
-    compact = np.ascontiguousarray(array.astype(np.uint8, copy=False))
+    compact = np.ascontiguousarray(array.astype(target_dtype, copy=False))
     return {
         "data": pybase64.b64encode(memoryview(compact)).decode("ascii"),
         "shape": list(compact.shape),
         "start": start,
+        "dtype": np.dtype(target_dtype).name,
     }
 
 
diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py
index 3e8431c12a..592681f0cc 100644
--- a/src/prime_rl/orchestrator/trajectories.py
+++ b/src/prime_rl/orchestrator/trajectories.py
@@ -247,7 +247,11 @@ def prepare_step_tokens(step: vf.TrajectoryStep, step_idx: int) -> dict[str, Any
             routed_experts_start = None
             if routed_experts_payload is not None:
                 decoded_routed_experts = pybase64.b64decode_as_bytearray(routed_experts_payload["data"])
-                routed_experts = np.frombuffer(decoded_routed_experts, dtype=np.uint8).reshape(
+                # dtype rides the payload so >256-expert MoEs (uint16) and
+                # int32 captures decode correctly; default uint8 for payloads
+                # serialized before the dtype field existed.
+                re_dtype = np.dtype(routed_experts_payload.get("dtype", "uint8"))
+                routed_experts = np.frombuffer(decoded_routed_experts, dtype=re_dtype).reshape(
                     routed_experts_payload["shape"]
                 )
                 routed_experts_start = routed_experts_payload["start"]

From cbc27924b45439b3dc2b9800911deeb0f22f7359 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Wed, 10 Jun 2026 09:08:13 -0700
Subject: [PATCH 11/19] feat(inference): forward moe_backend to vLLM so
 router-replay captures real routing

---
 .../prime-rl-configs/src/prime_rl/configs/inference.py    | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
index 9579259366..203cc37f83 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -342,6 +342,13 @@ class InferenceConfig(BaseConfig):
     enable_return_routed_experts: bool = False
     """Return routed experts in responses. Forwarded as ``--enable-return-routed-experts``."""
 
+    moe_backend: str = "auto"
+    """MoE kernel backend, forwarded as vLLM ``--moe-backend``. The default fused
+    backends (e.g. FlashInfer TRTLLM) compute gating+routing+experts in one kernel and
+    bypass the Python ``BaseRouter`` that the routed-experts capture hook binds to, so
+    capture yields all-zero routing. Use a non-fused backend (``triton``) when
+    ``enable_return_routed_experts``/router replay is on so real routing is captured."""
+
     enable_fp32_lm_head: bool = True
     """Run the lm_head projection in fp32 via a native bf16×bf16 → fp32 GEMM (``torch.mm`` with ``out_dtype=torch.float32``). Stabilizes logprob precision under FP8/bf16 inference, matching SGLang's ``--enable-fp32-lm-head``. Implemented as a monkey-patch over vLLM's LogitsProcessor, activated by setting ``additional_config["fp32_lm_head"] = True`` on the vLLM config."""
 
@@ -503,6 +510,7 @@ def to_vllm(self) -> Namespace:
             "gpu_memory_utilization": "gpu_memory_utilization",
             "api_server_count": "api_server_count",
             "enable_return_routed_experts": "enable_return_routed_experts",
+            "moe_backend": "moe_backend",
             "enable_expert_parallel": "enable_expert_parallel",
             "all2all_backend": "all2all_backend",
             "enable_eplb": "enable_eplb",

From 788281c00a949255016762ff6acda8384ee296d7 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Wed, 10 Jun 2026 11:14:31 -0700
Subject: [PATCH 12/19] fix(routed_experts): int32 fallback for >65535 experts,
 normalize uint16 to int32, auto-triton on router replay

---
 .../src/prime_rl/configs/rl.py                | 12 +++++++++++-
 src/prime_rl/inference/vllm/routed_experts.py | 19 +++++++++++++------
 src/prime_rl/orchestrator/trajectories.py     |  7 +++++++
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
index dab46a9ce1..6aa3cd1ff7 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -428,9 +428,19 @@ def auto_setup_router_replay(self):
                         stacklevel=2,
                     )
                 self.inference.enable_return_routed_experts = True
+                # Fused MoE backends (FlashInfer/TRTLLM) bypass the BaseRouter
+                # capture hook, yielding all-zero routing. Force a non-fused
+                # backend for capture unless the user set one explicitly.
+                if self.inference.moe_backend == "auto":
+                    warnings.warn(
+                        "Router replay is enabled; setting inference.moe_backend='triton' "
+                        "(fused MoE backends bypass routed-experts capture).",
+                        stacklevel=2,
+                    )
+                    self.inference.moe_backend = "triton"
             else:
                 warnings.warn(
-                    "Router replay is enabled, but inference is not configured. When manually starting the inference server, make sure to pass `--enable-return-routed-experts` to the vLLM server.",
+                    "Router replay is enabled, but inference is not configured. When manually starting the inference server, make sure to pass `--enable-return-routed-experts` and `--moe-backend triton` to the vLLM server.",
                     stacklevel=2,
                 )
         return self
diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py
index 421a5a6c96..a3fce14211 100644
--- a/src/prime_rl/inference/vllm/routed_experts.py
+++ b/src/prime_rl/inference/vllm/routed_experts.py
@@ -15,17 +15,24 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A
     array = np.asarray(routed_experts)
     assert array.ndim == 3
     assert np.issubdtype(array.dtype, np.integer)
-    # Narrow to the smallest unsigned int that holds the expert ids, matching
-    # vLLM's RoutedExpertsManager (uint8 for <=256 experts, uint16 otherwise).
-    # This keeps the wire payload compact while still supporting >256-expert
-    # MoEs (e.g. Kimi-K2) that overflow uint8. The dtype rides the payload so
-    # the consumer decodes with the right element type.
+    # Narrow to the smallest int that holds the expert ids, matching vLLM's
+    # RoutedExpertsManager (uint8 for <=256 experts, uint16 otherwise). Keeps
+    # the wire payload compact while supporting >256-expert MoEs (e.g. Kimi-K2)
+    # that overflow uint8. The dtype rides the payload so the consumer decodes
+    # with the right element type.
     if array.size:
         assert array.min() >= 0
         max_id = int(array.max())
     else:
         max_id = 0
-    target_dtype = np.uint8 if max_id <= np.iinfo(np.uint8).max else np.uint16
+    if max_id <= np.iinfo(np.uint8).max:
+        target_dtype = np.uint8
+    elif max_id <= np.iinfo(np.uint16).max:
+        target_dtype = np.uint16
+    else:
+        # Beyond uint16 (>65535 experts) astype(uint16) would wrap and corrupt
+        # routing; fall back to int32 (consumer decodes via the dtype field).
+        target_dtype = np.int32
 
     compact = np.ascontiguousarray(array.astype(target_dtype, copy=False))
     return {
diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py
index 592681f0cc..f5e11984cd 100644
--- a/src/prime_rl/orchestrator/trajectories.py
+++ b/src/prime_rl/orchestrator/trajectories.py
@@ -254,6 +254,13 @@ def prepare_step_tokens(step: vf.TrajectoryStep, step_idx: int) -> dict[str, Any
                 routed_experts = np.frombuffer(decoded_routed_experts, dtype=re_dtype).reshape(
                     routed_experts_payload["shape"]
                 )
+                # The trainer (batch.py row sizing + tensor decode) supports
+                # uint8/int16/int32 (all torch-safe). Normalize a uint16 wire
+                # dtype (vLLM/dynamo >256-expert capture) to int32 so it never
+                # reaches the trainer; expert ids are non-negative so this is
+                # lossless.
+                if routed_experts.dtype == np.uint16:
+                    routed_experts = routed_experts.astype(np.int32)
                 routed_experts_start = routed_experts_payload["start"]
 
             return {

From 13e411d62c8009f1570b8fdbf73ceff2c1686a1d Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Wed, 10 Jun 2026 11:25:00 -0700
Subject: [PATCH 13/19] fix(routed_experts): preserve per-model dtype so batch
 packing stays consistent

---
 src/prime_rl/inference/vllm/routed_experts.py | 29 +++++++------------
 1 file changed, 11 insertions(+), 18 deletions(-)

diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py
index a3fce14211..0bbd625f1a 100644
--- a/src/prime_rl/inference/vllm/routed_experts.py
+++ b/src/prime_rl/inference/vllm/routed_experts.py
@@ -15,31 +15,24 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A
     array = np.asarray(routed_experts)
     assert array.ndim == 3
     assert np.issubdtype(array.dtype, np.integer)
-    # Narrow to the smallest int that holds the expert ids, matching vLLM's
-    # RoutedExpertsManager (uint8 for <=256 experts, uint16 otherwise). Keeps
-    # the wire payload compact while supporting >256-expert MoEs (e.g. Kimi-K2)
-    # that overflow uint8. The dtype rides the payload so the consumer decodes
-    # with the right element type.
     if array.size:
         assert array.min() >= 0
-        max_id = int(array.max())
+    # Preserve vLLM's per-MODEL dtype (RoutedExpertsManager uses uint8 for
+    # <=256 experts, uint16 otherwise), so every sample of a model shares one
+    # dtype. Re-narrowing per-sample on the observed max id would emit uint8
+    # for one sample and uint16 for another of the SAME >256-expert model,
+    # which the trainer's same-dtype routed-experts packing rejects. Only an
+    # unexpectedly wide capture dtype (e.g. an int64 buffer) is capped to int32
+    # (still consistent per model). The dtype rides the payload.
+    if array.dtype in (np.uint8, np.uint16, np.int16, np.int32):
+        compact = np.ascontiguousarray(array)
     else:
-        max_id = 0
-    if max_id <= np.iinfo(np.uint8).max:
-        target_dtype = np.uint8
-    elif max_id <= np.iinfo(np.uint16).max:
-        target_dtype = np.uint16
-    else:
-        # Beyond uint16 (>65535 experts) astype(uint16) would wrap and corrupt
-        # routing; fall back to int32 (consumer decodes via the dtype field).
-        target_dtype = np.int32
-
-    compact = np.ascontiguousarray(array.astype(target_dtype, copy=False))
+        compact = np.ascontiguousarray(array.astype(np.int32, copy=False))
     return {
         "data": pybase64.b64encode(memoryview(compact)).decode("ascii"),
         "shape": list(compact.shape),
         "start": start,
-        "dtype": np.dtype(target_dtype).name,
+        "dtype": compact.dtype.name,
     }
 
 

From 2c61937f9aac66b0aa8402a5faee76dd05c1c03b Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Wed, 10 Jun 2026 12:18:34 -0700
Subject: [PATCH 14/19] fix(k8s): drop trainer nvidia.com/gpu request when
 trainer DRA resourceClaim enabled

---
 k8s/prime-rl/templates/deployment.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml
index f40aa8046e..f3c01a9ab7 100644
--- a/k8s/prime-rl/templates/deployment.yaml
+++ b/k8s/prime-rl/templates/deployment.yaml
@@ -394,13 +394,13 @@ spec:
         {{- end }}
         resources:
           requests:
-            {{- if .Values.trainer.gpu.enabled }}
+            {{- if and .Values.trainer.gpu.enabled (not .Values.trainer.resourceClaim.enabled) }}
             nvidia.com/gpu: {{ .Values.trainer.gpu.count }}
             {{- end }}
             memory: {{ .Values.trainer.resources.requests.memory }}
             cpu: {{ .Values.trainer.resources.requests.cpu }}
           limits:
-            {{- if .Values.trainer.gpu.enabled }}
+            {{- if and .Values.trainer.gpu.enabled (not .Values.trainer.resourceClaim.enabled) }}
             nvidia.com/gpu: {{ .Values.trainer.gpu.count }}
             {{- end }}
             {{- if .Values.trainer.resources.limits }}

From 1b5917a83ce7bda48fc1a3c1ba3598192d05ab1a Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Wed, 10 Jun 2026 12:44:39 -0700
Subject: [PATCH 15/19] fix(k8s): set backend=dynamo in the dynamo-deploy
 client example

---
 k8s/dynamo-deploy/prime-rl-configs.yaml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/k8s/dynamo-deploy/prime-rl-configs.yaml b/k8s/dynamo-deploy/prime-rl-configs.yaml
index 9cbfe50af9..11a8cc8c53 100644
--- a/k8s/dynamo-deploy/prime-rl-configs.yaml
+++ b/k8s/dynamo-deploy/prime-rl-configs.yaml
@@ -36,6 +36,9 @@ data:
     hard_threshold = 0.0
 
     [client]
+    # Use the Dynamo admin/rollout transport (/engine/* + nvext), not the
+    # default vLLM frontend — required for the Dynamo /v1/rl/* admin routes below.
+    backend = "dynamo"
     base_url = ["http://prime-rl-dynamo-frontend.<your-namespace>.svc.cluster.local:8000/v1"]
     # Admin endpoints (/v1/rl/*) are served natively by the Dynamo Rust frontend
     # (DYN_ENABLE_RL=true). No separate admin-stub service needed.

From 6f4256055ce28083f55265c2eda7e678ed071937 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Thu, 11 Jun 2026 02:34:20 -0700
Subject: [PATCH 16/19] chore(deps): bump verifiers/renderers submodules to
 rl-sdk-4 routed_experts tips

---
 deps/renderers | 2 +-
 deps/verifiers | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/deps/renderers b/deps/renderers
index 5dbf4941f1..b62aabf944 160000
--- a/deps/renderers
+++ b/deps/renderers
@@ -1 +1 @@
-Subproject commit 5dbf4941f118a126c3be2031cf9d6b4b68cf23fe
+Subproject commit b62aabf9447d4ccfef92abaec83b17905986dafb
diff --git a/deps/verifiers b/deps/verifiers
index d713edc7ab..b31ff2d767 160000
--- a/deps/verifiers
+++ b/deps/verifiers
@@ -1 +1 @@
-Subproject commit d713edc7ab58b4f3f2ad79e410f09d2c4166042c
+Subproject commit b31ff2d767f482178ecb68ba73ff44a67ec1a7eb

From f8f42f82374e212843b5f7d49b5dc35aac47e295 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Thu, 11 Jun 2026 03:43:05 -0700
Subject: [PATCH 17/19] chore(deps): restore submodules and .gitmodules to
 upstream, dropping deps changes from PR

---
 .gitmodules                | 4 ++--
 deps/pydantic-config       | 2 +-
 deps/renderers             | 2 +-
 deps/research-environments | 2 +-
 deps/verifiers             | 2 +-
 5 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.gitmodules b/.gitmodules
index d139f9810f..2041f460ee 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,9 +1,9 @@
 [submodule "verifiers"]
 	path = deps/verifiers
-	url = git@github.com:biswapanda/verifiers.git
+	url = git@github.com:PrimeIntellect-ai/verifiers.git
 [submodule "renderers"]
 	path = deps/renderers
-	url = git@github.com:biswapanda/renderers.git
+	url = git@github.com:PrimeIntellect-ai/renderers.git
 [submodule "research-environments"]
 	path = deps/research-environments
 	url = git@github.com:PrimeIntellect-ai/research-environments.git
diff --git a/deps/pydantic-config b/deps/pydantic-config
index 8fc28bf218..896ade4e69 160000
--- a/deps/pydantic-config
+++ b/deps/pydantic-config
@@ -1 +1 @@
-Subproject commit 8fc28bf21828b20ddd6812d27b240daec484dc86
+Subproject commit 896ade4e69d8d8dff2d4b0a431b7e1c7c12d638f
diff --git a/deps/renderers b/deps/renderers
index b62aabf944..e6dba5ad6c 160000
--- a/deps/renderers
+++ b/deps/renderers
@@ -1 +1 @@
-Subproject commit b62aabf9447d4ccfef92abaec83b17905986dafb
+Subproject commit e6dba5ad6c50ca83d4ffa462145037082542e52a
diff --git a/deps/research-environments b/deps/research-environments
index 3c58236993..c752781984 160000
--- a/deps/research-environments
+++ b/deps/research-environments
@@ -1 +1 @@
-Subproject commit 3c58236993eef79897ec9c5552d36aa591d178cc
+Subproject commit c752781984c1b4fbb0a3d7f4aac1e7ed67cc749e
diff --git a/deps/verifiers b/deps/verifiers
index b31ff2d767..05c66c2358 160000
--- a/deps/verifiers
+++ b/deps/verifiers
@@ -1 +1 @@
-Subproject commit b31ff2d767f482178ecb68ba73ff44a67ec1a7eb
+Subproject commit 05c66c235875d785754f2b7078db0e7deeddbeae

From 8c837e6b7d81137c097358e81ea054271dedfe33 Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Thu, 11 Jun 2026 03:46:15 -0700
Subject: [PATCH 18/19] rm extra files

---
 Dockerfile.dynamo                 |  79 ------------
 k8s/dynamo-deploy/admin-stub.yaml |  73 -----------
 scripts/vllm-pr-39366.patch       | 206 ------------------------------
 tools/dynamo/admin_stub.py        |  56 --------
 4 files changed, 414 deletions(-)
 delete mode 100644 Dockerfile.dynamo
 delete mode 100644 k8s/dynamo-deploy/admin-stub.yaml
 delete mode 100644 scripts/vllm-pr-39366.patch
 delete mode 100644 tools/dynamo/admin_stub.py

diff --git a/Dockerfile.dynamo b/Dockerfile.dynamo
deleted file mode 100644
index 574976f49b..0000000000
--- a/Dockerfile.dynamo
+++ /dev/null
@@ -1,79 +0,0 @@
-# syntax=docker/dockerfile:1.4
-# Dockerfile.dynamo — layer ai-dynamo onto a prime-rl image WITHOUT reinstalling vLLM.
-#
-# The prime-rl base already ships vLLM 0.20.2 (+ vLLM PR #39366 two-phase pause),
-# torch, flashinfer, DeepGEMM. We build the dynamo Rust bindings (ai-dynamo-runtime,
-# via maturin) at DYNAMO_REF, then install the dynamo Python package — which provides
-# BOTH `dynamo.frontend` and `dynamo.vllm` (hatch packages = components/src/dynamo) —
-# with `--no-deps` so the base's vLLM / torch / transformers are NEVER touched.
-# A curated set of dynamo runtime deps (explicitly excluding vllm/torch/ray) is added
-# via `uv pip` (the prime-rl venv is uv-managed and has no `pip` binary).
-#
-# Build (BuildKit; run in the arm64 dind builder):
-#   DOCKER_BUILDKIT=1 docker build -f Dockerfile.dynamo \
-#     --build-arg BASE_IMAGE=nvcr.io/nvidian/dynamo-dev/biswa:prime-rl-97950abd-20260531-arm64 \
-#     --build-arg DYNAMO_REF=ecae3569926410ef33b4d3d13c7d6a1b89789bb0 \
-#     -t nvcr.io/nvidian/dynamo-dev/biswa:prime-rl-97950abd-dynamo-ecae3569-arm64 .
-#
-# DYNAMO_REF may be any commit/branch/tag on https://github.com/ai-dynamo/dynamo
-# (e.g. bis/rl-workers-discovery tip ecae3569…, or a release tag v1.2.0).
-
-ARG BASE_IMAGE=nvcr.io/nvidian/dynamo-dev/biswa:prime-rl-97950abd-20260531-arm64
-ARG DYNAMO_REPO=https://github.com/ai-dynamo/dynamo.git
-ARG DYNAMO_REF=ecae3569926410ef33b4d3d13c7d6a1b89789bb0
-ARG CARGO_BUILD_JOBS=16
-
-# ===== Stage 1: build ai-dynamo-runtime Rust bindings wheel =====
-FROM ubuntu:24.04 AS dynamo-builder
-ARG DYNAMO_REPO
-ARG DYNAMO_REF
-ARG CARGO_BUILD_JOBS
-ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS}
-ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update && apt-get install -y --no-install-recommends \
-      ca-certificates curl build-essential pkg-config libclang-dev protobuf-compiler git \
-      python3 python3-dev python3-venv \
- && rm -rf /var/lib/apt/lists/* \
- && curl -fsSL https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable
-ENV CARGO_HOME=/root/.cargo RUSTUP_HOME=/root/.rustup PATH=/root/.cargo/bin:${PATH}
-RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
-    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
-    cargo install maturin --locked
-RUN git clone "${DYNAMO_REPO}" /build/dynamo && cd /build/dynamo && git checkout "${DYNAMO_REF}"
-RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
-    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
-    --mount=type=cache,target=/build/dynamo/lib/bindings/python/target,sharing=locked \
-    cd /build/dynamo/lib/bindings/python \
- && maturin build --release --out /build/dist
-
-# ===== Stage 2: prime-rl base + dynamo (reuses base vLLM, no reinstall) =====
-FROM ${BASE_IMAGE}
-USER root
-ENV DYNAMO_HOME=/opt/dynamo
-COPY --from=dynamo-builder /build/dynamo /opt/dynamo
-COPY --from=dynamo-builder /build/dist/*.whl /tmp/dynamo-wheels/
-
-# 1) ai-dynamo-runtime (Rust bindings) + dynamo python pkg (frontend + vllm modules).
-#    --no-deps: do NOT pull vllm/torch/transformers (keep the prime-rl base's patched stack).
-RUN uv pip install --python /app/.venv/bin/python --no-cache /tmp/dynamo-wheels/*.whl \
- && cd /opt/dynamo \
- && uv pip install --python /app/.venv/bin/python --no-cache --no-deps -e . \
- && rm -rf /tmp/dynamo-wheels
-
-# 2) dynamo runtime deps the base may lack — EXPLICITLY excluding vllm / torch / ray.
-#    uvloop + nixl are required by the dynamo.vllm worker; the rest are frontend/runtime.
-RUN uv pip install --python /app/.venv/bin/python --no-cache \
-      uvloop "nixl[cu12]<=0.10.1" \
-      "fastapi==0.120.1" "uvicorn==0.38.0" httpx \
-      "msgspec>=0.19.0" pyzmq "prometheus_client>=0.23.1" \
-      "aiohttp>=3.9.0,<4.0" "blake3>=1.0.0,<2.0.0" \
-      "kubernetes>=32.0.1,<33.0.0" \
-      opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp
-
-# NOTE: etcd + nats-server are intentionally NOT installed in this image.
-# In the k8s deployment dynamo uses the external dynamo-platform services
-# (e.g. NATS_SERVER=nats://dynamo-platform-nats...:4222 and the platform etcd),
-# so shipping the static binaries in the worker image is unnecessary bloat.
-
-USER appuser
-WORKDIR /app
diff --git a/k8s/dynamo-deploy/admin-stub.yaml b/k8s/dynamo-deploy/admin-stub.yaml
deleted file mode 100644
index ceeb3e7a87..0000000000
--- a/k8s/dynamo-deploy/admin-stub.yaml
+++ /dev/null
@@ -1,73 +0,0 @@
-# Optional admin-stub Deployment + Service.
-# kubectl apply -f admin-stub.yaml -n <your-namespace>
-#
-# Only needed if your Dynamo build does NOT serve /v1/rl/* natively
-# (i.e. older builds without DYN_ENABLE_RL=true). With a recent Dynamo,
-# point `admin_base_url` directly at the Dynamo frontend and skip this
-# manifest entirely.
-apiVersion: v1
-kind: ConfigMap
-metadata:
-  name: admin-stub-script
-  namespace: <your-namespace>
-data:
-  admin_stub.py: |
-    from http.server import HTTPServer, BaseHTTPRequestHandler
-    class H(BaseHTTPRequestHandler):
-        def do_POST(self):
-            n = int(self.headers.get("Content-Length", 0))
-            body = self.rfile.read(n) if n else b""
-            print(f"[stub] POST {self.path} body={body[:200]}")
-            self.send_response(200)
-            self.end_headers()
-            self.wfile.write(b"OK")
-        def do_GET(self):
-            self.send_response(200)
-            self.end_headers()
-            self.wfile.write(b"OK")
-    HTTPServer(("0.0.0.0", 8001), H).serve_forever()
----
-apiVersion: apps/v1
-kind: Deployment
-metadata:
-  name: admin-stub
-  namespace: <your-namespace>
-spec:
-  replicas: 1
-  selector:
-    matchLabels:
-      app: admin-stub
-  template:
-    metadata:
-      labels:
-        app: admin-stub
-    spec:
-      containers:
-        - name: stub
-          image: python:3.12-slim
-          command: ["python3", "/scripts/admin_stub.py"]
-          ports:
-            - containerPort: 8001
-          volumeMounts:
-            - name: script
-              mountPath: /scripts
-          resources:
-            requests:
-              memory: "64Mi"
-              cpu: "50m"
-      volumes:
-        - name: script
-          configMap:
-            name: admin-stub-script
----
-apiVersion: v1
-kind: Service
-metadata:
-  name: admin-stub
-  namespace: <your-namespace>
-spec:
-  selector:
-    app: admin-stub
-  ports:
-    - port: 8001
-      targetPort: 8001
diff --git a/scripts/vllm-pr-39366.patch b/scripts/vllm-pr-39366.patch
deleted file mode 100644
index f8d973eff3..0000000000
--- a/scripts/vllm-pr-39366.patch
+++ /dev/null
@@ -1,206 +0,0 @@
-diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
-index 4f903eeefa6d..6ba392802e31 100644
---- a/vllm/config/parallel.py
-+++ b/vllm/config/parallel.py
-@@ -663,6 +663,33 @@ def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool:
-         aggregated_has_unfinished = bool(tensor.item())
-         return aggregated_has_unfinished
- 
-+    @staticmethod
-+    def sync_dp_state(
-+        dp_group: ProcessGroup, has_unfinished: bool, pending_pause: bool
-+    ) -> tuple[bool, bool]:
-+        """Combined all-reduce for DP state synchronization.
-+
-+        Uses a single SUM all-reduce on a 2-element tensor:
-+          [0] = 1 if this rank has unfinished work, else 0.
-+                SUM > 0 ≡ logical OR across ranks → any rank has work.
-+          [1] = 1 if this rank has a pending pause request, else 0.
-+                SUM == dp_size ≡ all ranks reached pause consensus.
-+
-+        has_unfinished_global is true if any rank has unfinished work,
-+        or if some ranks are waiting for a pause consensus.
-+
-+        Returns:
-+            (has_unfinished_global, pause_consensus)
-+        """
-+        tensor = torch.tensor(
-+            [int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu"
-+        )
-+        torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=dp_group)
-+        dp_size = dp_group.size()
-+        pause_count = tensor[1].item()
-+        has_unfinished_global = tensor[0].item() > 0 or pause_count % dp_size != 0
-+        return has_unfinished_global, pause_count == dp_size
-+
-     @staticmethod
-     def sync_kv_cache_memory_size(dp_group: ProcessGroup, kv_cache_memory: int) -> int:
-         if kv_cache_memory == -1:
-diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
-index 36864ba738bf..11c5ee19a664 100644
---- a/vllm/v1/engine/core.py
-+++ b/vllm/v1/engine/core.py
-@@ -1571,7 +1571,8 @@ def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
- 
-         pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
-         self.scheduler.set_pause_state(pause_state)
--        if not self.has_work():
-+
-+        if self._pause_complete():
-             if clear_cache:
-                 self._reset_caches()
-             return None
-@@ -1580,6 +1581,13 @@ def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
-         self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))
-         return future
- 
-+    def _pause_complete(self) -> bool:
-+        """Returns True if the pause has fully completed and the caller can
-+        return ``None`` synchronously; False if the pause is still pending
-+        and the caller should register an idle-state callback to finish it.
-+        """
-+        return not self.has_work()
-+
-     def _send_finish_outputs_to_client(
-         self, req_ids: list[str], client_index: int, finish_reason: FinishReason
-     ) -> None:
-@@ -1635,6 +1643,14 @@ def __init__(
-         self.current_wave = 0
-         self.last_counts = (0, 0)
- 
-+        # Two-phase pause protocol state. When pending_pause is True, the
-+        # engine keeps stepping (dummy batches) while waiting for all DP
-+        # ranks to also set pending_pause. Once all ranks agree via
-+        # all-reduce, ignore_start_dp_wave is set so that stale
-+        # START_DP_WAVE messages cannot re-wake the engines.
-+        self.pending_pause = False
-+        self.ignore_start_dp_wave = False
-+
-         from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
- 
-         self.eep_scaling_state: ElasticEPScalingState | None = None
-@@ -1664,6 +1680,7 @@ def _init_data_parallel(self, vllm_config: VllmConfig):
-         assert 0 <= local_dp_rank <= dp_rank < dp_size
- 
-         self.dp_rank = dp_rank
-+        self.dp_size = dp_size
-         dp_group, dp_store = parallel_config.stateless_init_dp_group(return_store=True)
-         self.dp_group, self.dp_store = dp_group, dp_store
- 
-@@ -1672,6 +1689,24 @@ def shutdown(self):
-         if dp_group := getattr(self, "dp_group", None):
-             stateless_destroy_torch_distributed_process_group(dp_group)
- 
-+    def _pause_complete(self) -> bool:
-+        """Two-phase DP-aware pause.
-+
-+        Phase 1: Set local pause state and ``pending_pause`` flag. If the
-+        engines are idle, kick-start them by setting ``engines_running`` to
-+        True so ranks enter the stepping loop and reach the all-reduce
-+        consensus checkpoint in ``_has_global_unfinished_reqs``.
-+
-+        Phase 2 (in ``_has_global_unfinished_reqs``): Once the all-reduce
-+        confirms that **all** ranks have ``pending_pause`` set, collectively
-+        stop stepping and set ``ignore_start_dp_wave`` so that stale
-+        ``START_DP_WAVE`` messages cannot re-wake any engine.
-+        """
-+        self.pending_pause = True
-+        self.engines_running = True
-+
-+        return False
-+
-     def add_request(self, request: Request, request_wave: int = 0):
-         super().add_request(request, request_wave)
-         if self.has_coordinator and request_wave != self.current_wave:
-@@ -1681,36 +1716,60 @@ def add_request(self, request: Request, request_wave: int = 0):
-                 not self.engines_running
-                 and self.scheduler.pause_state == PauseState.UNPAUSED
-             ):
--                self.engines_running = True
-                 # Request received for an already-completed wave, notify
-                 # front-end that we need to start the next one.
-+                self.engines_running = True
-                 self.output_queue.put_nowait(
-                     (-1, EngineCoreOutputs(start_wave=self.current_wave))
-                 )
- 
-     def resume_scheduler(self):
--        super().resume_scheduler()
--        if (
--            self.has_coordinator
--            and not self.engines_running
--            and self.scheduler.has_unfinished_requests()
--        ):
--            # Wake up other DP engines.
--            self.output_queue.put_nowait(
--                (-1, EngineCoreOutputs(start_wave=self.current_wave))
-+        if self.pending_pause or (self.engines_running and self.ignore_start_dp_wave):
-+            raise RuntimeError(
-+                "resume_scheduler called while pause is still in "
-+                "flight. Wait for the pause future to resolve before "
-+                "resuming."
-             )
-+        if self.engines_running:
-+            logger.debug("Resume called while engines are not paused, ignoring.")
-+            return
-+
-+        super().resume_scheduler()
-+        self.ignore_start_dp_wave = False
-+
-+        # Barrier: wait for all DP ranks to have resumed (and cleared
-+        # ignore_start_dp_wave) before any rank starts stepping. Uses
-+        # the existing all-reduce which is safe because engines are
-+        # stopped.
-+        has_global_unfinished = ParallelConfig.has_unfinished_dp(
-+            self.dp_group, self.scheduler.has_unfinished_requests()
-+        )
-+
-+        if has_global_unfinished:
-+            self.engines_running = True
-+
-+    def barrier(self):
-+        """Blocking barrier on the DP process group (test-only utility)."""
-+        import torch.distributed as dist
-+
-+        dist.barrier(group=self.dp_group)
- 
-     def _handle_client_request(
-         self, request_type: EngineCoreRequestType, request: Any
-     ) -> None:
-         if request_type == EngineCoreRequestType.START_DP_WAVE:
-+            if self.ignore_start_dp_wave:
-+                return
-             new_wave, exclude_eng_index = request
-             if exclude_eng_index != self.engine_index and (
-                 new_wave >= self.current_wave
-             ):
-                 self.current_wave = new_wave
-                 if not self.engines_running:
--                    logger.debug("EngineCore starting idle loop for wave %d.", new_wave)
-+                    logger.debug(
-+                        "EngineCore starting idle loop for wave %d.",
-+                        new_wave,
-+                    )
-                     self.engines_running = True
-         else:
-             super()._handle_client_request(request_type, request)
-@@ -1790,7 +1849,18 @@ def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
-         if self.step_counter % 32 != 0:
-             return True
- 
--        return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished)
-+        has_unfinished, pause_consensus = ParallelConfig.sync_dp_state(
-+            self.dp_group,
-+            has_unfinished=local_unfinished,
-+            pending_pause=self.pending_pause,
-+        )
-+
-+        if pause_consensus:
-+            self.ignore_start_dp_wave = True
-+            self.pending_pause = False
-+            logger.debug("DP pause consensus reached, ignoring START_DP_WAVE.")
-+
-+        return has_unfinished
- 
-     def reinitialize_distributed(
-         self, reconfig_request: ReconfigureDistributedRequest
diff --git a/tools/dynamo/admin_stub.py b/tools/dynamo/admin_stub.py
deleted file mode 100644
index 8a088c2e46..0000000000
--- a/tools/dynamo/admin_stub.py
+++ /dev/null
@@ -1,56 +0,0 @@
-#!/usr/bin/env python3
-"""
-Stub HTTP server for prime-rl admin endpoints.
-
-NOTE: As of dynamo#8630 (bis/parity-tokenize-tcp), Dynamo's Rust frontend
-implements these routes natively at /v1/rl/* when DYN_ENABLE_RL=true:
-  POST /v1/rl/load_lora_adapter
-  POST /v1/rl/unload_lora_adapter
-  GET  /v1/rl/health
-
-For K8s and any deployment with a real Dynamo frontend, point admin_base_url
-at the Dynamo service (e.g. http://<frontend-svc>:8000/v1/rl). This stub is
-kept as a local development fallback for running the orchestrator without a
-live Dynamo instance.
-
-Usage:
-    python tools/dynamo/admin_stub.py
-    python tools/dynamo/admin_stub.py --port 8001
-"""
-import argparse
-
-from aiohttp import web
-
-
-async def pause(request):
-    print("[stub] POST /pause - OK")
-    return web.Response(status=200, text="OK")
-
-
-async def resume(request):
-    print("[stub] POST /resume - OK")
-    return web.Response(status=200, text="OK")
-
-
-async def update_weights(request):
-    body = await request.json()
-    print(f"[stub] POST /update_weights weight_dir={body.get('weight_dir')} - OK (weights not reloaded)")
-    return web.Response(status=200, text="OK")
-
-
-async def health(request):
-    return web.Response(status=200, text="OK")
-
-
-app = web.Application()
-app.router.add_post("/pause", pause)
-app.router.add_post("/resume", resume)
-app.router.add_post("/update_weights", update_weights)
-app.router.add_get("/health", health)
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Dynamo admin stub server")
-    parser.add_argument("--port", type=int, default=8001, help="Port to listen on")
-    args = parser.parse_args()
-    print(f"[stub] Dynamo admin stub server starting on port {args.port}...")
-    web.run_app(app, port=args.port)

From 08bb4ea2c4e39db6d1b28913ff920872c70413dd Mon Sep 17 00:00:00 2001
From: Biswa Panda <biswa.panda@gmail.com>
Date: Thu, 11 Jun 2026 03:47:10 -0700
Subject: [PATCH 19/19] rm unnecessary files

---
 packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 1f740c2785..00a5325b7f 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -639,12 +639,6 @@ def validate_lora_adapter_saving(self):
                 )
         return self
 
-    # NOTE: HEAD's validate_weight_broadcast_type (nccl => async level 1) and
-    # validate_broadcast_keep_recent (keep_recent >= max_async_level) were dropped
-    # in the rl-sdk-4 merge: main removed trainer.max_async_level (the staleness
-    # window is now orchestrator.max_off_policy_steps). Re-add equivalent guards at
-    # the RLConfig level (which sees both trainer + orchestrator) if needed.
-
     @model_validator(mode="after")
     def validate_opt_and_fsdp_offload(self):
         if self.optim.type == "muon" and self.model.fsdp_cpu_offload: