Train from inline multimodal rollouts

src/prime_rl/inference/vllm/serving_tokens.py

@@ -20,6 +20,9 @@
    ``get_max_tokens``); we keep an equivalent guard so callers that omit
    ``max_tokens`` don't truncate at vLLM's 16-token ``SamplingParams`` default.
 
+4. Raw multimodal input — port vLLM's ``content_parts`` handling onto the
+   pinned release and return the effective, expanded ``prompt_token_ids``.
+
 Everything else (request/response schema, sampling params, error handling)
 delegates to upstream so we track future vLLM changes for free.
 """
@@ -31,6 +34,7 @@
 from typing import Any
 
 from fastapi import Request
+from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker
 from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
     PromptTokenUsageInfo,
@@ -44,6 +48,7 @@
 )
 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
 from vllm.entrypoints.serve.utils.api_utils import get_max_tokens
+from vllm.inputs import TokensPrompt
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
@@ -56,6 +61,7 @@ class PrimeRlGenerateResponseChoice(GenerateResponseChoice):
 
 class PrimeRlGenerateResponse(GenerateResponse):
     choices: list[PrimeRlGenerateResponseChoice]
+    prompt_token_ids: list[int] | None = None
     # Upstream ``GenerateResponse`` doesn't declare a ``usage`` field, so the
     # parent ``ServingTokens.serve_tokens_full_generator`` constructs it and
     # Pydantic silently drops it on serialization. Declare it here so the
@@ -78,6 +84,7 @@ def post_process(self, response: GenerateResponse) -> PrimeRlGenerateResponse:
             choices=choices,
             prompt_logprobs=response.prompt_logprobs,
             kv_transfer_params=response.kv_transfer_params,
+            prompt_token_ids=getattr(response, "prompt_token_ids", None),
         )
 
 
@@ -147,6 +154,14 @@ async def _client_set_max_tokens(raw_request: Request | None) -> bool:
     return isinstance(sp, dict) and "max_tokens" in sp
 
 
+async def _raw_content_parts(raw_request: Request | None) -> list[dict[str, Any]]:
+    if raw_request is None:
+        return []
+    body = await raw_request.json()
+    parts = body.get("content_parts") if isinstance(body, dict) else None
+    return parts if isinstance(parts, list) else []
+
+
 class PrimeRlServingTokens(ServingTokens):
     """ServingTokens + DP-rank routing + compact routed experts + max_tokens defaulting."""
 
@@ -180,8 +195,8 @@ async def serve_tokens(
         # (a) inject ``data_parallel_rank`` from the inbound header into
         # ``engine_client.generate``; (b) default ``sampling_params.max_tokens``
         # to ``max_model_len - prompt_len`` when the caller didn't set it; and
-        # (c) dispatch to our overridden response builder so ``routed_experts``
-        # makes it into the JSON.
+        # (c) accept raw content parts; (d) dispatch to our overridden response
+        # builder so routed experts and effective prompt IDs make it into JSON.
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
             return error_check_ret
@@ -199,7 +214,30 @@ async def serve_tokens(
 
         # Build the engine input — features-aware (MM) or text-only fallback.
         # Identical to upstream so we keep tracking it.
-        if features := request.features:
+        content_parts = await _raw_content_parts(raw_request)
+        if content_parts and request.features:
+            return self.create_error_response("content_parts and features are mutually exclusive")
+        if content_parts:
+            tracker = AsyncMultiModalItemTracker(self.model_config)
+            mm_parser = tracker.create_parser()
+            for part in content_parts:
+                part_type = part.get("type", "")
+                url = part.get("url")
+                uuid = part.get("uuid")
+                if part_type == "image_url":
+                    mm_parser.parse_image(url, uuid)
+                elif part_type == "audio_url":
+                    mm_parser.parse_audio(url, uuid)
+                elif part_type == "video_url":
+                    mm_parser.parse_video(url, uuid)
+            mm_data, mm_uuids = await tracker.resolve_items()
+            prompt = TokensPrompt(prompt_token_ids=request.token_ids)
+            if mm_data:
+                prompt["multi_modal_data"] = mm_data
+            if mm_uuids:
+                prompt["multi_modal_uuids"] = mm_uuids
+            (engine_input,) = await self.online_renderer.renderer.render_cmpl_async([prompt])
+        elif features := request.features:
             from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item
             from vllm.inputs import mm_input
             from vllm.multimodal.inputs import (
@@ -259,7 +297,7 @@ async def serve_tokens(
             sampling_params.max_tokens = get_max_tokens(
                 max_model_len=self.model_config.max_model_len,
                 max_tokens=None,
-                input_length=len(request.token_ids),
+                input_length=self._extract_prompt_len(engine_input),
                 default_sampling_params=diff_sp,
                 override_max_tokens=override,
             )
@@ -301,7 +339,11 @@ async def serve_tokens(
             )
 
         return await self.serve_tokens_full_generator(
-            request, result_generator, request_id, model_name, request_metadata
+            request,
+            result_generator,
+            request_id,
+            model_name,
+            request_metadata,
         )
 
     async def serve_tokens_full_generator(  # type: ignore[override]
@@ -352,5 +394,6 @@ async def serve_tokens_full_generator(  # type: ignore[override]
 
         if final_capture.final_res is not None:
             response.usage = _build_usage(final_capture.final_res)
+            response.prompt_token_ids = final_capture.final_res.prompt_token_ids
 
         return response

src/prime_rl/multimodal/__init__.py

@@ -0,0 +1,9 @@
+from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, MultimodalAdapter
+from prime_rl.multimodal.registry import get_multimodal_adapter
+
+__all__ = [
+    "ForwardPolicy",
+    "MaterializedMM",
+    "MultimodalAdapter",
+    "get_multimodal_adapter",
+]

src/prime_rl/multimodal/base.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import torch
+from PIL.Image import Image
+
+
+@dataclass(frozen=True)
+class ForwardPolicy:
+    pass_position_ids: bool = True
+    requires_mm_token_type_ids: bool = False
+    defer_context_parallelism: bool = False
+
+
+@dataclass(frozen=True)
+class MaterializedMM:
+    kwargs: dict[str, torch.Tensor]
+    forward_policy: ForwardPolicy
+
+
+class MultimodalAdapter(Protocol):
+    model_types: frozenset[str]
+    forward_policy: ForwardPolicy
+
+    def materialize(
+        self,
+        image_processor: Any,
+        images: list[Image],
+        placeholder_lengths: list[int],
+    ) -> MaterializedMM: ...
+
+
+def required_tensors(values: Any, keys: tuple[str, ...]) -> dict[str, torch.Tensor]:
+    data = dict(values)
+    missing = [key for key in keys if key not in data]
+    if missing:
+        raise ValueError(f"Image processor did not return {', '.join(missing)}")
+    return {key: torch.as_tensor(data[key]).contiguous() for key in keys}

src/prime_rl/multimodal/kimi_k25.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Any
+
+from PIL.Image import Image
+
+from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, required_tensors
+
+
+class KimiK25Adapter:
+    model_types = frozenset({"kimi_k25"})
+    forward_policy = ForwardPolicy()
+
+    def materialize(
+        self,
+        image_processor: Any,
+        images: list[Image],
+        placeholder_lengths: list[int],
+    ) -> MaterializedMM:
+        preprocess = getattr(image_processor, "preprocess", None)
+        if preprocess is None:
+            raise ValueError("Kimi image processor is missing preprocess")
+        media = [{"type": "image", "image": image} for image in images]
+        kwargs = required_tensors(
+            preprocess(media, return_tensors="pt"),
+            ("pixel_values", "grid_thws"),
+        )
+        lengths = [1] * len(kwargs["grid_thws"].reshape(-1, 3))
+        if lengths != placeholder_lengths:
+            raise ValueError(
+                f"Kimi image placeholder lengths differ from vLLM: expected {placeholder_lengths}, got {lengths}"
+            )
+        return MaterializedMM(kwargs=kwargs, forward_policy=self.forward_policy)

src/prime_rl/multimodal/qwen_vl.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Any
+
+from PIL.Image import Image
+
+from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, required_tensors
+
+
+class QwenVLAdapter:
+    model_types = frozenset({"qwen3_vl", "qwen3_vl_moe", "qwen3_5", "qwen3_5_moe"})
+    forward_policy = ForwardPolicy(
+        pass_position_ids=False,
+        requires_mm_token_type_ids=True,
+        defer_context_parallelism=True,
+    )
+
+    def materialize(
+        self,
+        image_processor: Any,
+        images: list[Image],
+        placeholder_lengths: list[int],
+    ) -> MaterializedMM:
+        kwargs = required_tensors(
+            image_processor(images=images, return_tensors="pt"),
+            ("pixel_values", "image_grid_thw"),
+        )
+        merge_size = int(image_processor.merge_size)
+        lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)]
+        if lengths != placeholder_lengths:
+            raise ValueError(
+                f"Qwen image placeholder lengths differ from vLLM: expected {placeholder_lengths}, got {lengths}"
+            )
+        return MaterializedMM(kwargs=kwargs, forward_policy=self.forward_policy)

src/prime_rl/multimodal/registry.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from prime_rl.multimodal.base import MultimodalAdapter
+from prime_rl.multimodal.kimi_k25 import KimiK25Adapter
+from prime_rl.multimodal.qwen_vl import QwenVLAdapter
+
+_ADAPTERS = (QwenVLAdapter(), KimiK25Adapter())
+_BY_MODEL_TYPE: dict[str, MultimodalAdapter] = {
+    model_type: adapter for adapter in _ADAPTERS for model_type in adapter.model_types
+}
+
+
+def get_multimodal_adapter(model_type: str) -> MultimodalAdapter:
+    try:
+        return _BY_MODEL_TYPE[model_type]
+    except KeyError as exc:
+        raise NotImplementedError(f"Raw image training is not implemented for model type {model_type!r}") from exc

src/prime_rl/orchestrator/trajectories.py

@@ -8,9 +8,8 @@
 live on `vf.Trace` itself.
 
 Training is renderer-only across every mode (RL/OPD student, SFT teacher), so every node
-always carries its tokens — no backfill needed. For multimodal rollouts the branch also carries
-the images it introduced (`branch.multi_modal_data`), rebuilt here into the flat `mm_kwargs` /
-`mm_token_type_ids` the trainer forwards.
+always carries its tokens — no backfill needed. Multimodal RL keeps the inline image URLs on
+the messages and pairs them with vLLM's expanded image-token runs here.
 """
 
 from __future__ import annotations
