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

[hallerite]

What changes and why

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. This PR moves prime-rl onto it.

Server

• Drop serving_generate.py + the /v1/generate route. vLLM 0.20's build_app already attaches /inference/v1/generate via attach_disagg_router — the bespoke handler is now redundant.
• Subclass ServingTokens as PrimeRlServingTokens. Two prime-rl features aren't covered by the upstream protocol:
  1. data_parallel_rank routing — read from the X-data-parallel-rank header and forwarded to engine_client.generate. The DP-replicated inference servers we run need this to target a specific replica.
  2. routed_experts per-token export — surfaced on each choice when the engine runs with enable_return_routed_experts=True. The trainer's router-replay path consumes this.
• custom_init_app_state swaps the upstream serving_tokens instance for our subclass after init_app_state.

Orchestrator

• compute_teacher_logprobs 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]] back to the list[float] callers expect.
• Bypasses the SDK's parse layer without losing auth. The OpenAI SDK's cast_to enforces openai.BaseModel; vLLM's GenerateResponse is plain pydantic.BaseModel and the SDK rejects it with TypeError. Calling client.post(..., cast_to=httpx.Response, body=...) runs the full request pipeline (auth_headers, retries, timeouts, idempotency keys) and just hands back the raw response for us to validate ourselves — going through client._client.post directly would skip auth.

Tests

• Replace test_serving_generate.py with test_serving_tokens.py covering the prime-rl deltas (routed_experts encoding, response shape stability).
• Update test_teacher_logprobs.py for the new endpoint URL, payload shape, and response unwrap, and assert the AsyncOpenAI SDK doesn't double-prefix the absolute URL.

Renderers / verifiers pins

• Pin renderers==0.1.6 (PyPI). First release after the renderers/verifiers monorepo split — same code as 9acdc60, now a published wheel. Declared as a direct prime-rl dependency since it was previously pulled transitively through verifiers' in-tree workspace package.
• Bump verifiers to 7bdc769 for the post-split main. The matching renderers release carries the client-side switch to /inference/v1/generate.

Test plan

Unit

• 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 smoke matrix (5 steps each, all green, no NaN, no error)

Test	Model	Path under test	Transport	Result
A	Qwen3-4B-Instruct-2507	TITO /v1/chat/completions/tokens (agentic, multi-turn)	NCCL	✅ 5/5 steps, 9 960 TITO calls
B	Qwen3-30B-A3B-Thinking-2507 (MoE, EP=8)	Renderer → /inference/v1/generate	NCCL	✅ 5/5 steps, 1 481 generate calls, 3 weight broadcasts
C	Qwen3-30B-A3B-Thinking-2507 (MoE, EP=8)	TITO config (single-turn → MITO /v1/chat/completions fallback)¹	NCCL	✅ 5/5 steps, 1 156 chat calls
D	Qwen3-0.6B (live inference server)	compute_teacher_logprobs → /inference/v1/generate (prompt_logprobs=1)	n/a	✅ Flat list[float], first slot 0.0, all values ≤ 0

¹ reverse-text is single-turn, so the TITO client takes the MITO fallback at openai_chat_completions_token_client.py:114. A and B together cover both the TITO /tokens wire path and the new disagg path; C confirms 30B MoE weight reload doesn't regress on the unchanged TITO config.

WandB:

• A: 4b-pr-smoke-agentic-5step
• B: pr-smoke-30b-rendered-rt-1010
• C: pr-smoke-30b-tito-rt-1010

🤖 Generated with Claude Code

---

Note

Medium Risk
Moderate risk because it removes the legacy /v1/generate API and swaps in a custom subclass of vLLM’s ServingTokens, changing request/response wire shapes for renderer/teacher paths and generation defaults (max_tokens).

Overview
Unifies prime-RL’s token-level generation onto vLLM 0.20’s built-in /inference/v1/generate by removing the bespoke /v1/generate handler (serving_generate.py and route) and swapping state.serving_tokens to a new PrimeRlServingTokens subclass during server init.

PrimeRlServingTokens forwards X-data-parallel-rank into engine_client.generate, preserves routed_experts export in non-streaming responses, and adds server-side defaulting for missing sampling_params.max_tokens to avoid vLLM’s implicit 16-token cap. The chat-with-tokens path reuses the new routed-experts capture base.

The orchestrator’s compute_teacher_logprobs is updated to call /inference/v1/generate with the new payload shape (token_ids + nested sampling_params) and to parse/flatten vLLM’s upstream prompt_logprobs format; tests are updated accordingly. Dependencies are adjusted by pinning renderers==0.1.6 (PyPI) and bumping verifiers to a newer git rev, with uv.lock regenerated.

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

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

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

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

[cursor]

Missing HTTP status code check on raw response

Medium Severity

Using cast_to=httpx.Response bypasses the OpenAI SDK's status-code validation in _process_response (which short-circuits and returns the raw response for this cast type). The comment says "just hands back the raw response for us to validate ourselves," but http_response.status_code is never checked before calling GenerateResponse.model_validate_json(http_response.content). If the server returns a 4xx/5xx error, model_validate_json will attempt to parse the error body as a GenerateResponse, producing a confusing pydantic.ValidationError instead of a clear HTTP error.

Additional Locations (1)

• src/prime_rl/orchestrator/utils.py#L104-L118

 

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

[cursor]

Full upstream method copy creates fragile maintenance burden

Medium Severity

serve_tokens duplicates ~116 lines of the upstream ServingTokens.serve_tokens verbatim to inject just two changes: data_parallel_rank at line 260 and max_tokens defaulting at lines 229-237. The entire engine-input construction, sampling-params setup, logging, tracing, and streaming-vs-full dispatch are copied. Any future upstream change (e.g., new parameters, bug fixes, flow changes) will silently diverge from this fork, making it a high-risk maintenance surface.

 

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