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

slime/backends/vllm_utils/vllm_engine.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
 import ipaddress
+import json
 import logging
 import multiprocessing
 import os
 import time
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
 from urllib.parse import quote
 
 import requests
@@ -308,9 +313,6 @@ def launch_server_process(
     # unless the user already passed --vllm-max-model-len explicitly.
     if args.rollout_max_context_len is not None and getattr(args, "vllm_max_model_len", None) is None:
         cmd += ["--max-model-len", str(args.rollout_max_context_len)]
-    if getattr(args, "use_rollout_routing_replay", False):
-        cmd += ["--enable-return-routed-experts"]
-
     # vime-preferred defaults — must be explicitly forwarded because the vllm-side
     # default would otherwise apply (the generic forwarder skips values that equal
     # action.default).
@@ -336,6 +338,10 @@ def _user_overrode(dest: str) -> bool:
         _, action = entry
         return getattr(args, dest, action.default) != action.default
 
+    # MoE routing replay (vLLM 0.22+, PR #39568): routed experts on ``/inference/v1/generate``.
+    if getattr(args, "use_rollout_routing_replay", False):
+        cmd += ["--enable-return-routed-experts"]
+
     # 1) gpu_memory_utilization: vllm default 0.92 OOMs in colocate training; vime ships 0.55.
     if _user_overrode("vllm_gpu_memory_utilization"):
         gpu_mem = args.vllm_gpu_memory_utilization
@@ -405,6 +411,83 @@ def _redact_cmd_for_log(cmd: list[str]) -> str:
     return " ".join(parts)
 
 
+def _response_json_or_fallback(response) -> dict:
+    """Parse JSON from an HTTP response, returning a structured error dict on failure."""
+    try:
+        data = response.json()
+    except (ValueError, json.JSONDecodeError):
+        return {"ok": False, "error": "Invalid JSON response", "raw": getattr(response, "text", "")}
+    if not isinstance(data, dict):
+        return {"ok": False, "error": "Response is not a dictionary", "data": data}
+    return data
+
+
+def _routing_rows_from_http_payload(value: Any) -> np.ndarray | None:
+    """Decode vLLM routed-experts HTTP field (base64 npy or nested list)."""
+    import base64
+    import io
+
+    import numpy as np
+
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return np.load(io.BytesIO(base64.b64decode(value)), allow_pickle=False)
+    if isinstance(value, list):
+        return np.asarray(value, dtype=np.int32)
+    return None
+
+
+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 _merge_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)
+
+    merged = _merge_generate_routed_experts(
+        body,
+        choice,
+        gen_token_count=num_gen,
+        prompt_token_count=num_prompt,
+    )
+    if merged 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(merged.shape[0])
+    if expected_rows > 0 and n_rows not in (expected_rows, expected_rows + 1):
+        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 +534,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 +652,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 +667,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 +890,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 +915,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,7 +774,17 @@ 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:
+        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
+        elif samples[0].rollout_routed_experts is not None:
             train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples]
 
         if samples[0].train_metadata is not None: