feat: unify renderer + teacher endpoints onto vLLM 0.20 /inference/v1/generate

[hallerite]

Summary

vLLM 0.20 ships a generic tokens-in / tokens-out endpoint at /inference/v1/generate (vllm.entrypoints.serve.disagg.serving.ServingTokens) that supersedes the bespoke /v1/generate handler prime-rl maintained on top of vllm 0.19. Replace it.

Net effect: −1 endpoint, ~−275 LoC, no functional change for callers.

What's in the PR

Server side

• Drop src/prime_rl/inference/vllm/serving_generate.py and the /v1/generate route in server.py — vLLM 0.20's build_app already attaches /inference/v1/generate via attach_disagg_router.
• Subclass upstream's ServingTokens with PrimeRlServingTokens to preserve two prime-rl features the upstream protocol doesn't natively cover:
  1. data_parallel_rank routing — read from the X-data-parallel-rank header and forwarded to engine_client.generate. The DP-replicated inference servers prime-rl runs need this to target a specific replica.
  2. routed_experts per-token export — surfaced on each choice when the engine is launched with enable_return_routed_experts=True. This is what the trainer's router-replay path consumes.
• custom_init_app_state swaps the upstream serving_tokens instance for our subclass after init_app_state.

Orchestrator side

• compute_teacher_logprobs in orchestrator/utils.py points at /inference/v1/generate, builds the upstream payload (token_ids + nested sampling_params), and re-flattens prompt_logprobs from the upstream list[dict[token_id, Logprob]] shape back to the list[float] callers expect.

Tests

• Replace test_serving_generate.py (class deleted) with test_serving_tokens.py — exercises the prime-rl deltas (routed_experts encoding, response shape stability).
• Update test_teacher_logprobs.py to expect the new endpoint URL, payload shape, and response unwrap.

Renderers / verifiers pins

• Pin renderers to PyPI ==0.1.6 (the first release after the renderers/verifiers monorepo split — same code as 9acdc60, just shipped as a published wheel) and declare it as a direct prime-rl dependency. Previously transitively pulled via verifiers' in-tree workspace package.
• Bump verifiers to 7bdc769 to pick up the post-split main. The renderers release carries the matching client-side switch to /inference/v1/generate.

Test plan

• tests/unit/inference/test_serving_tokens.py (4 tests) and tests/unit/orchestrator/test_teacher_logprobs.py (1 test) pass locally against vllm 0.20.
• Server module imports cleanly (prime_rl.inference.vllm.server, prime_rl.inference.vllm.serving_tokens).
• Companion verifiers tests (packages/renderers/tests/test_client.py) green.
• uv.lock regenerated against renderers==0.1.6 (PyPI) and verifiers 7bdc769.
• E2E text-only renderer + TITO rollouts against a live vllm 0.20 server. configs/multi_reverse_text/rl.toml × 20 steps on 2× RTX PRO 6000. Renderer run: 2688 calls to /inference/v1/generate, all steps green, eval Avg@4=0.85. TITO run: 20/20 steps green, eval Avg@4=0.85 (reverse-text is single-turn so the TITO client falls back to MITO /v1/chat/completions per openai_chat_completions_token_client.py:114 — exercises the routing flag but not the /tokens wire path).

Out of scope (follow-ups)

• VLM rollouts via the renderer. The new endpoint handles MM features end-to-end (it accepts pre-built MultiModalFeatures), but validate_renderer_vs_vlm in configs/orchestrator.py still blocks the combination. Lifting the ban needs the renderer client to build features client-side (HF processor → MultiModalKwargsItem → base64 msgpack). Cleaner as a separate PR after this lands.
• mismatch_kl is unrelated. The drift identified in MM_KL_INVESTIGATION_SUMMARY.md lives in vLLM's bf16 forward kernel, not the wire protocol — this migration neither helps nor hurts it.

🤖 Generated with Claude Code

---

Note

Medium Risk
Swaps the token-level generation API from a custom /v1/generate implementation to vLLM 0.20’s /inference/v1/generate, which can break callers or response expectations if any edge cases differ. Adds a custom ServingTokens subclass that intercepts headers and post-processes outputs, so integration correctness depends on matching vLLM’s evolving protocol.

Overview
Moves token-level generation off the legacy /v1/generate endpoint onto vLLM 0.20’s /inference/v1/generate. The bespoke serving_generate.py handler and its route are removed, and server.py now swaps vLLM’s serving_tokens instance to a new PrimeRlServingTokens wrapper during app init.

PrimeRlServingTokens preserves prime-RL-specific behavior on the new endpoint: forwards X-data-parallel-rank into engine_client.generate, re-exports per-token routed_experts in responses, and applies server-side defaulting for sampling_params.max_tokens when omitted (avoiding vLLM’s 16-token default).

The orchestrator’s compute_teacher_logprobs is updated to call /inference/v1/generate with the new request schema (token_ids + nested sampling_params) and to re-flatten upstream prompt_logprobs back into the legacy list-of-floats shape expected by callers. Tests are updated accordingly (replace test_serving_generate.py with test_serving_tokens.py, adjust teacher logprobs expectations), and dependencies are updated by pinning renderers==0.1.6 (PyPI) and bumping verifiers to a new git rev with lockfile refresh.

^{Reviewed by Cursor Bugbot for commit 913cc4c. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Duplicated _encode_routed_experts logic across two files

Low Severity

The standalone _encode_routed_experts function in serving_tokens.py has identical logic to the instance method _RoutedExpertsCapture._encode_routed_experts in serving_chat_with_tokens.py. Both base85-encode a numpy array and return a {"data": ..., "shape": ...} dict. One shared utility function could serve both callers, reducing the risk of inconsistent fixes if the encoding format ever changes.

 

^{Reviewed by Cursor Bugbot for commit e16f639. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 913cc4c. Configure here.}

[cursor]

Missing engine abort on client disconnect

Medium Severity

When a client disconnects mid-generation, serve_tokens_full_generator catches asyncio.CancelledError and returns an error response, but never calls engine_client.abort(request_id). The deleted serving_generate.py explicitly called await self.engine_client.abort(request_id) before re-raising, ensuring the engine stopped processing the request. Without the abort, the inference engine may continue consuming GPU compute for requests whose clients are long gone, which compounds under high concurrency (2k+ simultaneous rollouts).

 

^{Reviewed by Cursor Bugbot for commit 913cc4c. Configure here.}