feat: per-model client-side renderers for RL rollouts (+ /v1/generate)

docs/on_policy_distillation.md

@@ -71,7 +71,6 @@ Use this mode when you want to train from teacher-generated completions directly
 type = "sft"
 
 [orchestrator]
-use_token_client = false
 
 [orchestrator.teacher_rollout_model.client]
 base_url = ["https://your-openai-compatible-endpoint/v1"]

examples/alphabet_sort/sft_distill_hard.toml

@@ -32,7 +32,6 @@ save_adapter_separately = true
 [orchestrator]
 batch_size = 256
 rollouts_per_example = 4
-use_token_client = false
 
 [orchestrator.sampling]
 max_tokens = 512

pyproject.toml

@@ -123,7 +123,7 @@ override-dependencies = [
 
 [tool.uv.sources]
 torch = { index = "pytorch-cu128" }
-verifiers = { git = "https://github.com/PrimeIntellect-ai/verifiers.git", rev = "0760204" }
+verifiers = { git = "https://github.com/PrimeIntellect-ai/verifiers.git", branch = "renderers" }
 torchtitan = { git = "https://github.com/pytorch/torchtitan", rev = "a1fdd7e" }
 dion = { git = "https://github.com/samsja/dion.git", rev = "d891eeb" }
 transformers = { git = "https://github.com/huggingface/transformers.git", rev = "c1c3424" }

skills/entrypoints/SKILL.md

@@ -17,6 +17,20 @@ uv run rl @ examples/reverse_text/rl.toml @ examples/reverse_text/slurm_rl.toml
 uv run rl @ examples/reverse_text/rl.toml --dry-run # generate scripts without running
 ```
 
+When running from a fresh worktree, include optional extras required by the config. For example, VLM color-codeword debugging needs FlashAttention and the env packages:
+
+```bash
+uv run --extra flash-attn --extra envs rl @ configs/multimodal/rl_color_codeword_test.toml
+```
+
+Short debug runs can write large filesystem weight-broadcast snapshots. If `/` is low on space, put the debug `--output-dir` on `/dev/shm` and clean it up afterward.
+
+For hard-distillation with `orchestrator.teacher_rollout_model`, omit `[inference]`. That mode talks directly to the external OpenAI-compatible rollout endpoint and reconstructs training tokens from messages; it does not use the local direct renderer client or `/v1/generate`.
+
+For LoRA RL smoke runs against local vLLM, set `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True` in the environment or adapter load/update requests will fail.
+
+For ad-hoc `CliAgentEnv`/`ComposableEnv` sandbox smoke tests that call back through Prime Tunnel using Python stdlib `urllib`, set an explicit non-default `User-Agent` header such as `curl/7.81.0`. The tunnel can reject Python's default `urllib` user agent from inside Prime sandboxes with HTTP 403 / code 1010. Prefer uploading an agent script into the sandbox and running `python3 /tmp/agent.py` over passing a heredoc directly as a `start_background_job` command.
+
 - **Config:** `RLConfig` (`src/prime_rl/configs/rl.py`)
 - **Entrypoint:** `src/prime_rl/entrypoints/rl.py`
 - **SLURM:** yes — single-node and multi-node
@@ -49,7 +63,7 @@ uv run inference --model.name Qwen/Qwen3-0.6B --model.enforce-eager
 Always use the `inference` entrypoint — never `vllm serve` directly.
 
 Custom endpoints beyond standard OpenAI API:
-- `/v1/chat/completions/tokens` — accepts token IDs as prompt input
+- `/v1/generate` — accepts token IDs as prompt input, with optional images
 - `/update_weights` — hot-reload model weights from the trainer
 - `/load_lora_adapter` — load LoRA adapters at runtime
 - `/init_broadcaster` — initialize weight broadcast for distributed training

src/prime_rl/configs/rl.py

@@ -441,11 +441,6 @@ def validate_external_rollout_mode(self):
                 "External rollout mode does not use the local inference server."
             )
 
-        if self.orchestrator.use_token_client:
-            raise ValueError(
-                "orchestrator.use_token_client must be false when orchestrator.teacher_rollout_model is configured."
-            )
-
         return self
 
     ### Auto-setup and validate shared configs

src/prime_rl/configs/shared.py

@@ -124,6 +124,18 @@ class BaseModelConfig(BaseConfig):
         ),
     ] = False
 
+    renderer: Annotated[
+        str,
+        Field(
+            description=(
+                "Renderer to use for chat template tokenization. "
+                "Options: 'auto' (detect from tokenizer), 'qwen3', 'qwen3_vl', "
+                "'qwen3.5', 'glm5', 'glm4.5', 'minimax-m2', 'deepseek_v3', "
+                "'kimi_k2', 'kimi_k25', 'nemotron3', 'gpt_oss', 'default'."
+            ),
+        ),
+    ] = "auto"
+
     vlm: Annotated[
         "VLMConfig | None",
         Field(

src/prime_rl/inference/vllm/server.py

@@ -1,19 +1,16 @@
 from argparse import Namespace
-from http import HTTPStatus
 from typing import Any
 
 import uvloop
-from fastapi import APIRouter, Depends, Request
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
 from starlette.datastructures import State
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.openai.api_server import init_app_state
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest
 from vllm.entrypoints.utils import load_aware_call, with_cancellation
 from vllm.logger import init_logger
@@ -139,9 +136,8 @@ def resolve_tool_call_parser(model_name: str, tool_call_parser: str | None) -> s
     monkey_patch_load_lora_adapter,
     monkey_patch_tokenize_params_validation,
 )
-from prime_rl.inference.vllm.serving_chat_with_tokens import (
-    ChatCompletionRequestWithTokens,
-    OpenAIServingChatWithTokens,
+from prime_rl.inference.vllm.serving_chat_with_routed_experts import (
+    OpenAIServingChatWithRoutedExperts,
 )
 
 # NOTE: Fix harmony stop token propagation for GPT-OSS models
@@ -178,8 +174,25 @@ def models(request: Request) -> OpenAIServingModels:
 }
 
 
-def chat_with_tokens(request: Request) -> OpenAIServingChatWithTokens | None:
-    return request.app.state.openai_serving_chat_with_tokens
+def generate_handler(request: Request):
+    return request.app.state.openai_serving_generate
+
+
+@router.post("/v1/generate")
+@with_cancellation
+@load_aware_call
+async def _generate(raw_request: Request):
+    from prime_rl.inference.vllm.serving_generate import GenerateRequest
+
+    body = await raw_request.json()
+    request = GenerateRequest(**body)
+    handler = generate_handler(raw_request)
+    if handler is None:
+        return JSONResponse({"error": "generate endpoint not available"}, status_code=503)
+    result = await handler.generate(request, raw_request)
+    if isinstance(result, dict) and "error" in result:
+        return JSONResponse(result, status_code=500)
+    return JSONResponse(content=result.model_dump())
 
 
 @router.post("/pause")
@@ -234,32 +247,6 @@ async def init_broadcaster(request: Request):
     return {"status": "ok"}
 
 
-@router.post(
-    "/v1/chat/completions/tokens",
-    dependencies=[Depends(validate_json_request)],
-    responses={
-        HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-    },
-)
-@with_cancellation
-@load_aware_call
-async def _chat_with_tokens(request: ChatCompletionRequestWithTokens, raw_request: Request):
-    handler = chat_with_tokens(raw_request)
-    if handler is None:
-        return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
-    generator = await handler.create_chat_completion_with_tokens(request, raw_request)
-    if isinstance(generator, ErrorResponse):
-        return JSONResponse(content=generator.model_dump(), status_code=generator.error.code)
-
-    elif isinstance(generator, ChatCompletionResponse):
-        return JSONResponse(content=generator.model_dump())
-
-    return StreamingResponse(content=generator, media_type="text/event-stream")
-
-
 async def custom_init_app_state(
     engine_client: EngineClient,
     state: State,
@@ -269,18 +256,23 @@ async def custom_init_app_state(
     """
     Modifies init_app_state:
     1. Call the original init_app_state to set up standard state.
-    2. Replace the serving_chat with our OpenAIServingChatWithTokens wrapper.
+    2. Replace the serving_chat with our routed-experts wrapper so chat responses
+       can still expose routed experts when requested.
     """
     await init_app_state(engine_client, state, args, supported_tasks)
 
     if "generate" in supported_tasks and state.openai_serving_chat is not None:
         original_chat = state.openai_serving_chat
-        serving_chat = object.__new__(OpenAIServingChatWithTokens)
+        serving_chat = object.__new__(OpenAIServingChatWithRoutedExperts)
         serving_chat.__dict__.update(original_chat.__dict__)
         state.openai_serving_chat = serving_chat
-        state.openai_serving_chat_with_tokens = serving_chat
     else:
-        state.openai_serving_chat_with_tokens = None
+        serving_chat = None
+
+    # /v1/generate endpoint — tokens + optional images, no chat template
+    from prime_rl.inference.vllm.serving_generate import OpenAIServingGenerate
+
+    state.openai_serving_generate = OpenAIServingGenerate(engine_client, chat_handler=serving_chat)
 
 
 import vllm.entrypoints.openai.api_server

src/prime_rl/inference/vllm/serving_chat_with_routed_experts.py

@@ -0,0 +1,73 @@
+import base64
+from collections.abc import AsyncGenerator, AsyncIterator
+
+import numpy as np
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest, ChatCompletionResponse
+from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse, RequestResponseMetadata
+from vllm.logger import init_logger
+from vllm.outputs import RequestOutput
+from vllm.reasoning import ReasoningParser
+
+logger = init_logger(__name__)
+
+
+class _RoutedExpertsCapture:
+    def __init__(self, generator: AsyncGenerator[RequestOutput, None]):
+        self._generator = generator
+        self.routed_experts: dict[int, dict] = {}
+
+    def _encode_routed_experts(self, arr: np.ndarray) -> dict:
+        return {
+            "data": base64.b85encode(arr.tobytes()).decode("ascii"),
+            "shape": list(arr.shape),
+        }
+
+    async def __aiter__(self):
+        async for request_output in self._generator:
+            for output in request_output.outputs:
+                if output.routed_experts is not None:
+                    self.routed_experts[output.index] = self._encode_routed_experts(output.routed_experts)
+            yield request_output
+
+    def post_process(self, response: ChatCompletionResponse):
+        for choice in response.choices:
+            if choice.index in self.routed_experts:
+                choice.routed_experts = self.routed_experts[choice.index]
+
+
+class OpenAIServingChatWithRoutedExperts(OpenAIServingChat):
+    """OpenAI chat serving wrapper that preserves routed expert metadata."""
+
+    async def chat_completion_full_generator(
+        self,
+        request: ChatCompletionRequest,
+        result_generator: AsyncIterator[RequestOutput],
+        request_id: str,
+        model_name: str,
+        conversation,
+        tokenizer,
+        request_metadata: RequestResponseMetadata,
+        reasoning_parser: ReasoningParser | None = None,
+    ) -> ErrorResponse | ChatCompletionResponse:
+        if self.model_config.enable_return_routed_experts:
+            capture = _RoutedExpertsCapture(result_generator)
+            result_generator = capture
+        else:
+            capture = None
+
+        response = await super().chat_completion_full_generator(
+            request,
+            result_generator,
+            request_id,
+            model_name,
+            conversation,
+            tokenizer,
+            request_metadata,
+            reasoning_parser,
+        )
+
+        if capture and isinstance(response, ChatCompletionResponse):
+            capture.post_process(response)
+
+        return response