@@ -20,33 +19,48 @@
 import numpy as np
 import verifiers.v1 as vf
 
-from prime_rl.transports.batch import TrainingSample
-from prime_rl.transports.batch.types import EncodedTensor, RoutedExperts
+from prime_rl.transports.batch import MMImageRef, MMRefs, TrainingSample
+from prime_rl.transports.batch.types import RoutedExperts
 from prime_rl.utils.logger import get_logger
 
 
-def _to_numpy(val) -> np.ndarray:
-    """A renderer mm item value (torch tensor or numpy array) -> a contiguous numpy array."""
-    if hasattr(val, "detach"):  # torch tensor
-        val = val.detach().cpu().numpy()
-    return np.ascontiguousarray(val)
-
-
-def _encode_mm_kwargs(mm_items: dict[str, list[dict]]) -> dict[str, EncodedTensor] | None:
-    """Concatenate the branch's per-image renderer items into the flat `mm_kwargs` the trainer
-    forwards — one `EncodedTensor` per kwarg key (e.g. `pixel_values`, `image_grid_thw`), images
-    cat'd along dim 0 in branch token order. Model-agnostic: the keys are whatever the processor
-    emits. Returns None when there are no items."""
-    bins: dict[str, list[np.ndarray]] = {}
-    for items in mm_items.values():  # per modality
-        for item in items:  # per image
-            for key, val in item.items():
-                bins.setdefault(key, []).append(_to_numpy(val))
-    encoded: dict[str, EncodedTensor] = {}
-    for key, arrs in bins.items():
-        arr = np.concatenate(arrs, axis=0)
-        encoded[key] = EncodedTensor(dtype=str(arr.dtype), shape=list(arr.shape), data=arr.tobytes())
-    return encoded or None
+def _image_urls(branch: vf.Branch) -> list[str]:
+    urls: list[str] = []
+    for node in branch.nodes:
+        content = node.message.content
+        if not isinstance(content, list):
+            continue
+        for part in content:
+            if getattr(part, "type", None) == "image_url":
+                urls.append(part.image_url.url)
+    return urls
+
+
+def _image_runs(token_types: list[int]) -> list[tuple[int, int]]:
+    runs: list[tuple[int, int]] = []
+    start: int | None = None
+    for index, token_type in enumerate([*token_types, 0]):
+        if token_type == 1 and start is None:
+            start = index
+        elif token_type != 1 and start is not None:
+            runs.append((start, index - start))
+            start = None
+    return runs
+
+
+def _build_mm_refs(urls: list[str], token_types: list[int]) -> MMRefs | None:
+    runs = _image_runs(token_types)
+    if len(urls) != len(runs):
+        raise ValueError(
+            f"Inline image count does not match expanded placeholder runs: images={len(urls)}, runs={len(runs)}"
+        )
+    if not urls:
+        return None
+    return MMRefs(
+        images=[
+            MMImageRef(url=url, offset=offset, length=length) for url, (offset, length) in zip(urls, runs, strict=True)
+        ]
+    )
 
 
 def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExperts | None:
