From a827863353119038cbd53e1e8ec0ff5bcde68b57 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 6 Oct 2026 11:42:06 +0000
Subject: [PATCH 1/5] feat(inference): router replay with disaggregated P/D on
 vLLM 0.31

vLLM 0.31 captures routed experts through the AuxOutput connector and
rejects it with P/D connectors: a decode instance never computes a
remotely prefilled prompt, so its worker would materialize missing rows
(BlockObjectStoreError), and the first forward starts mid-block.

Carry our upstream AuxOutput fix as a plugin patch:
- the config check skips the P/D connector rejection for NixlConnector
  (alone or inside MultiConnector); every other check still runs;
- requests that arrive with do_remote_prefill emit rows only from their
  first forward (num_computed_tokens then); the floor survives preemption;
- the worker does not buffer or publish the partially known first block.

The pinned router drops as many leading decode rows as prefill returned,
so the decode server left-pads its rows to routed_experts_prompt_start
(the same payload size as 0.30's full-length decode rows). Remove the
config rejection added by the 0.31 bump.
---
 docs/inference.md                             |   4 +-
 .../src/prime_rl/configs/inference.py         |  10 --
 .../src/prime_rl/configs/rl.py                |  15 ---
 src/prime_rl/inference/patches.py             | 100 ++++++++++++++++++
 src/prime_rl/inference/vllm/routed_experts.py |  24 ++++-
 src/prime_rl/inference/vllm/serving_tokens.py |   1 +
 .../unit/inference/test_routed_experts_pd.py  |  59 +++++++++++
 tests/unit/test_configs.py                    |   5 +-
 8 files changed, 188 insertions(+), 30 deletions(-)
 create mode 100644 tests/unit/inference/test_routed_experts_pd.py

diff --git a/docs/inference.md b/docs/inference.md
index d99fc1f00c2..134b0ef0b53 100644
--- a/docs/inference.md
+++ b/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. The pinned router (v0.2.2) merges uint8 routing only, so P/D router replay does not support models with more than 256 experts until [router#61](https://github.com/PrimeIntellect-ai/router/pull/61) ships.
 
 ```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`.
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 c982b8495c3..928d41609ea 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -508,16 +508,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."""
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 c7349f7f30d..90b9efb33ac 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -594,21 +594,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 (
diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py
index 06d200a417f..fdffc4ca19e 100644
--- a/src/prime_rl/inference/patches.py
+++ b/src/prime_rl/inference/patches.py
@@ -18,6 +18,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()
@@ -183,6 +184,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.
 
diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py
index 5f73512fe5a..bc8e352a5e4 100644
--- a/src/prime_rl/inference/vllm/routed_experts.py
+++ b/src/prime_rl/inference/vllm/routed_experts.py
@@ -33,16 +33,36 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A
     }
 
 
+def pad_to_start(routed_experts: Any, num_rows: int) -> Any:
+    """Left-pad routing rows with zero rows up to ``num_rows``.
+
+    A P/D decode instance returns rows only from its first forward (the last prompt token). The
+    router (v0.2.2) drops as many leading decode rows as the prefill instance returned, so the
+    decode rows must start at ``routed_experts_prompt_start`` like the prefill rows. The padding
+    is never used: the router replaces it with the prefill rows.
+    """
+    if routed_experts is None or len(routed_experts) >= num_rows:
+        return routed_experts
+    array = np.asarray(routed_experts)
+    padding = np.zeros((num_rows - len(array), *array.shape[1:]), dtype=array.dtype)
+    return np.concatenate((padding, array))
+
+
 class RoutedExpertsCapture:
-    def __init__(self, generator: AsyncIterator[RequestOutput], start: int = 0):
+    def __init__(self, generator: AsyncIterator[RequestOutput], start: int = 0, remote_prefill: bool = False):
         self._generator = generator
         self._start = start
+        self._remote_prefill = remote_prefill
         self.routed_experts: dict[int, dict[str, Any]] = {}
 
     async def __aiter__(self):
         async for request_output in self._generator:
             for output in request_output.outputs:
-                encoded = serialize_routed_experts(getattr(output, "routed_experts", None), start=self._start)
+                routed_experts = getattr(output, "routed_experts", None)
+                if self._remote_prefill:
+                    num_rows = len(request_output.prompt_token_ids) + len(output.token_ids) - 1 - self._start
+                    routed_experts = pad_to_start(routed_experts, num_rows)
+                encoded = serialize_routed_experts(routed_experts, start=self._start)
                 if encoded is not None:
                     self.routed_experts[output.index] = encoded
             yield request_output
diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py
index f3a9b3a0b40..8febc135bc0 100644
--- a/src/prime_rl/inference/vllm/serving_tokens.py
+++ b/src/prime_rl/inference/vllm/serving_tokens.py
@@ -70,6 +70,7 @@ async def serve_tokens_full_generator(  # type: ignore[override]
             routed_experts = _GenerateRoutedExpertsCapture(
                 result_generator,
                 start=request.sampling_params.routed_experts_prompt_start,
+                remote_prefill=bool((request.kv_transfer_params or {}).get("do_remote_prefill")),
             )
             result_generator = routed_experts
 
diff --git a/tests/unit/inference/test_routed_experts_pd.py b/tests/unit/inference/test_routed_experts_pd.py
new file mode 100644
index 00000000000..20252af287d
--- /dev/null
+++ b/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])
diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py
index d08424a1f5e..f5356ead877 100644
--- a/tests/unit/test_configs.py
+++ b/tests/unit/test_configs.py
@@ -1011,11 +1011,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},
     )

From b1b2742c59eb9f3bf7391e15954b971cd835a100 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 12 Oct 2026 18:22:21 +0000
Subject: [PATCH 2/5] refactor(inference): send P/D decode routing rows from
 their first forward

Anchor the routing rows at the end (start = prompt + completion - 1 - rows), so
a P/D decode response no longer zero-pads its rows to the prompt start. The
router cuts decode rows by position since v0.2.3 (router#61).
---
 docs/inference.md                             |  2 +-
 src/prime_rl/inference/vllm/routed_experts.py | 31 +++++--------------
 src/prime_rl/inference/vllm/serving_tokens.py |  6 +---
 3 files changed, 9 insertions(+), 30 deletions(-)

diff --git a/docs/inference.md b/docs/inference.md
index 134b0ef0b53..5fa14c3437e 100644
--- a/docs/inference.md
+++ b/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). Under disaggregated P/D the router takes the prompt routing from the prefill instance and the completion routing from the decode instance. The pinned router (v0.2.2) merges uint8 routing only, so P/D router replay does not support models with more than 256 experts until [router#61](https://github.com/PrimeIntellect-ai/router/pull/61) ships.
+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. This needs router v0.2.3 or newer: decode responses carry routing rows only from their first forward, and the router cuts them by position.
 
 ```toml
 [trainer]
diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py
index bc8e352a5e4..b368799bcbb 100644
--- a/src/prime_rl/inference/vllm/routed_experts.py
+++ b/src/prime_rl/inference/vllm/routed_experts.py
@@ -33,36 +33,19 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A
     }
 
 
-def pad_to_start(routed_experts: Any, num_rows: int) -> Any:
-    """Left-pad routing rows with zero rows up to ``num_rows``.
-
-    A P/D decode instance returns rows only from its first forward (the last prompt token). The
-    router (v0.2.2) drops as many leading decode rows as the prefill instance returned, so the
-    decode rows must start at ``routed_experts_prompt_start`` like the prefill rows. The padding
-    is never used: the router replaces it with the prefill rows.
-    """
-    if routed_experts is None or len(routed_experts) >= num_rows:
-        return routed_experts
-    array = np.asarray(routed_experts)
-    padding = np.zeros((num_rows - len(array), *array.shape[1:]), dtype=array.dtype)
-    return np.concatenate((padding, array))
-
-
 class RoutedExpertsCapture:
-    def __init__(self, generator: AsyncIterator[RequestOutput], start: int = 0, remote_prefill: bool = False):
+    def __init__(self, generator: AsyncIterator[RequestOutput]):
         self._generator = generator
-        self._start = start
-        self._remote_prefill = remote_prefill
         self.routed_experts: dict[int, dict[str, Any]] = {}
 
     async def __aiter__(self):
         async for request_output in self._generator:
             for output in request_output.outputs:
                 routed_experts = getattr(output, "routed_experts", None)
-                if self._remote_prefill:
-                    num_rows = len(request_output.prompt_token_ids) + len(output.token_ids) - 1 - self._start
-                    routed_experts = pad_to_start(routed_experts, num_rows)
-                encoded = serialize_routed_experts(routed_experts, start=self._start)
-                if encoded is not None:
-                    self.routed_experts[output.index] = encoded
+                if routed_experts is None:
+                    continue
+                # Rows end at the last forwarded token. A P/D decode instance's rows begin at its first
+                # forward (the last prompt token); the router cuts them by this start (router >= v0.2.3).
+                start = len(request_output.prompt_token_ids) + len(output.token_ids) - 1 - len(routed_experts)
+                self.routed_experts[output.index] = serialize_routed_experts(routed_experts, start=start)
             yield request_output
diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py
index 8febc135bc0..95b25add1bc 100644
--- a/src/prime_rl/inference/vllm/serving_tokens.py
+++ b/src/prime_rl/inference/vllm/serving_tokens.py
@@ -67,11 +67,7 @@ async def serve_tokens_full_generator(  # type: ignore[override]
     ) -> ErrorResponse | GenerateResponse:
         routed_experts: _GenerateRoutedExpertsCapture | None = None
         if self.engine_client.vllm_config.aux_output_config.enable_return_routed_experts:
-            routed_experts = _GenerateRoutedExpertsCapture(
-                result_generator,
-                start=request.sampling_params.routed_experts_prompt_start,
-                remote_prefill=bool((request.kv_transfer_params or {}).get("do_remote_prefill")),
-            )
+            routed_experts = _GenerateRoutedExpertsCapture(result_generator)
             result_generator = routed_experts
 
         response = await super().serve_tokens_full_generator(

From 978146cf52c97ec0ef3094e9737cd30a48fbe9cd Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 13 Oct 2026 10:15:48 +0000
Subject: [PATCH 3/5] refactor(inference): take the end-anchored routing start
 from the packed /generate path

The start moves into _PackedOutputs (#3900), so routed_experts.py and
serving_tokens.py stay as on main.
---
 src/prime_rl/inference/vllm/routed_experts.py | 13 +++++--------
 src/prime_rl/inference/vllm/serving_tokens.py |  5 ++++-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py
index b368799bcbb..5f73512fe5a 100644
--- a/src/prime_rl/inference/vllm/routed_experts.py
+++ b/src/prime_rl/inference/vllm/routed_experts.py
@@ -34,18 +34,15 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A
 
 
 class RoutedExpertsCapture:
-    def __init__(self, generator: AsyncIterator[RequestOutput]):
+    def __init__(self, generator: AsyncIterator[RequestOutput], start: int = 0):
         self._generator = generator
+        self._start = start
         self.routed_experts: dict[int, dict[str, Any]] = {}
 
     async def __aiter__(self):
         async for request_output in self._generator:
             for output in request_output.outputs:
-                routed_experts = getattr(output, "routed_experts", None)
-                if routed_experts is None:
-                    continue
-                # Rows end at the last forwarded token. A P/D decode instance's rows begin at its first
-                # forward (the last prompt token); the router cuts them by this start (router >= v0.2.3).
-                start = len(request_output.prompt_token_ids) + len(output.token_ids) - 1 - len(routed_experts)
-                self.routed_experts[output.index] = serialize_routed_experts(routed_experts, start=start)
+                encoded = serialize_routed_experts(getattr(output, "routed_experts", None), start=self._start)
+                if encoded is not None:
+                    self.routed_experts[output.index] = encoded
             yield request_output
diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py
index 95b25add1bc..f3a9b3a0b40 100644
--- a/src/prime_rl/inference/vllm/serving_tokens.py
+++ b/src/prime_rl/inference/vllm/serving_tokens.py
@@ -67,7 +67,10 @@ async def serve_tokens_full_generator(  # type: ignore[override]
     ) -> ErrorResponse | GenerateResponse:
         routed_experts: _GenerateRoutedExpertsCapture | None = None
         if self.engine_client.vllm_config.aux_output_config.enable_return_routed_experts:
-            routed_experts = _GenerateRoutedExpertsCapture(result_generator)
+            routed_experts = _GenerateRoutedExpertsCapture(
+                result_generator,
+                start=request.sampling_params.routed_experts_prompt_start,
+            )
             result_generator = routed_experts
 
         response = await super().serve_tokens_full_generator(

From efab8d0ca224ee7f76d393ff59084b712033aa0a Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 14:12:39 +0000
Subject: [PATCH 4/5] chore(deps): pin vllm-router to v0.2.3

---
 pyproject.toml |  4 ++--
 uv.lock        | 24 ++++++++++++------------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 727780269d0..91ba6770981 100644
--- a/pyproject.toml
+++ b/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'" },
diff --git a/uv.lock b/uv.lock
index b95756c9f43..5e9e521cf7a 100644
--- a/uv.lock
+++ b/uv.lock
@@ -5050,8 +5050,8 @@ all = [
     { name = "torchvision", marker = "sys_platform == 'linux'" },
     { name = "vllm", version = "0.31.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.31.0/vllm-0.31.0-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "vllm", version = "0.31.0", source = { 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' and sys_platform == 'linux'" },
-    { name = "vllm-router", version = "0.2.2", source = { 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 != 'x86_64' and sys_platform == 'linux'" },
-    { name = "vllm-router", version = "0.2.2", source = { 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' and sys_platform == 'linux'" },
+    { name = "vllm-router", version = "0.2.3", source = { 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 != 'x86_64' and sys_platform == 'linux'" },
+    { name = "vllm-router", version = "0.2.3", source = { 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' and sys_platform == 'linux'" },
 ]
 dashboard = [
     { name = "fastapi" },
@@ -5063,8 +5063,8 @@ disagg = [
     { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "nixl", marker = "sys_platform == 'linux'" },
-    { name = "vllm-router", version = "0.2.2", source = { 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 != 'x86_64' and sys_platform == 'linux'" },
-    { name = "vllm-router", version = "0.2.2", source = { 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' and sys_platform == 'linux'" },
+    { name = "vllm-router", version = "0.2.3", source = { 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 != 'x86_64' and sys_platform == 'linux'" },
+    { name = "vllm-router", version = "0.2.3", source = { 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' and sys_platform == 'linux'" },
 ]
 dynamo = [
     { name = "ai-dynamo", marker = "sys_platform == 'linux'" },
@@ -5202,9 +5202,9 @@ requires-dist = [
     { name = "vllm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/vllm-project/vllm/releases/download/v0.31.0/vllm-0.31.0-cp38-abi3-manylinux_2_28_aarch64.whl" },
     { name = "vllm", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.31.0" },
     { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/vllm-project/vllm/releases/download/v0.31.0/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl" },
-    { name = "vllm-router", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.2/vllm_router-0.2.2-cp38-abi3-manylinux_2_28_aarch64.whl" },
+    { name = "vllm-router", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl" },
     { name = "vllm-router", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'disagg'" },
-    { name = "vllm-router", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", 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" },
+    { name = "vllm-router", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", 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" },
     { name = "wandb", specifier = ">=0.26.1" },
     { name = "wandb-workspaces", specifier = ">=0.4.3" },
 ]
@@ -8233,8 +8233,8 @@ provides-extras = ["zen", "bench", "tensorizer", "fastsafetensors", "instanttens
 
 [[package]]
 name = "vllm-router"
-version = "0.2.2"
-source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.2/vllm_router-0.2.2-cp38-abi3-manylinux_2_28_aarch64.whl" }
+version = "0.2.3"
+source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
@@ -8247,7 +8247,7 @@ dependencies = [
     { name = "uvicorn" },
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.2/vllm_router-0.2.2-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4e12c3e9e8123556b261b2ed79998ad256698b7e64da5afe4fe69e733e49b2d3" },
+    { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0f1b0ab898cee9ced6a68f45de7e68280711bb99c404b62e581c87b6890342d0" },
 ]
 
 [package.metadata]
@@ -8266,8 +8266,8 @@ provides-extras = ["dev"]
 
 [[package]]
 name = "vllm-router"
-version = "0.2.2"
-source = { 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" }
+version = "0.2.3"
+source = { 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" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
@@ -8280,7 +8280,7 @@ dependencies = [
     { name = "uvicorn" },
 ]
 wheels = [
-    { 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", hash = "sha256:2132bda10632efc5a726c69c7e3e7a18ea70731c5ca1ec2528227aa0fb634eb8" },
+    { 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", hash = "sha256:972faf118888b073214c7880d69b81895cb4d7f4402b5445f842eaf88a89f708" },
 ]
 
 [package.metadata]

From e586ef860e6d679481ca89d275430c3edb8fc93a Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 17:22:52 +0100
Subject: [PATCH 5/5] Update inference.md

Signed-off-by: faresobeid <111092724+faresobeid@users.noreply.github.com>
---
 docs/inference.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/inference.md b/docs/inference.md
index 5fa14c3437e..e15658e267b 100644
--- a/docs/inference.md
+++ b/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). Under disaggregated P/D the router takes the prompt routing from the prefill instance and the completion routing from the decode instance. This needs router v0.2.3 or newer: decode responses carry routing rows only from their first forward, and the router cuts them by position.
+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]