[Feature] Routing replay (R3) for vLLM rollout (vLLM 0.22+, /inference/v1/generate)

slime/backends/vllm_utils/vllm_engine.py

@@ -405,6 +405,51 @@ def _redact_cmd_for_log(cmd: list[str]) -> str:
     return " ".join(parts)
 
 
+def _verify_generate_routed_experts(base_url: str, model: str, timeout_s: float = 120.0) -> None:
+    """Smoke-check MoE routing via ``/inference/v1/generate`` (same path as R3 rollout; vLLM 0.22+)."""
+    from slime.rollout.vllm_rollout import _decode_generate_routed_experts
+
+    base = base_url.rstrip("/")
+    token_ids = [1, 2, 3, 4, 5]
+    payload = {
+        "model": model,
+        "token_ids": token_ids,
+        "sampling_params": {
+            "max_tokens": 8,
+            "temperature": 0.0,
+            "logprobs": 1,
+        },
+    }
+    response = requests.post(f"{base}/inference/v1/generate", json=payload, timeout=timeout_s)
+    response.raise_for_status()
+    body = response.json()
+    choice = (body.get("choices") or [{}])[0]
+    out_ids = choice.get("token_ids") or []
+    num_gen = len(out_ids)
+    num_prompt = len(token_ids)
+
+    arr = _decode_generate_routed_experts(choice)
+    if arr is None:
+        raise RuntimeError(
+            "vLLM /inference/v1/generate returned no routed-experts fields. "
+            "Requires vLLM 0.22+ with --enable-return-routed-experts (use_rollout_routing_replay)."
+        )
+
+    expected_rows = num_prompt + num_gen - 1 if num_prompt > 0 and num_gen > 0 else 0
+    n_rows = int(arr.shape[0])
+    if expected_rows > 0 and n_rows != expected_rows:
+        raise RuntimeError(
+            f"vLLM routing replay smoke check: routing rows {n_rows} != "
+            f"expected {expected_rows} (prompt+gen len(tokens)-1). "
+            "Requires vLLM 0.22+ (PR #39568)."
+        )
+    logger.info(
+        "vLLM routing replay smoke check OK (/inference/v1/generate): routing rows=%s " "(expected %s)",
+        n_rows,
+        expected_rows,
+    )
+
+
 def _wait_server_healthy(base_url: str, process: multiprocessing.Process | None, timeout_s: float = 300.0) -> None:
     """Wait until the vLLM server responds on ``GET /health`` (SGLang stacks typically use ``GET /health_generate``)."""
     start = time.time()