@@ -115,19 +129,21 @@ def trace_to_samples(trace: vf.Trace, *, env_name: str = "") -> list[TrainingSam
     `branch.sampled_mask` / `branch.logprobs`), so a sample carries it directly: `mask` marks
     the trainable (model-sampled) tokens, the context tokens between completions stay masked
     out. Errored traces are dropped upstream (`TrainSink.process_episode`), so no error
-    handling happens here. A branch carrying images also gets `mm_kwargs` (the concatenated
-    pixel tensors) and `mm_token_type_ids` (`branch.mm_token_type_ids`, computed from the
-    trace's renderer-stamped `mm_token_type_id_map`). Branches with no sampled tokens
-    (e.g. an openai client carrying none) yield nothing.
+    handling happens here. A branch carrying images gets raw image refs paired with expanded
+    placeholder ranges and `mm_token_type_ids`. Branches with no sampled tokens yield nothing.
     """
     samples: list[TrainingSample] = []
     trained_loss_nodes: dict[str, set[int]] = {"rl": set(), "ce": set(), "ref_kl": set()}
     for branch, mask in iter_trainable_branches(trace):
         token_ids = branch.token_ids
-        mm_kwargs: dict[str, EncodedTensor] | None = None
-        mmd = branch.multi_modal_data
-        if mmd is not None:
-            mm_kwargs = _encode_mm_kwargs(mmd.mm_items)
+        mm_token_type_ids: list[int] | None = None
+        mm_refs: MMRefs | None = None
+        image_urls = _image_urls(branch)
+        if image_urls:
+            mm_token_type_ids = branch.mm_token_type_ids
+            if mm_token_type_ids is None:
+                raise ValueError("Inline images have no expanded multimodal prompt tokens")
+            mm_refs = _build_mm_refs(image_urls, mm_token_type_ids)
         samples.append(
             TrainingSample(
                 token_ids=token_ids,
@@ -136,8 +152,8 @@ def trace_to_samples(trace: vf.Trace, *, env_name: str = "") -> list[TrainingSam
                 temperatures=[],  # filled by TrainSink.process_group
                 env_name=env_name,
                 ref_logprobs=branch.reference_logprobs,
-                mm_kwargs=mm_kwargs,
-                mm_token_type_ids=branch.mm_token_type_ids,
+                mm_refs=mm_refs,
+                mm_token_type_ids=mm_token_type_ids,
                 routed_experts=_encode_routed_experts(branch.routed_experts, len(token_ids)),
                 rl_weights=_loss_weights(branch, "rl", trained_loss_nodes["rl"]),
                 ce_weights=_loss_weights(branch, "ce", trained_loss_nodes["ce"]),