[Feature]Routing replay (R3) for vLLM rollout (vLLM v0.21.0)

[CalvinXKY]

Purpose

Implements Phase 1 of RFC #32: MoE routing replay (--use-rollout-routing-replay) on the vLLM rollout path, pinned to vLLM 0.21.0.

SGLang already fills sample.rollout_routed_experts with shape (len(tokens) - 1, num_layers, moe_router_topk). vLLM 0.21.0 records routing internally but /inference/v1/generate did not expose it, and vllm_rollout only handled base64 npy on choices[].routed_experts.

Follow-up (not this PR): colocate E2E aligned with tests/test_qwen3_30B_A3B_r3.py; upgrade to latest vLLM (RFC Phase 2).

What's included

• docker/patch/latest/vllm.patch — add prompt_routed_experts + choices[].routed_experts on /inference/v1/generate; bake in docker/Dockerfile; docs in docker/README.md.
• slime/rollout/vllm_rollout.py — merge prompt + gen routing (nested list or base64 npy); assign Megatron-shaped rollout_routed_experts.
• slime/backends/vllm_utils/vllm_engine.py — when R3 enabled: --enable-return-routed-experts, --no-async-scheduling (unless overridden); --enable-expert-parallel + linear placement only if --vllm-enable-expert-parallel.
• Tests — tests/unit/rollout/test_vllm_rollout.py; shape checks in tests/test_vllm_generate_endpoint.py (R3 case).

Test plan

Model: Qwen3-30B-A3B (num_layers=48, moe_router_topk=8)

Unit

pytest tests/unit/rollout/test_vllm_rollout.py -k routed_experts

Integration (GPU + model)

pytest tests/test_vllm_generate_endpoint.py -k r3

Manual — decoupled (Gate B, per RFC #32)

1. Image with patch applied (docker build -f docker/Dockerfile ... or runtime patch per docker/README.md).
2. Rollout flags: --use-rollout-routing-replay --vllm-enable-expert-parallel (MoE multi-GPU), no --colocate.
3. Confirm /inference/v1/generate returns prompt_routed_experts and choices[].routed_experts.
4. Short train/rollout run: Megatron accepts rollout_routed_experts, shape (len(tokens)-1, 48, 8).

Manual — colocate (Gate C, follow-up)

Same layout as tests/test_qwen3_30B_A3B_r3.py with vLLM rollout args; not required to merge this PR.

Closes #32 (Phase 1)

[gemini-code-assist]

Code Review

This pull request adds support for exposing and processing MoE routing information (both prompt and generation routed experts) from vLLM 0.21.0. It includes a Dockerfile patch step, updates to the vLLM server launch arguments, and logic to merge prompt and generation routing arrays. The review feedback highlights two key issues: first, the patch utility might not be pre-installed in the Docker base image, which could cause build failures; second, empty routing lists could lead to dimension mismatches or index errors during array concatenation, which can be resolved with safer list-length checks.

[gemini-code-assist]

If prompt_routed_experts (pre) or routed_experts (gen) is returned as an empty list [] (which can happen if routing is empty or disabled for a segment), _vllm_routed_experts_payload_to_array will return a 1D empty array of shape (0,). Concatenating a 1D array with a 3D array of shape (gen_len, num_layers, top_k) will raise a ValueError due to dimension mismatch. Additionally, if both are empty, parts will be empty, causing an IndexError when accessing parts[0]. Checking len(...) > 0 before appending ensures we only concatenate non-empty 3D arrays.

Suggested change

	if pre is not None:
	parts.append(_vllm_routed_experts_payload_to_array(pre))
	if gen is not None:
	parts.append(_vllm_routed_experts_payload_to_array(gen))
	arr = np.concatenate(parts, axis=0) if len(parts) > 1 else parts[0]
	if pre is not None and len(pre) > 0:
	parts.append(_vllm_routed_experts_payload_to_array(pre))
	if gen is not None and len(gen) > 0:
	parts.append(_vllm_routed_experts_payload_to_array(gen))
	if not parts:
	return
	arr = np.concatenate(parts, axis=0) if len(parts) > 1 else parts[0]

[gemini-code-assist]

The patch utility is invoked here to apply vllm.patch. However, patch is not explicitly installed in the apt-get install command on line 20, and minimal base images (such as vllm/vllm-openai) typically do not have it pre-installed. This could cause the Docker build to fail with a FileNotFoundError when trying to run the patch command. Please ensure patch is added to the list of packages installed via apt-get on line 20.

[CalvinXKY]

Comparison method of return values for vLLM vs SGLang requests:

Start two services on a single machine and send corresponding requests to both services.

export CUDA_VISIBLE_DEVICES=0,1,2,3
VLLM_SERVER_DEV_MODE=1 python -m vllm.entrypoints.openai.api_server \
  --model /data/nfs_87/model/Qwen3-30B-A3B \
  --served-model-name Qwen3-30B-A3B-vllm \
  --host 0.0.0.0 \
  --port 8000 \
  --tensor-parallel-size 4 \
  --data-parallel-size 1 \
  --enable-expert-parallel \
  --trust-remote-code \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.75 \
  --enable-return-routed-experts \
  --no-async-scheduling

curl -s http://127.0.0.1:8000/v1/completions   -H "Content-Type: application/json"   -d '{
    "model": "Qwen3-30B-A3B-vllm",
    "prompt": "What is 12 + 15?",
    "max_tokens": 32,
    "temperature": 1.0,
    "logprobs": 1,
    "return_token_ids": true
  }' | python3 -m json.tool > /data/nfs_87/xky/logs/vllm_curl.log