@@ -451,8 +496,18 @@ def __init__(
         self.node_rank = 0
 
     def _http_base(self) -> str:
+        if not hasattr(self, "server_host") or not hasattr(self, "server_port"):
+            raise RuntimeError("Call init() before using the vLLM HTTP client.")
         return f"http://{self.server_host}:{self.server_port}"
 
+    def _weight_transfer_http_timeout(self) -> float:
+        return float(
+            os.environ.get(
+                "SLIME_VLLM_WEIGHT_TRANSFER_UPDATE_TIMEOUT_SEC",
+                os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"),
+            )
+        )
+
     def init(
         self,
         dist_init_addr,
@@ -559,7 +614,10 @@ def _init_normal(self) -> None:
             visible_devices=visible_devices,
             model_path=self.model_path,
         )
-        _wait_server_healthy(self._http_base(), process=self.process)
+        base = self._http_base()
+        _wait_server_healthy(base, process=self.process)
+        if getattr(self.args, "use_rollout_routing_replay", False):
+            _verify_generate_routed_experts(base, self.model_path)
 
     def _post_json(self, endpoint: str, payload: dict, timeout: float) -> requests.Response:
         url = f"{self._http_base()}/{endpoint.lstrip('/')}"
@@ -571,13 +629,11 @@ def _post_vllm_update_weights_http(self, update_info: dict) -> dict:
         Caller must invoke ``start_weight_update`` / ``finish_weight_update`` around a batch of
         ``/update_weights`` calls (see ``UpdateWeightFromTensor`` / ``UpdateWeightFromDistributed``).
         """
-        timeout_s = float(
-            os.environ.get(
-                "SLIME_VLLM_WEIGHT_TRANSFER_UPDATE_TIMEOUT_SEC",
-                os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"),
-            )
+        response = self._post_json(
+            "update_weights",
+            {"update_info": update_info},
+            timeout=self._weight_transfer_http_timeout(),
         )
-        response = self._post_json("update_weights", {"update_info": update_info}, timeout=timeout_s)
         response.raise_for_status()
         try:
             return response.json()
@@ -796,11 +852,10 @@ def init_weight_transfer_engine(self, payload: dict) -> dict:
 
     def start_weight_update(self, is_checkpoint_format: bool = False) -> dict:
         """``POST /start_weight_update`` — signals vLLM to enter IPC weight-update mode."""
-        update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"))
         response = self._post_json(
             "start_weight_update",
             {"is_checkpoint_format": is_checkpoint_format},
-            timeout=update_timeout_s,
+            timeout=self._weight_transfer_http_timeout(),
         )
         response.raise_for_status()
         try:
@@ -822,8 +877,7 @@ def finish_weight_update(self, weight_version: str | None = None) -> dict:
         (e.g. ``/root/models/Qwen2.5-0.5B-Instruct``), never matching the
         updater's integer version (``"1"``, ``"2"``, …).
         """
-        update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"))
-        response = self._post_json("finish_weight_update", {}, timeout=update_timeout_s)
+        response = self._post_json("finish_weight_update", {}, timeout=self._weight_transfer_http_timeout())
         response.raise_for_status()
         # Record the new version only after the POST succeeded — if the engine
         # never actually exited weight-update mode, ``_weight_version`` must not

slime/ray/rollout.py

@@ -774,8 +774,16 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
         if samples[0].rollout_log_probs is not None:
             train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples]
 
-        if samples[0].rollout_routed_experts is not None:
-            train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples]
+        if getattr(self.args, "use_rollout_routing_replay", False):
+            routed = [sample.rollout_routed_experts for sample in samples]
+            missing = [i for i, r in enumerate(routed) if r is None]
+            if missing:
+                raise ValueError(
+                    f"use_rollout_routing_replay: {len(missing)}/{len(samples)} samples missing "
+                    "rollout_routed_experts (see rollout logs for vLLM routing replay errors). "
+                    "Ensure vLLM 0.22+ serves with --enable-return-routed-experts."
+                )
+            train_data["rollout_routed_experts"] = routed
 
         if samples[0].train_metadata is not None:
             train_data["metadata"] = [sample.train_metadata for sample in samples]

slime/rollout/vllm_rollout.py

