[Feature] Routing replay (R3) for vLLM rollout (vLLM 0.22+, /inference/v1/generate) - #49
Conversation
There was a problem hiding this comment.
Code Review
This pull request updates the vLLM backend and rollout components to support MoE routing replay with vLLM 0.22+. It introduces logic to merge prompt and generation routed experts from the /inference/v1/generate endpoint, align routing rows, and perform smoke checks during engine initialization. It also consolidates weight transfer HTTP timeout retrieval and updates unit tests. The reviewer identified several opportunities to clean up the codebase, including removing unused helper functions (_response_json_or_fallback and _routing_rows_from_http_payload) and their associated imports, as well as eliminating a duplicate helper function (_encode_routed_npy) in the test suite.
| from typing import TYPE_CHECKING, Any | ||
|
|
||
| if TYPE_CHECKING: | ||
| import numpy as np |
There was a problem hiding this comment.
make sense, could you remove this?
| 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 _encode_routed_npy(arr: np.ndarray) -> str: | ||
| buf = io.BytesIO() | ||
| np.save(buf, arr) | ||
| encoded = base64.b64encode(buf.getvalue()).decode("ascii") | ||
| return base64.b64encode(buf.getvalue()).decode("ascii") |
| "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: |
There was a problem hiding this comment.
Kinda weird, can we avoid passing rollout_routed_experts when use_rollout_routing_replay is false? I think it's inconsistent.
There was a problem hiding this comment.
Kinda weird, can we avoid passing
rollout_routed_expertswhenuse_rollout_routing_replayis false? I think it's inconsistent.
Good catch, removed
| _, action = entry | ||
| return getattr(args, dest, action.default) != action.default | ||
|
|
||
| # MoE routing replay (vLLM 0.22+, PR #39568): routed experts on ``/inference/v1/generate``. |
There was a problem hiding this comment.
I don't think we need to move the code
There was a problem hiding this comment.
I don't think we need to move the code
Agreed
| try: | ||
| data = response.json() | ||
| except (ValueError, json.JSONDecodeError): | ||
| return {"ok": False, "error": "Invalid JSON response", "raw": getattr(response, "text", "")} |
| return None | ||
|
|
||
|
|
||
| def _verify_generate_routed_experts(base_url: str, model: str, timeout_s: float = 120.0) -> None: |
There was a problem hiding this comment.
I’m not sure _verify_generate_routed_experts() belongs in VLLMEngine.init(). I think we should move it to integration test
There was a problem hiding this comment.
I’m not sure
_verify_generate_routed_experts()belongs inVLLMEngine.init(). I think we should move it to integration test
That's a fair point — the main reason I kept it in init() is to fail fast at engine startup rather than discovering a broken routing configuration only when the first rollout batch arrives.
There was a problem hiding this comment.
I think this belongs in a smoke/integration test rather than production init code. rollout already validates missing rollout_routed_experts before training, so we still fail on the first real rollout anyway.
| args, | ||
| sample, | ||
| output, | ||
| choice, |
There was a problem hiding this comment.
One issue with the current ordering: _apply_vllm_routed_experts() checks sample.status == Sample.Status.ABORTED to allow abort + 0-token samples without routed experts, but sample.update_from_meta_info(args, meta) is called only after _apply_vllm_routed_experts().
So if vLLM returns finish_reason="abort" with no generated tokens/routed_experts, the sample is still PENDING when _apply_vllm_routed_experts() runs, and the intended abort guard will not trigger. This can raise RuntimeError during partial rollout abort/drain.
Could we either call sample.update_from_meta_info(args, meta) before _apply_vllm_routed_experts(), or make _apply_vllm_routed_experts() check the parsed finish reason directly instead of relying on sample.status?
There was a problem hiding this comment.
One issue with the current ordering:
_apply_vllm_routed_experts()checkssample.status == Sample.Status.ABORTEDto allow abort + 0-token samples without routed experts, butsample.update_from_meta_info(args, meta)is called only after_apply_vllm_routed_experts().So if vLLM returns
finish_reason="abort"with no generated tokens/routed_experts, the sample is stillPENDINGwhen_apply_vllm_routed_experts()runs, and the intended abort guard will not trigger. This can raiseRuntimeErrorduring partial rollout abort/drain.Could we either call
sample.update_from_meta_info(args, meta)before_apply_vllm_routed_experts(), or make_apply_vllm_routed_experts()check the parsed finish reason directly instead of relying onsample.status?
Fixed by calling sample.update_from_meta_info() before _apply_vllm_routed_experts(), so the abort guard based on sample.status works correctly. Also added a regression test.
| from typing import TYPE_CHECKING, Any | ||
|
|
||
| if TYPE_CHECKING: | ||
| import numpy as np |
There was a problem hiding this comment.
make sense, could you remove this?
| """Decode vLLM routed-experts HTTP field (base64 npy or nested list).""" | ||
| import base64 | ||
| import io | ||
|
|
| 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``. | ||
| def _merge_generate_routed_experts( |
There was a problem hiding this comment.
Looking at vLLM’s /inference/v1/generate implementation, the response contract for routed experts seems much simpler than the compatibility logic here.
The protocol defines choices[].routed_experts as a single str | None, documented as base64-encoded .npy bytes with decoded shape (num_tokens - 1, num_layers, num_experts_per_tok):
https://github.com/vllm-project/vllm/blob/71d810bbf44b34f3a019730a6878fbcbf2480499/vllm/entrypoints/serve/disagg/protocol.py#L157-L172
The serving code also always encodes it with np.save(...) + base64, and puts it only on choice.routed_experts:
https://github.com/vllm-project/vllm/blob/71d810bbf44b34f3a019730a6878fbcbf2480499/vllm/entrypoints/serve/disagg/serving.py#L277-L294
The vLLM test decodes it the same way:
https://github.com/vllm-project/vllm/blob/71d810bbf44b34f3a019730a6878fbcbf2480499/tests/entrypoints/serve/disagg/test_return_routed_experts.py#L70-L75
Given that contract, do we need to support nested-list payloads, prompt_routed_experts, or split/trim/concat logic in VIME? I think this path can be simplified to only read choice["routed_experts"], decode the base64 .npy, validate ndim == 3, and require the row count to match len(sample.tokens) - 1.
There was a problem hiding this comment.
Looking at vLLM’s
/inference/v1/generateimplementation, the response contract for routed experts seems much simpler than the compatibility logic here.The protocol defines
choices[].routed_expertsas a singlestr | None, documented as base64-encoded.npybytes with decoded shape(num_tokens - 1, num_layers, num_experts_per_tok): https://github.com/vllm-project/vllm/blob/71d810bbf44b34f3a019730a6878fbcbf2480499/vllm/entrypoints/serve/disagg/protocol.py#L157-L172The serving code also always encodes it with
np.save(...)+ base64, and puts it only onchoice.routed_experts: https://github.com/vllm-project/vllm/blob/71d810bbf44b34f3a019730a6878fbcbf2480499/vllm/entrypoints/serve/disagg/serving.py#L277-L294The vLLM test decodes it the same way: https://github.com/vllm-project/vllm/blob/71d810bbf44b34f3a019730a6878fbcbf2480499/tests/entrypoints/serve/disagg/test_return_routed_experts.py#L70-L75
Given that contract, do we need to support nested-list payloads,
prompt_routed_experts, or split/trim/concat logic in VIME? I think this path can be simplified to only readchoice["routed_experts"], decode the base64.npy, validatendim == 3, and require the row count to matchlen(sample.tokens) - 1.
Simplified to match vLLM's documented contract
| return | ||
| if sample.status == Sample.Status.ABORTED and sample.response_length == 0: | ||
| return | ||
| pre = output.get("prompt_routed_experts") |
There was a problem hiding this comment.
vllm now doesn't have prompt_router experts so I think we can just simplify this logic.
There was a problem hiding this comment.
vllm now doesn't have prompt_router experts so I think we can just simplify this logic.
Simplified to match vLLM's documented contract
| gen_url = f"{base}/inference/v1/generate" | ||
| with trace_span(sample, "vllm_mm_generate", attrs={"max_tokens": params["max_new_tokens"]}): | ||
| output = await post(gen_url, generate_body, headers=headers) | ||
| request_prompt_len = len(generate_body.get("token_ids") or []) |
There was a problem hiding this comment.
vllm now doesn't have prompt_router experts so I think we can just simplify this logic.
There was a problem hiding this comment.
vllm now doesn't have prompt_router experts so I think we can just simplify this logic.
Simplified to match vLLM's documented contract
ff5b248 to
af58b9d
Compare
Address non-controversial PR feedback by deleting unused vLLM engine helper code and reusing the existing routed-experts encoder helper in rollout unit tests.
af58b9d to
ec0814c
Compare
| sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)) | ||
|
|
||
|
|
||
| def _vllm_expected_routed_rows_from_tokens(token_count: int) -> int: |
There was a problem hiding this comment.
Done. Also cleaned up some dead code left over from debugging.
| @pytest.mark.unit | ||
| def test_verify_generate_routed_experts_accepts_single_buffer(monkeypatch): | ||
| prompt_toks = 5 | ||
| gen_toks = 3 |
There was a problem hiding this comment.
We can remove the unit test and have an integration test about it.
There was a problem hiding this comment.
| ) | ||
| _wait_server_healthy(self._http_base(), process=self.process) | ||
| base = self._http_base() | ||
| _wait_server_healthy(base, process=self.process) |
There was a problem hiding this comment.
I think we can revert here.
…e/v1/generate) (#49) * vllm engine supports router replay * feat(r3): vLLM 0.22+ generate API for MoE routing replay * chore(review): remove dead helpers and duplicate test encoder Address non-controversial PR feedback by deleting unused vLLM engine helper code and reusing the existing routed-experts encoder helper in rollout unit tests. * remove dead code * refactor(r3): trim redundant routed replay code --------- Co-authored-by: aoshen02 <aoshen@inferact.ai>
…ounterpart) This test was vime-specific (created in #49, not from slime) and was never registered in CI. Remove it to keep the test tree aligned with slime@44d29ee. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…ms (#218) * fix: complete slime #1985 sync — remove TIGHT_HOST_MEMORY + align params PR #214 synced #1985's TIGHT_DEVICE_MEMORY removal and batch-size shrinks but missed 10 files that used TIGHT_HOST_MEMORY (a different env-var guard with the same pre-#1985 pattern). Also aligns n-samples-per-prompt (8→4) and num-critic-only-steps (3→2) that were left at pre-#1985 values. Files: test_moonlight_16B_A3B{,_r3}, test_qwen3_{0.6B_parallel_check, 30B_A3B{,_r3}, 4B_ckpt, 4B_ppo{,_disaggregate,_train_critic_only}, 4B_streaming_partial_rollout}, test_qwen2.5_0.5B_ppo_critic_only_short Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: remove leftover blank lines from TIGHT_ variable removal #214 and the preceding commit removed TIGHT_DEVICE_MEMORY and TIGHT_HOST_MEMORY definitions but left behind an extra blank line in the header area of 11 test files. Collapse double-blank to single-blank to match slime's spacing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * revert: remove pre-ported slime #2016 content from http_utils http_utils.py had get_rollout_num_engines() and run_router improvements that were pre-ported from slime #2016 (post-cutoff 44d29ee). Revert to match slime@44d29ee baseline so the diff stays clean; #2016 will be synced as a whole when its turn comes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: remove test_vllm_generate_endpoint.py (vime-native, no slime counterpart) This test was vime-specific (created in #49, not from slime) and was never registered in CI. Remove it to keep the test tree aligned with slime@44d29ee. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * sync: port gpu_lock_exec subprocess signal forwarding (slime #1945) PR #155 synced #1945's DistOptim checkpoint rider but missed the gpu_lock_exec.py change: replaces os.execvp with subprocess.Popen + proper signal forwarding (SIGINT/SIGTERM/SIGHUP) and fd_lock cleanup. This prevents orphaned GPU-holding processes when CI runners are cancelled. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: complete #1985 sync + fix NUM_GPUS placement + pre-commit - Remove TIGHT_HOST_MEMORY from 10 test files (missed in #214) - Align n-samples-per-prompt (8→4) and num-critic-only-steps (3→2) - Remove leftover blank lines from TIGHT_ variable removal (11 files) - Revert http_utils.py pre-ported #2016 content to slime@44d29ee baseline - Remove test_vllm_generate_endpoint.py (vime-native, no slime counterpart) - Port gpu_lock_exec subprocess signal forwarding (slime #1945) - Align remaining test params (over-sampling-batch-size, max-tokens-per-gpu) - Full test_qwen3_4B_ckpt.py #1945 sync (optimizer placement CLI) - Fix NUM_GPUS=0 placement in 7 CPU test files (was inside decorators) - NamedTemporaryFile multi-line formatting All pre-commit checks pass. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>



Purpose
Implements RFC #32 Phase 2 / follow-up to #34: MoE routing replay (
--use-rollout-routing-replay) on the vLLM rollout path using vLLM ≥ 0.22 native support (vllm#39568).SGLang already fills
sample.rollout_routed_expertswith shape(len(tokens) - 1, num_layers, moe_router_topk). On vLLM 0.21.x, routing required a site-packages patch and/or/v1/completionsworkarounds. From vLLM 0.22+,/inference/v1/generatereturns routed experts (often as a single base64 npy buffer onchoices[].routed_experts).This PR drops the disagg patch path and standardizes rollout + engine smoke checks on
/inference/v1/generate.Supersedes the approach in #34 (v0.21.0 +
docker/patch/latest/vllm.patch+ completions fallback).What's included
Rollout (
slime/rollout/vllm_rollout.py)POST /inference/v1/generateonly (removed/v1/completionsbranch)._merge_generate_routed_experts: merge split or single-buffer routing; align tolen(tokens) - 1.Engine (
slime/backends/vllm_utils/vllm_engine.py)use_rollout_routing_replay: inject--enable-return-routed-expertsonly (no 0.21.x--no-async-scheduling/--no-enable-prefix-cachingguards)._verify_generate_routed_experts: startup smoke on/inference/v1/generate.Tests
tests/unit/rollout/test_vllm_rollout.py— merge/apply routing, single-buffer case.tests/unit/backends/vllm_utils/test_vllm_engine.py—_verify_generate_routed_experts.tests/test_vllm_generate_endpoint.py—qwen3-30b-a3b-r3integration.Requirements
--vllm-enable-expert-parallel--use-rollout-routing-replayVerification logs
Environment: 8×A100, container
vime_v22, Qwen3-30B-A3B, train/rollout 4+4, vLLM0.21.1rc1.dev38(includes #39568 generate routing). NoSplit sizes/ missing-routing errors observed.1) Local unit tests
2) Full verify — HTTP + rollout E2E
Phase A — standalone vLLM HTTP smoke (
:8000)Phase B —
debug-rollout-only+.ptshape check3) Production-shaped training — 4+4 R3 smoke
Training completed multiple rollout→train steps without routing-shape failures.
Test plan
Unit
pytest tests/unit/rollout/test_vllm_rollout.py -k "routed or merge_generate" pytest tests/unit/backends/vllm_utils/test_vllm_engine.py -k verify_generateIntegration (GPU)
Manual
--enable-return-routed-experts(no patch).--use-rollout-routing-replay --vllm-enable-expert-parallel, 4+4 layout.rollout_routed_expertsshape(len(tokens)-1, num_layers, moe_router_topk).Out of scope
v0.22.0image tag.tests/test_qwen3_30B_A3B_r3.py.docker/patchtree.Related
Related to #32 (Phase 2).