curl -s http://127.0.0.1:8000/inference/v1/generate \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen3-30B-A3B-vllm",
    "token_ids": [151644, 872, 374, 220, 717, 489, 717, 30],
    "sampling_params": {
      "max_tokens": 32,
      "temperature": 1.0,
      "top_p": 1.0,
      "logprobs": 1
    },
    "return_routed_experts": true
  }' | python3 -m json.tool > /data/nfs_87/xky/logs/vllm_curl.log
  
export CUDA_VISIBLE_DEVICES=4,5,6,7
python -m sglang.launch_server \
  --model-path /data/nfs_87/model/Qwen3-30B-A3B \
  --served-model-name Qwen3-30B-A3B-sglang \
  --host 0.0.0.0 \
  --port 8001 \
  --tp 4 \
  --ep 1 \
  --context-length 8192 \
  --mem-fraction-static 0.75 \
  --trust-remote-code \
  --enable-return-routed-experts 
  
curl -s http://127.0.0.1:8001/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen3-30B-A3B",
    "prompt": "What is 12 + 15?",
    "max_tokens": 32,
    "temperature": 1.0,
    "logprobs": 1,
    "return_token_ids": true,
    "return_routed_experts": true
  }' | python3 -m json.tool  > /data/nfs_87/xky/logs/sglang_curl.log
 
 
curl -s http://127.0.0.1:8001/generate \
  -H "Content-Type: application/json" \
  -d '{
    "text": "What is 12 + 15?",
    "sampling_params": { "max_new_tokens": 32, "temperature": 1.0 },
    "return_routed_experts": true
  }' | python3 -m json.tool > /data/nfs_87/xky/logs/sglang_curl.log
  
  
curl -s http://127.0.0.1:8000/inference/v1/generate \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen3-30B-A3B-vllm",
    "token_ids": [151644, 872, 374, 220, 717, 489, 717, 30],
    "sampling_params": {
      "max_tokens": 32,
      "temperature": 1.0,
      "top_p": 1.0,
      "logprobs": 1
    },
    "return_routed_experts": true
  }' | python3 -m json.tool > /data/nfs_87/xky/logs/vllm_curl2.log
  

Aligning return values after modifying vllm rollout:

[CalvinXKY]

Smoke test:

Result:

---

verify_rollout_routed_experts_pt.py