@@ -161,49 +161,70 @@ def _vllm_meta_from_generate_choice(args: Namespace, choice: dict, usage: dict |
 
 
 def _decode_vllm_routed_experts(value: str) -> np.ndarray:
+    """Decode vLLM routed-experts field returned as base64 ``.npy`` bytes."""
     raw = base64.b64decode(value.encode("ascii"), validate=True)
     return np.load(io.BytesIO(raw), allow_pickle=False)
 
 
 def _apply_vllm_routed_experts(
     args: Namespace,
     sample: Sample,
-    _output: dict,
     choice: dict,
 ) -> None:
-    """Populate ``sample.rollout_routed_experts`` from vLLM ``choices[].routed_experts`` when enabled.
+    """Populate ``sample.rollout_routed_experts`` from vLLM ``/inference/v1/generate`` (R3 only).
 
-    vLLM ``/inference/v1/generate`` returns routed experts as a base64 encoded
-    ``.npy`` payload on each response choice when the server is launched with
-    ``--enable-return-routed-experts``.
+    vLLM's contract is a single base64 `.npy` buffer on `choices[].routed_experts` with decoded
+    shape `(len(tokens) - 1, num_layers, top_k)`.
     """
     if not getattr(args, "use_rollout_routing_replay", False):
         return
-    gen_re = choice.get("routed_experts")
-    if gen_re is None:
+
+    routed = choice.get("routed_experts")
+    if sample.status == Sample.Status.ABORTED and sample.response_length == 0:
         return
-    arr = _decode_vllm_routed_experts(gen_re)
-    n_tok = len(sample.tokens)
-    expected_rows = max(0, n_tok - 1)
+    if routed is None:
+        raise RuntimeError(
+            "vLLM routing replay: missing choices[0].routed_experts on /inference/v1/generate response. "
+            "Check vLLM 0.22+ was launched with --enable-return-routed-experts."
+        )
+    if not isinstance(routed, str):
+        raise RuntimeError(
+            f"vLLM routing replay: choices[0].routed_experts must be base64 npy str, got {type(routed)}"
+        )
+
+    arr = _decode_vllm_routed_experts(routed)
     if arr.ndim != 3:
-        logger.warning(f"Unexpected routed_experts ndim={arr.ndim} shape={arr.shape}")
-        return
-    if arr.shape[0] == n_tok:
-        arr = arr[:-1]
-    elif arr.shape[0] != expected_rows:
-        logger.warning(
-            f"routed_experts row count {arr.shape[0]} not in {{{expected_rows}, {n_tok}}}; "
-            "skipping rollout_routed_experts assign",
+        raise RuntimeError(f"vLLM routing replay: routed_experts ndim={arr.ndim}, expected 3, shape={arr.shape}")
+
+    expected_rows = max(0, len(sample.tokens) - 1)
+    if arr.shape[0] != expected_rows:
+        raise RuntimeError(
+            f"vLLM routing replay: routed_experts rows {arr.shape[0]} != expected {expected_rows} (len(tokens)-1)."
         )
-        return
+
     nl = getattr(args, "num_layers", None)
     mtk = getattr(args, "moe_router_topk", None)
     if nl is not None and mtk is not None and (arr.shape[1] != nl or arr.shape[2] != mtk):
-        logger.warning(
-            f"routed_experts shape {arr.shape} does not match args (num_layers={nl}, moe_router_topk={mtk})",
-        )
-        return
-    sample.rollout_routed_experts = arr
+        raise RuntimeError(f"vLLM routing replay: routed_experts shape {arr.shape} != (rows,{nl},{mtk}) from args.")
+    sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True))
+
+
+def _vllm_expected_routed_rows_from_tokens(token_count: int) -> int:
+    """Return expected routed-experts row count given total tokens."""
+    return max(0, token_count - 1)
+
+
+def _decode_generate_routed_experts(choice: dict[str, Any]) -> np.ndarray | None:
+    """Decode `choices[].routed_experts` (base64 npy) from vLLM generate response."""
+    routed = choice.get("routed_experts")
+    if routed is None:
+        return None
+    if not isinstance(routed, str):
+        raise RuntimeError(f"vLLM routed_experts must be base64 npy str, got {type(routed)}")
+    arr = _decode_vllm_routed_experts(routed)
+    if arr.ndim != 3:
+        raise RuntimeError(f"vLLM routed_experts ndim={arr.ndim}, expected 3, shape={arr.shape}")
+    return arr
 
 
 def _inference_generate_tokens_and_logprobs(choice: dict[str, Any]) -> tuple[list[int], list[float]]:
@@ -491,9 +512,8 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         sample.rollout_log_probs = []
     sample.rollout_log_probs += new_response_log_probs
 
-    _apply_vllm_routed_experts(args, sample, output, choice)
-
     sample.update_from_meta_info(args, meta)
+    _apply_vllm_routed_experts(args, sample, choice)
     return sample
 
 

tests/test_vllm_generate_endpoint.py

@@ -181,7 +181,13 @@ def _execute_case(case: VLLMGenerateCase):
         assert len(sample.rollout_log_probs) == sample.response_length
         assert sample.status in (Sample.Status.COMPLETED, Sample.Status.TRUNCATED)
         if case.use_rollout_routing_replay:
-            assert sample.rollout_routed_experts is not None
+            re = sample.rollout_routed_experts
+            assert re is not None
+            assert re.ndim == 3
+            expected_rows = len(sample.tokens) - 1
+            assert (
+                re.shape[0] == expected_rows
+            ), f"rollout_routed_experts rows {re.shape[0]} != len(tokens)-1 ({expected_rows})"
     finally:
         _stop_process_tree(process)
 

tests/unit/backends/vllm_utils/test_vllm_engine.py

@@ -2,9 +2,12 @@
 
 from __future__ import annotations
 
+import base64
 import dataclasses
+import io
 import json
 
+import numpy as np
 import pytest
 import requests
 import torch
@@ -178,34 +181,6 @@ def test_weight_transfer_http_timeout_fallback_to_legacy_env(vllm_engine, monkey
     assert vllm_engine._weight_transfer_http_timeout() == 42.0
 
 
-@pytest.mark.unit
-def test_response_json_or_fallback_parses_dict():
-    response = _MockResponse(json_data={"status": "ready"})
-    assert mod._response_json_or_fallback(response) == {"status": "ready"}
-
-
-@pytest.mark.unit
-def test_response_json_or_fallback_non_dict_wrapped():
-    response = _MockResponse()
-    response.json = lambda: ["a", "b"]  # type: ignore[method-assign]
-    assert mod._response_json_or_fallback(response) == {
-        "ok": False,
-        "error": "Response is not a dictionary",
-        "data": ["a", "b"],
-    }
-
-
-@pytest.mark.unit
-def test_response_json_or_fallback_invalid_json():
-    response = _MockResponse(text="not-json")
-    response.json = lambda: (_ for _ in ()).throw(ValueError("no json"))  # type: ignore[method-assign]
-    assert mod._response_json_or_fallback(response) == {
-        "ok": False,
-        "error": "Invalid JSON response",
-        "raw": "not-json",
-    }
-
-
 @pytest.mark.unit
 def test_http_base_requires_init(vllm_args):
     from slime.backends.vllm_utils.vllm_engine import VLLMEngine
@@ -329,3 +304,36 @@ def test_init_weights_update_group_raises_after_three_failures(vllm_engine, monk
             group_name="g",
             backend="nccl",
         )
+
+
+@pytest.mark.unit
+def test_verify_generate_routed_experts_accepts_single_buffer(monkeypatch):
+    prompt_toks = 5
+    gen_toks = 3
+    nrow = prompt_toks + gen_toks - 1
+    combined = np.zeros((nrow, 2, 1), dtype=np.int32)
+    buf = io.BytesIO()
+    np.save(buf, combined)
+    encoded = base64.b64encode(buf.getvalue()).decode("ascii")
+
+    def fake_post(url, json=None, timeout=None):
+        assert url.endswith("/inference/v1/generate")
+        return _MockResponse(
+            json_data={
+                "choices": [{"token_ids": list(range(gen_toks)), "routed_experts": encoded}],
+            },
+        )
+
+    monkeypatch.setattr(mod.requests, "post", fake_post)
+    mod._verify_generate_routed_experts("http://127.0.0.1:8000", "test-model")
+
+
+@pytest.mark.unit
+def test_verify_generate_routed_experts_raises_when_routing_missing(monkeypatch):
+    monkeypatch.setattr(
+        mod.requests,
+        "post",
+        lambda *a, **k: _MockResponse(json_data={"choices": [{"token_ids": [1, 2]}]}),
+    )
+    with pytest.raises(RuntimeError, match="no routed-experts"):
+        mod._verify_generate_routed_experts("http://127.0.0.1:8000", "test-model")