feat(inference): router replay on the V2 runner everywhere

docs/inference.md

@@ -287,7 +287,7 @@ On the CLI the same keys are available as `--inference.vllm.max-num-seqs 256` (o
 
 Router replay works by capturing the expert routing decisions into a buffer. This buffer then gets sent to the trainer, which can use it instead of re-computing the routing. This lowers the trainer↔inference mismatch by an order of magnitude, resulting in more stable training.
 
-To enable router replay, you can set `inference.vllm.enable_return_routed_experts = true`. Capture runs on vLLM's V2 model runner (forced automatically). Router replay is not supported with disaggregated P/D: vLLM does not capture routed experts across NIXL KV transfer.
+To enable router replay, you can set `inference.vllm.enable_return_routed_experts = true`. Capture runs on vLLM's V2 model runner (forced automatically). Under disaggregated P/D the router takes the prompt routing from the prefill instance and the completion routing from the decode instance. 
 
 ```toml
 [trainer]
@@ -313,6 +313,6 @@ top_k = 512   # optional, defaults to 512 under truncation (bounds each sampling
 
 That's all — there are no replay flags. Truncated train sampling makes the inference server return sampling masks (`inference.enable_return_sampling_mask`, auto-enabled), and the trainer replays them. Train-sampling `top_k` above 512 is rejected: the trainer pads each micro batch to its largest sampling mask, so the bound caps trainer memory. Configs that would break under renormalized logprobs are rejected: `opd`/`opsd` and truncation knobs smuggled via `extra_body`. Frozen-source envs are exempt.
 
-Capture runs on vLLM's V2 model runner (forced automatically) and is engine-wide. It can run together with router replay. While sampling capture is on, vLLM rejects any request with `temperature <= 0` or without an effective `top_k > 0` — eval sampling against the training server must set `top_k` (the model's generation config often supplies one) and a non-zero temperature.
+Capture runs on vLLM's V2 model runner (forced automatically) and is engine-wide. It can run together with router replay, including under disaggregated P/D. While sampling capture is on, vLLM rejects any request with `temperature <= 0` or without an effective `top_k > 0` — eval sampling against the training server must set `top_k` (the model's generation config often supplies one) and a non-zero temperature.
 
 When launching the inference server standalone, set `inference.enable_return_sampling_mask = true` yourself; clients must sample with `0 < top_k <= 512`.

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

@@ -477,16 +477,6 @@ def validate_llmd_no_routed_experts(self):
             )
         return self
 
-    @model_validator(mode="after")
-    def validate_disaggregated_no_routed_experts(self):
-        """vLLM rejects routed-expert capture with NIXL KV transfer."""
-        if self.deployment.type == "disaggregated" and self.vllm.enable_return_routed_experts:
-            raise ValueError(
-                "Routed-expert return (enable_return_routed_experts) is not supported with disaggregated P/D: "
-                "vLLM does not capture routed experts across NIXL KV transfer."
-            )
-        return self
-
     @model_validator(mode="after")
     def validate_router_deployment(self):
         """The llm-d router (EPP + Envoy) is launched by the SLURM templates only; multi-node deployments need a router to front the per-rank engines."""

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

@@ -561,21 +561,6 @@ def auto_setup_sampling_mask_capture(self):
             )
         return self
 
-    @model_validator(mode="after")
-    def validate_disaggregated_no_routed_experts(self):
-        """Runs after ``auto_setup_router_replay``, which sets the inference flag after InferenceConfig's validators."""
-        inference = self.inference
-        if (
-            inference is not None
-            and inference.deployment.type == "disaggregated"
-            and inference.vllm.enable_return_routed_experts
-        ):
-            raise ValueError(
-                "Router replay (inference.vllm.enable_return_routed_experts / trainer.enable_router_replay) "
-                "is not supported with disaggregated P/D: vLLM does not capture routed experts across NIXL KV transfer."
-            )
-        return self
-
     @model_validator(mode="after")
     def validate_router_replay_without_kv_offload(self):
         if (

pyproject.toml

@@ -276,8 +276,8 @@ sonic-moe = { git = "https://github.com/S1ro1/sonic-moe", rev = "0bf06ac" }
 quack-kernels = { git = "https://github.com/S1ro1/quack", rev = "f62009f" }
 flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", subdirectory = "flash_attn/cute", rev = "0251105" }
 vllm-router = [
-    { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.2/vllm_router-0.2.2-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.2/vllm_router-0.2.2-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 vllm = [
     { url = "https://github.com/vllm-project/vllm/releases/download/v0.31.0/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },

src/prime_rl/inference/patches.py

@@ -21,6 +21,7 @@ def apply_shared_vllm_patches():
     _patch_qwen35_moe_lora_format()
     monkey_patch_nano_v3_reasoning_parser()
     monkey_patch_minimax_m2_think_end_passthrough()
+    monkey_patch_routed_experts_with_pd_connectors()
     monkey_patch_kv_xfer_finished_tolerate_freed()
     monkey_patch_online_fp8_parameter_cast()
     monkey_patch_deepseek_v4_allowed_layer_types()
@@ -380,6 +381,105 @@ def _patched_config():
     minimax_m2.minimax_m2_config = _patched_config
 
 
+def monkey_patch_routed_experts_with_pd_connectors():
+    """Let vLLM's routed-experts capture (AuxOutput connector) run on NIXL P/D instances.
+
+    vLLM 0.31 rejects routed-experts capture with P/D connectors, because a decode instance never
+    computes a remotely prefilled prompt, so it has no rows for it. This carries our upstream fix
+    (AuxOutput patch 02): a request on a decode instance emits rows only from its first forward on
+    this instance, and the worker does not buffer the
+    partially known first block. The router takes the prompt rows from the prefill instance. Delete
+    once vLLM supports routed experts with P/D connectors.
+    """
+    from vllm.config.vllm import VllmConfig
+    from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
+    from vllm.distributed.aux_output_connector.worker import AuxOutputWorkerConnector
+
+    original_verify = VllmConfig._verify_aux_output_compatibility
+    if getattr(original_verify, "_prime_rl_allows_pd", False):
+        return
+
+    def _verify_aux_output_compatibility(config):
+        kv_transfer_config = config.kv_transfer_config
+        if kv_transfer_config is None or not kv_transfer_config.has_connector("NixlConnector"):
+            return original_verify(config)
+        # The P/D connector check is the last one; every other check still runs.
+        config.kv_transfer_config = None
+        try:
+            return original_verify(config)
+        finally:
+            config.kv_transfer_config = kv_transfer_config
+
+    original_build_connector_meta = AuxOutputSchedulerConnector.build_connector_meta
+
+    def build_connector_meta(self, scheduler_output, requests):
+        metadata = original_build_connector_meta(self, scheduler_output, requests)
+        floors = _emit_floors(self)
+        for request_id, emit_start in metadata.requests.items():
+            request = requests[request_id]
+            if not _is_pd_decode_request(request):
+                continue
+            # Fixed at the first forward (built before the step advances num_computed_tokens).
+            # Normally prompt_len - 1; lower after a failed KV load, where this instance
+            # recomputes those rows itself.
+            floor = floors.setdefault(request_id, request.num_computed_tokens)
+            metadata.requests[request_id] = max(emit_start, floor)
+        return metadata
+
+    original_request_finished = AuxOutputSchedulerConnector.request_finished
+
+    def request_finished(self, request):
+        # Preemption keeps the floor: a resumed request may hit prompt blocks loaded from the prefill
+        # instance, which have no rows here.
+        if request.is_finished():
+            _emit_floors(self).pop(request.request_id, None)
+        return original_request_finished(self, request)
+
+    original_worker_init = AuxOutputWorkerConnector.__init__
+
+    def worker_init(self, *args, **kwargs):
+        original_worker_init(self, *args, **kwargs)
+        if self._buffer is not None:
+            skip_remote_prefill_partial_block(self)
+
+    _verify_aux_output_compatibility._prime_rl_allows_pd = True
+    VllmConfig._verify_aux_output_compatibility = _verify_aux_output_compatibility
+    AuxOutputSchedulerConnector.build_connector_meta = build_connector_meta
+    AuxOutputSchedulerConnector.request_finished = request_finished
+    AuxOutputWorkerConnector.__init__ = worker_init
+
+
+def skip_remote_prefill_partial_block(worker) -> None:
+    """Make an AuxOutput worker skip the mid-block start of a remotely prefilled request's first forward."""
+    from vllm.utils.math_utils import cdiv
+
+    buffer = worker._buffer
+    capture = buffer.capture
+    block_size = buffer.block_size
+
+    def _capture(request_id, token_start, rows):
+        state = worker._requests[request_id]
+        if state.capture_cursor is None and token_start % block_size and state.emit_cursor >= token_start:
+            # The rows before this token exist only on the prefill instance, so this block is never
+            # published. Buffer from the next block; this step's rows are still emitted directly.
+            state.prime_rl_buffer_start = cdiv(token_start, block_size) * block_size
+        skip = min(max(getattr(state, "prime_rl_buffer_start", 0) - token_start, 0), len(rows))
+        return capture(request_id, token_start + skip, rows[skip:])
+
+    buffer.capture = _capture
+
+
+def _emit_floors(connector) -> dict[str, int]:
+    """First output token of each remotely prefilled request on this instance."""
+    return connector.__dict__.setdefault("_prime_rl_emit_floors", {})
+
+
+def _is_pd_decode_request(request) -> bool:
+    """A request on a P/D decode instance (NIXL clears ``do_remote_prefill`` once scheduled, not ``remote_engine_id``)."""
+    params = request.kv_transfer_params or {}
+    return bool(params.get("remote_engine_id")) and not params.get("do_remote_decode")
+
+
 def monkey_patch_strip_routed_experts_from_chat():
     """Drop routed_experts from chat-completions responses.
 

tests/unit/inference/test_routed_experts_pd.py

@@ -0,0 +1,59 @@
+from types import SimpleNamespace
+
+import numpy as np
+
+from prime_rl.inference.patches import (
+    _emit_floors,
+    monkey_patch_routed_experts_with_pd_connectors,
+    skip_remote_prefill_partial_block,
+)
+
+
+def test_remote_prefill_rows_start_at_first_forward():
+    from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
+
+    monkey_patch_routed_experts_with_pd_connectors()
+    connector = AuxOutputSchedulerConnector()
+    request = SimpleNamespace(
+        request_id="r",
+        block_hashes=[b"a" * 32],
+        num_tokens=6,
+        num_output_tokens=0,
+        num_computed_tokens=5,  # tokens [0, 5) were loaded from the prefill instance
+        sampling_params=SimpleNamespace(routed_experts_prompt_start=0),
+        kv_transfer_params={"do_remote_prefill": False, "remote_engine_id": "p0"},
+        is_finished=lambda: False,
+    )
+
+    def emit_start():
+        step = SimpleNamespace(num_scheduled_tokens={"r": request.num_tokens - request.num_computed_tokens})
+        return connector.build_connector_meta(step, {"r": request}).requests["r"]
+
+    assert emit_start() == 5
+    # A preempted request recomputes from scratch but still emits from its first forward.
+    connector.request_finished(request)
+    request.num_computed_tokens = 0
+    assert emit_start() == 5
+    request.is_finished = lambda: True
+    connector.request_finished(request)
+    assert not _emit_floors(connector)
+
+
+def test_remote_prefill_partial_first_block_is_not_buffered():
+    from vllm.distributed.aux_output_connector.routed_experts import RoutedExpertsBuffer
+    from vllm.distributed.aux_output_connector.worker import _WorkerRequestState
+
+    rows = np.arange(12 * 3 * 2, dtype=np.uint8).reshape(12, 3, 2)
+    buffer = RoutedExpertsBuffer(np.dtype("uint8"), (3, 2), 4, 1, 8, 2)
+    state = _WorkerRequestState(emit_cursor=6)
+    worker = SimpleNamespace(_buffer=buffer, _requests={"r": state})
+    skip_remote_prefill_partial_block(worker)
+
+    # The first forward starts mid-block at token 6; without the skip the buffer asserts.
+    assert buffer.capture("r", 6, rows[6:7]) == []
+    state.capture_cursor = 7
+    assert buffer.capture("r", 7, rows[7:9]) == []
+    state.capture_cursor = 9
+    [(start, block)] = buffer.capture("r", 9, rows[9:12])
+    assert start == 8
+    np.testing.assert_array_equal(block, rows[8:12])

tests/unit/test_configs.py

@@ -1162,11 +1162,14 @@ def test_explicit_inference_parser_wins_over_auto():
     assert config.inference.vllm.tool_call_parser == "hermes"
 
 
-def test_combined_replay_uses_v2_runner(monkeypatch):
+@pytest.mark.parametrize("deployment", [{"type": "single_node"}, {"type": "disaggregated"}])
+def test_combined_replay_uses_v2_runner(monkeypatch, deployment):
     from prime_rl.inference.server import setup_vllm_env
 
     monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
     config = InferenceConfig(
+        deployment=deployment,
+        slurm={} if deployment["type"] == "disaggregated" else None,
         enable_return_sampling_mask=True,
         vllm={"enable_return_routed_experts": True},
     )