#!/usr/bin/env python3
"""Check rollout_*.pt from ``--save-debug-rollout-data`` for routing replay contract."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import torch


def _check_sample(sample: dict, *, num_layers: int, moe_router_topk: int, index: int) -> None:
    tokens = sample.get("tokens") or []
    re = sample.get("rollout_routed_experts")
    if re is None:
        raise SystemExit(f"sample[{index}]: rollout_routed_experts is None (routing replay not applied)")
    arr = np.asarray(re)
    if arr.ndim != 3:
        raise SystemExit(f"sample[{index}]: expected ndim=3, got {arr.ndim} shape={arr.shape}")
    expected_rows = max(0, len(tokens) - 1)
    if arr.shape[0] != expected_rows:
        raise SystemExit(
            f"sample[{index}]: shape[0]={arr.shape[0]} != len(tokens)-1={expected_rows} "
            f"(tokens={len(tokens)})"
        )
    if arr.shape[1] != num_layers or arr.shape[2] != moe_router_topk:
        raise SystemExit(
            f"sample[{index}]: shape={arr.shape} != ({expected_rows}, {num_layers}, {moe_router_topk})"
        )
    print(
        f"  sample[{index}]: OK shape={tuple(arr.shape)} "
        f"tokens={len(tokens)} response_length={sample.get('response_length', '?')}"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("pt_path", type=Path, help="Path to rollout_0.pt (or similar)")
    parser.add_argument("--num-layers", type=int, default=48)
    parser.add_argument("--moe-router-topk", type=int, default=8)
    parser.add_argument("--min-samples-with-routing", type=int, default=1)
    args = parser.parse_args()

    if not args.pt_path.is_file():
        print(f"ERROR: file not found: {args.pt_path}", file=sys.stderr)
        return 1

    data = torch.load(args.pt_path, map_location="cpu", weights_only=False)
    samples = data.get("samples") or []
    if not samples:
        print(f"ERROR: no samples in {args.pt_path}", file=sys.stderr)
        return 1

    print(f"Checking {args.pt_path} ({len(samples)} samples)")
    ok = 0
    for i, sample in enumerate(samples):
        if sample.get("rollout_routed_experts") is None:
            continue
        _check_sample(sample, num_layers=args.num_layers, moe_router_topk=args.moe_router_topk, index=i)
        ok += 1

    if ok < args.min_samples_with_routing:
        print(
            f"ERROR: only {ok} sample(s) have rollout_routed_experts "
            f"(need >= {args.min_samples_with_routing})",
            file=sys.stderr,
        )
        return 1

    print(f"PASS: {ok} sample(s) satisfy routing replay shape contract.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

---

verify_routing_replay_qwen3_30b.sh

set -euo pipefail

if grep -q $'\r' "$0" 2>/dev/null; then
  exec bash <(sed 's/\r$//' "$0") "$@"
fi

REPO_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." &>/dev/null && pwd)"
VIME_ROOT="${VIME_ROOT:-/data/nfs_87/xky/new_rl/vime}"
SCRIPT_DIR="${REPO_ROOT}/run_script"
VERIFY_MODE="${VERIFY_MODE:-all}"
SKIP_PATCH="${SKIP_PATCH:-0}"
SKIP_HTTP="${SKIP_HTTP:-0}"
VLLM_BASE_URL="${VLLM_BASE_URL:-http://127.0.0.1:8000}"
MODEL="${MODEL:-Qwen3-30B-A3B-vllm}"

ROLLOUT_NUM_GPUS="${ROLLOUT_NUM_GPUS:-4}"
ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-4}"
DEBUG_DATA_DIR="${DEBUG_DATA_DIR:-/data/nfs_87/xky/logs/routing_replay_verify}"
LOG_FILE="${LOG_FILE:-/data/nfs_87/xky/logs/verify_routing_replay_$(date +%Y%m%d_%H%M%S).log}"

HF_CKPT="${HF_CKPT:-/data/nfs_87/model/Qwen3-30B-A3B}"
REF_LOAD="${REF_LOAD:-/data/nfs_87/model/Qwen3-30B-A3B_torch_dist}"
PROMPT_DATA="${PROMPT_DATA:-/data/nfs_87/xky/datasets/dapo-math-17k/dapo-math-17k.jsonl}"

NUM_LAYERS=48
MOE_ROUTER_TOPK=8

log() { echo "[routing-replay-verify] $*"; }
die() { echo "[routing-replay-verify] ERROR: $*" >&2; exit 1; }

run_http_checks() {
  if [[ "${SKIP_HTTP}" == "1" ]]; then
    log "SKIP_HTTP=1: skip HTTP checks."
    return 0
  fi
  if ! curl -sf "${VLLM_BASE_URL}/v1/models" >/dev/null 2>&1; then
    log "vLLM not reachable at ${VLLM_BASE_URL}; skip HTTP (set VLLM_BASE_URL or start server)."
    return 0
  fi
  log "HTTP checks against ${VLLM_BASE_URL} ..."
  VLLM_BASE_URL="${VLLM_BASE_URL}" MODEL="${MODEL}" \
    bash "${SCRIPT_DIR}/apply_vllm_generate_routed_experts_patch.sh" --curl
}

run_rollout_smoke() {
  log "debug-rollout-only smoke (rollout_num_gpus=${ROLLOUT_NUM_GPUS}) ..."
  mkdir -p "${DEBUG_DATA_DIR}" "$(dirname "${LOG_FILE}")"

  export PYTHONPATH="${VIME_ROOT}:/root/Megatron-LM"
  export CUDA_DEVICE_MAX_CONNECTIONS=1
  export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
  unset PYTORCH_ALLOC_CONF
  export SLIME_NCCL_BRIDGE_CPU_FALLBACK="${SLIME_NCCL_BRIDGE_CPU_FALLBACK:-1}"
  export PYTHONUNBUFFERED=1

  # shellcheck source=/dev/null
  source "${VIME_ROOT}/scripts/models/qwen3-30B-A3B.sh"
  cd "${VIME_ROOT}"

  python train.py \
    --debug-rollout-only \
    --train-backend megatron \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 4 \
    --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" \
    --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE}" \
    "${MODEL_ARGS[@]}" \
    \
    --hf-checkpoint "${HF_CKPT}" \
    --ref-load "${REF_LOAD}" \
    --model-name qwen3moe \
    \
    --prompt-data "${PROMPT_DATA}" \
    --input-key prompt \
    --label-key label \
    --apply-chat-template \
    --rm-type deepscaler \
    \
    --num-rollout 1 \
    --rollout-batch-size 2 \
    --n-samples-per-prompt 1 \
    --rollout-max-response-len 128 \
    --rollout-temperature 0.0 \
    --global-batch-size 2 \
    \
    --use-rollout-routing-replay \
    --save-debug-rollout-data "${DEBUG_DATA_DIR}/rollout_{rollout_id}.pt" \
    \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 1 \
    --context-parallel-size 1 \
    --expert-model-parallel-size 4 \
    --expert-tensor-parallel-size 1 \
    --recompute-granularity full \
    --recompute-method uniform \
    --recompute-num-layers 1 \
    --use-dynamic-batch-size \
    --max-tokens-per-gpu 2048 \
    \
    --vllm-gpu-memory-utilization 0.70 \
    --vllm-max-model-len 8192 \
    --vllm-server-concurrency 32 \
    --vllm-enable-expert-parallel \
    --vllm-disable-custom-all-reduce \
    \
    --attention-dropout 0.0 \
    --hidden-dropout 0.0 \
    --attention-softmax-in-fp32 \
    --attention-backend flash \
    2>&1 | tee -a "${LOG_FILE}"

  PT_FILE="${DEBUG_DATA_DIR}/rollout_0.pt"
  [[ -f "${PT_FILE}" ]] || die "missing ${PT_FILE} after rollout"

  log "Validating saved rollout tensors in ${PT_FILE} ..."
  python3 "${SCRIPT_DIR}/verify_rollout_routed_experts_pt.py" \
    "${PT_FILE}" \
    --num-layers "${NUM_LAYERS}" \
    --moe-router-topk "${MOE_ROUTER_TOPK}" \
    --min-samples-with-routing 1

  log "Rollout path PASS: ${PT_FILE}"
}

# --- main ---
log "VIME_ROOT=${VIME_ROOT} VERIFY_MODE=${VERIFY_MODE}"
[[ -d "${VIME_ROOT}" ]] || die "VIME_ROOT not found: ${VIME_ROOT}"
[[ -f "${VIME_ROOT}/train.py" ]] || die "train.py missing under ${VIME_ROOT}"

if [[ "${SKIP_PATCH}" != "1" ]]; then
  log "Applying vLLM generate routed_experts patch ..."
  WORKSPACE_ROOT="${REPO_ROOT}" VIME_REPO_ROOT="${VIME_ROOT}" \
    bash "${SCRIPT_DIR}/apply_vllm_generate_routed_experts_patch.sh"
else
  log "SKIP_PATCH=1: skip vLLM patch."
fi

case "${VERIFY_MODE}" in
  http)
    run_http_checks
    ;;
  rollout)
    run_rollout_smoke
    ;;
  all)
    run_http_checks
    run_rollout_smoke
    ;;
  *)
    die "Unknown VERIFY_MODE=${VERIFY_MODE} (use all|http|rollout)"
    ;;
esac

log "All requested checks finished successfully."

[andakai]

Hi, overall it looks good to me. I left a few comments for your reference.

[andakai]

This part is a bit redundant with the generic vLLM arg init. _forward_vllm_cli_args(args, cmd) will deal with vllm_enable_expert_parallel. expert-placement-strategy is linear by default.

[CalvinXKY]

Agreed.

[andakai]

Mentioning “SGLang's contract” here is a bit confusing.

Maybe "The concatenated tensor must match the rollout_routed_experts contract consumed by Megatron routing replay:"?

[CalvinXKY]

Agreed.

[aoshen02]

Sounds it's very struggle so that we need to add these guard? Actually the pr vllm-project/vllm#39568 that I mentioned earlier has already enabled async scheduling and prefix caching, we can just remove these guard as 0.22.0 is expected to release today or tomorrow.

[CalvinXKY]

Agreed for vLLM ≥ 0.22 (#39568). vime Docker still pins 0.21.0; we keep --no-async-scheduling and --no-enable-prefix-caching as defaults under use_rollout_routing_replay until the base image bumps.