Support v1 raw multimodal image offload

[eligotts]

Summary

Wires Prime-RL v1 training onto the raw multimodal image offload contract used by the companion Renderers and Verifiers PRs.

Companion PRs:

• Renderers: [codex] Support raw image refs for multimodal rendering renderers#89
• Verifiers: [codex] Support raw image offload in v1 train client verifiers#1746

Current design:

• Prime-RL config owns only the image offload directory. There is no inline/offload mode knob; v1 raw multimodal uses offloaded image files.
• Renderers/verifiers are the only writer side that needs VF_RENDERER_IMAGE_OFFLOAD_DIR; they rewrite image parts to shared file:// assets and emit JSON-safe raw descriptors.
• Raw descriptors carry raw_image_uri as the sole image locator plus adapter-owned metadata (family, layout_fingerprint, payload). raw_image_id was removed to avoid two sources of truth.
• The vLLM payload is information-complete: renderers serialize mmraw:<base64-json-payload> refs carrying raw_image_uri, hash, family, fingerprint, modality, and payload. Raw-mm ref and descriptor version markers were removed because this branch has no compatibility contract yet.
• Prime inference materializes images from the URI in the request payload, so it no longer needs or derives VF_RENDERER_IMAGE_OFFLOAD_DIR.
• The trainer already consumes exact file:// URIs through MMRefs; build_mm_refs now builds those URIs from the raw descriptors instead of rescanning messages.
• Legacy processed multimodal tensors, cache-only None slots, inline base64 storage, and root-plus-id lookup paths are not supported.

Submodule pins:

• Renderers: a49e0fc (branch merged with renderers main at 5904fa2)
• Verifiers: 9bc3cc36 (branch merged with verifiers main at 5885ab9c)

Latest merge (7ea71a96f): merged prime-rl origin/main (22 commits, incl. the dispatch-gate fix #2937, 1-indexed steps #2896, vLLM 0.24 #2921, and the ExperimentalConfig removal #2931 — resolved in favor of keeping [multimodal] without the placeholders). Both dep submodules were merged with their own main and re-pinned.

Review hardening (latest push)

• MMRefs restructured from a modality-keyed descriptor dict + parallel lists to images: list[MMImageRef] (descriptor, hash, uri, placeholder range per image). Alignment is structural, so the count validators collapsed to one zip at build_mm_refs; truncation keeps a prefix of whole images and never splits a placeholder.
• build_mm_refs enforces sorted, non-overlapping placeholders; the orchestrator additionally validates every image ref's placeholder span lands on image-typed tokens (mm_token_type_ids == 1) before a sample ships, so offset drift anywhere upstream fails loudly.
• Kimi adapter reads the actual processor layout and relies on the fingerprint comparison alone (no hard equality against the renderer constant), matching the Qwen adapter.
• Renderers: fixed the bridge merge mutating the previous turn's persisted sidecar in place; Qwen resize math is now imported from transformers instead of ported; layout parity tests run against the real Qwen3-VL and Kimi-K2.5 processors.
• Verifiers: ingress offload covers every image part shape (image_url nested/direct, image direct, typed pydantic parts); renderers hard-require file:// sources.

Validation

• Renderers: uv run ruff check renderers/configs.py renderers/base.py renderers/client.py renderers/qwen3_vl.py renderers/qwen35.py renderers/kimi_k25.py renderers/mm_store.py renderers/__init__.py tests/test_renderer_config.py tests/test_multimodal_output_modes.py tests/test_client.py
• Renderers: uv run ruff format --check renderers/configs.py renderers/base.py renderers/client.py renderers/qwen3_vl.py renderers/qwen35.py renderers/kimi_k25.py renderers/mm_store.py renderers/__init__.py tests/test_renderer_config.py tests/test_multimodal_output_modes.py tests/test_client.py
• Renderers: uv run pytest tests/test_renderer_config.py tests/test_client.py::test_generate_serializes_raw_mm_refs tests/test_multimodal_output_modes.py -q -> 38 passed
• Verifiers commit hooks: ruff check, ruff format, ty (ci parity) passed
• Prime-RL: uv run ruff check src/prime_rl/entrypoints/rl.py src/prime_rl/inference/server.py src/prime_rl/inference/vllm/serving_tokens.py src/prime_rl/multimodal/schema.py src/prime_rl/orchestrator/trajectories.py src/prime_rl/utils/mm.py tests/unit/inference/test_serving_tokens.py tests/unit/orchestrator/test_batch.py tests/unit/utils/test_mm.py
• Prime-RL focused tests: uv run pytest tests/unit/utils/test_mm.py tests/unit/inference/test_serving_tokens.py::test_materialize_raw_image_ref_uses_generic_family_payload tests/unit/orchestrator/test_batch.py::test_prepare_sample_rejects_overlong_raw_mm_refs -> 4 passed
• Post-ref-version cleanup: uv run pytest tests/unit/inference/test_serving_tokens.py::test_materialize_raw_image_ref_uses_generic_family_payload -> 1 passed
• Latest sync: git diff --check, uv lock --check, and uv sync --all-extras --locked passed on the branch after merging origin/main.
• Latest push: renderers full suite 2182 passed (incl. Kimi-K2.5 parity/byte-parity against the pinned processor), verifiers 839 passed (pre-existing test_envs/test_opencode_rlm_env failures excluded, reproduced on the base branch), prime-rl tests/unit minus GPU-only tests/unit/train 410 passed.

Latest sync:

• Fast-forwarded local /home/ubuntu/prime-rl-main and merged latest prime-rl origin/main (f19ba721a) into this branch.
• Bumped Verifiers to 9b3e7ee, which contains both this PR's previous Verifiers pin (22c7cf4c) and current main's Verifiers pin (22d6333a). This resolves the deps/verifiers merge conflict and satisfies lean-v1's verifiers>=0.1.15.dev402 requirement.
• Kept Renderers at a7953b9, which contains main's renderers pin (a5efbb), the processed multimodal renderer output mode, and the cleaned lockfile diff.
• Refreshed uv.lock for the optional Renderers vision extra metadata so CI's uv sync --all-extras --locked path is in sync.

Update: merged main (2119bdbd8)

Reconciled with prime-rl main (79 commits: MXFP8, FA3-default, CP for hybrid+VLM SFT, VLM SFT, Task/TaskData, NIXL/ModelExpress, the seq_lens packed-sample contract, get_trace_path rollout-save refactor, deployment-config restructure, v0.7.0). Both companion PRs were reconciled with their own mains first and re-pinned (renderers e64cc58, verifiers 2b1627d03).

Resolution decisions:

• orchestrator/trajectories.py: main's iter_trainable_branches() (fork-dedup) loop is authoritative; this branch's _validate_image_spans + build_mm_refs spliced into it unchanged.
• trainer/model.py: kept this branch's ForwardPolicy adapter over main's "image_grid_thw" not in mm_kwargs string check in forward() — Kimi K2.5 (pass_position_ids_with_mm=True) and Qwen VL (False, requires_mm_token_type_ids=True) need opposite behavior, which a key-presence check cannot express. Main's MXFP8/CP/quantization work merged around it untouched.
• trainer/batch.py: took main's new mm_kwargs bin-packing machinery, plus a guard that raw-ref (mm_refs) samples never pack — with text or each other — because their placeholder offsets are sample-relative and _materialize_bin carries only the first sample's refs; packing would silently drop or misalign images (main's own packing test packed a text sample into an mm bin, which would have dropped the refs on this branch's path). Rewrote that test to assert the raw-ref contract.
• Two auto-merge-dropped imports caught by tests, not conflict markers: numpy in trainer/batch.py (main's _routed_experts_row_size rewrite) and FinishReason/Usage in verifiers graph.py.
• Verifiers bump crosses verifiers' interleaving-agents refactor (Add router_url support for elastic inference pool #2049), which renamed textarena's seat config: migrated env.agent → env.player in configs/basic/wordle/rl.toml and examples/basic/wordle/rl.toml (prime-rl main will need the same migration at its next verifiers bump).
• The seq_lens/mm-truncation interaction flagged in review scoping is safe: prepare_sample computes seq_lens=[len(input_ids)] after _truncate_mm_refs.

Validation: uv lock --check clean; uv sync --all-packages --all-extras clean; tests/unit minus GPU-only tests/unit/train: 393 passed; tests/unit/train collects (160 tests) with no import errors. Companion suites: renderers 2578 passed, verifiers full suite passing (minus PRIME_API_KEY-gated e2e).

Update: SFT renderer default + serving-layer materialize cache (a2fdf6a8, 76761553)

• SFT defaults to processed renderer output. The renderers library defaults multimodal_output='raw' (built for the RL offload path), but SFT's data path consumes processed pixel tensors straight from the renderer and has no raw-ref materializer. SFTConfig now defaults its renderer to 'processed' and hard-rejects an explicit 'raw'. RL keeps the library default of 'raw', untouched.
• Serving-layer materialize cache. Every request carries a raw ref for every image in its prompt (prior turns included), so multi-turn rollouts re-materialized the same images every turn — CPU work (read + sha256 + PIL decode + HF processor forward) vLLM's own caches can't reach, since production happens in PrimeRlServingTokens before vLLM sees the request. Added a byte-bounded, single-flight LRU keyed by (raw_ref, expected_placeholder_length, processor_model_name) (content-addressed — hits are provably identical work; evicted entries go cold, never stale). _decode_raw_mm_kwargs now gathers a request's images concurrently instead of sequentially. Failures are never cached. One knob: PRIME_RL_MM_MATERIALIZE_CACHE_GB (default 2.0, 0 disables — byte-identical to the uncached path, the kill switch). Logs mm materialize cache: hits/misses/hit_rate/bytes/evictions every 1000 lookups; the monitor-run skill documents the signature. Four unit tests: hit-skips-work, byte-budget eviction, failure-not-cached, single-flight.

Validation: full tests/unit minus GPU-only tests/unit/train: 397 passed.

Update: post-merge audit fixes (e0c22ee5)

An adversarial audit pass over the reconciliation (independent per-repo reviews diffing the merge against both parents, plus the first CI run) surfaced and fixed:

• SFT Qwen-VL position-ids regression (real bug): forward()'s no-policy ForwardPolicy() fallback inverted main's image_grid_thw gate — SFT callers don't thread an adapter policy, so multimodal SFT would have passed packed 1D position_ids and silently skipped the model's internal MRoPE construction. The fallback now reproduces the key-presence heuristic; RL's explicit adapter policies unaffected.
• Materialize-cache cancellation poisoning (real bug): a disconnecting owner request set CancelledError on the shared single-flight future, spuriously cancelling every deduped awaiter. Materialization now runs as a detached task; awaiters shield the shared future.
• CI runs the full tests/unit on GPU runners, not the CPU subset validated locally: adapted main's two eager-mm_kwargs packer tests to the raw-ref contract (raw samples never pack, even within a run — _mm_sample now carries mm_refs), added the missing seq_lens kwarg to the branch-only forward-policy test, extended the batch no-pack test to the mm+mm direction, and migrated main's new configs/ci/nightly-fft/wordle.toml (added after the merge point; CI tests the PR merged with latest main) to the env.player seat. Merged main again (45488ca04) to pick that file up.
• Hardening: _is_multimodal_sample also recognizes eager mm_kwargs samples (keeps main's packing-compatibility machinery live-correct); multimodal_sample_error requires mm_token_type_ids on raw mm_refs samples.

Audit verdicts elsewhere: renderers merge CLEAN (independent re-verification of every conflict site, 182-test focused suite + full 2578 green); verifiers merge CLEAN (the SkipJsonSchema adoption was verified necessary — the plain annotation breaks Trace.model_json_schema() on the numpy field).

Validation: tests/unit minus tests/unit/train/models: 484 passed locally (tests/unit/train/models kernels segfault on the local Blackwell box in a mamba backward test untouched by this PR — green in CI).

Update: restored main's multimodal packing (6a1a54bc)

The main merge had resolved the mm-packing conflict too conservatively — raw-ref samples never packed, regressing main's feature where multimodal samples pack with text spans and compatible mm samples from the same run/LoRA. Restored for the mm_refs representation:

• can_add: raw-ref samples are pack-compatible when they share an adapter family (the materializer enforces exactly one family per micro batch; grids/sizes vary freely within a family) — the simpler analogue of main's eager dtype/trailing-dim check. Run/LoRA gating, seq_lens boundaries, and cross-rank modality alignment as on main.
• _materialize_bin: merges refs across the bin's samples, rebasing each sample's placeholder offsets by its start position in the packed token stream. The trainer-side adapter already materializes a list of refs into concatenated tensors, so packed batches reach forward() shaped exactly like main's packed eager batches.
• Packing tests restored to main's semantics (pack within run with rebased offsets, family mismatch splits, never across runs).

Also mapped adapter validation errors (fingerprint/grid/placeholder-length mismatches) to invalid_mm_image_ref 400s in serving_tokens instead of escaping as 500s (ported from the inline bundle's review).

Update: ref-format and cache-key parity with the inline bundle (7c58fd6 / fd3f379d)

Three parity changes so the offload and inline bundles differ only in image storage:

• Unwrapped ref payload: raw_mm_ref serializes compact JSON instead of base64-wrapping it. Worth little here (the payload is a short file:// URI) but it keeps the ref format identical to the inline bundle, where the wrapper cost ~32 KiB per image slot.
• Materialize-cache key: adopted the inline bundle's hardened key — digest the whole ref rather than trusting the caller's outer metadata, plus feature_modality and trust_remote_code. The digest matters less here (keying on a short content-addressed ref was already sound), but trust_remote_code closes a real gap: it selects a different image processor, so two values could previously alias one cache entry. Ported the anti-aliasing test.
• Explicit env inheritance at the orchestrator launch: the subprocess previously inherited os.environ only as a side effect of build_run_asset_env returning dict(os.environ) — the same implicit coupling that broke inheritance outright in the inline bundle when that helper was removed. inherited_env is now spread first, as at the trainer and inference launches.

Also fixed a test-fixture flaw found while porting: _mm_features colored images by filename length, so two same-length names produced byte-identical images and a forged-hash case couldn't be expressed. It now seeds content explicitly.

Update: synced with main (2c3f6d6e8)

All three repos in this bundle re-merged with their own mains.

prime-rl (18 commits): vLLM 0.26 bump (#3166), task-centric orchestrator config restructure (train.source.env), RAE / role-conditioned advantage estimation (#3136), verifiers bump adopting runtime-on-agent + textarena player (#3137).

• serving_tokens.py: vLLM 0.26 moved the disagg serving modules under scale_out.token_in_token_out and changed mm_input() to take MultiModalKwargsItems. Adopted both — the wrapper is a genuine API change, not an eager-path detail — while leaving out main's decode_mm_kwargs_item / MultiModalKwargsItem imports, which only serve the processed-tensor path this branch replaces.
• wordle configs: main has since adopted the textarena player seat and restructured to [orchestrator.train.source.env], so main's versions supersede this branch's earlier env.agent → env.player migration.
• configs/orchestrator.py + rl.py: insertion collisions only — kept multimodal alongside main's rollout_transport and main's removal of the deprecated [[env]]/[sampling] shim.
• test_qwen3_vl_e2e.py stays deleted (main only refreshed its vLLM import paths; it round-trips the eager processed-features payload v1 rejects by design).

Validation: tests/unit minus GPU-only model kernels: 409 + 84 passed on vLLM 0.26; uv lock --check clean; ruff clean.

---

Note

High Risk
Touches the full RL multimodal path (transport, packing/truncation, inference materialization, trainer forward) and depends on companion renderer/verifier behavior; misaligned refs or offload paths can fail runs or silently mask loss via placeholders.

Overview
Prime-RL v1 multimodal RL now uses a raw image offload contract aligned with renderers/verifiers: images land once under a run-scoped directory, training samples carry mm_refs (descriptor + file:// URI + placeholder span) instead of encoded mm_kwargs, and processed pixel tensors are rejected on the v1 path.

Config and launcher: Shared [multimodal].offload_dir propagates to trainer, orchestrator, and inference; the RL launcher sets VF_RENDERER_IMAGE_OFFLOAD_DIR (protected from env_vars) for the orchestrator and SLURM templates. Trainer missing_mm_image_policy can zero-loss placeholder when assets disappear; SFT defaults renderer output to processed and rejects raw.

Inference: PrimeRlServingTokens materializes every raw ref (hash/fingerprint checks, Qwen VL + Kimi K2.5 adapters) before vLLM, with a byte-bounded LRU cache (PRIME_RL_MM_MATERIALIZE_CACHE_GB) and invalid_mm_image_ref 400s on validation failures.

Trainer/orchestrator: build_mm_refs + placeholder-span validation before samples ship; truncation/packing respect whole-image boundaries and rebase offsets when packing same-family refs. RawImageMaterializer and ForwardPolicy drive trainer-side materialization and correct position-id / mm_token_type_ids behavior per family. Metrics add time/mm_materialize and image materialize/placeholder counts; orchestrator timing maps verifiers’ agent phase to existing generation metrics.

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

[samsja]

can we put this in the multi modal config ?

[samsja]

not a fan of seting up env var, why can't we jsut read from the env var instead and dinamycally change it in the code if the env var is not present ?

[samsja]

I dont fully get this part

[eligotts]

Tech Spec: Serving-Layer Materialization Cache for Raw Multimodal Image Refs

Status: approved design, ready to implement
Target branch: feat/v1-raw-mm-offload (PR #2836) — this feature modifies code that only exists on that branch
Primary file: src/prime_rl/inference/vllm/serving_tokens.py
Estimated size: ~150 lines of implementation + 3–4 unit tests, single file plus tests

---

1. Background: the raw multimodal offload architecture

You need this context to understand why the cache is needed and where it sits. Skim the
three companion PRs if you want the full history:

• prime-rl #2836 — "Support v1 raw multimodal image offload"
• renderers #89 — raw image refs for rendering
• verifiers #1746 — raw image offload in the v1 train client

How images flow through the system

Multimodal RL rollouts (e.g. browser agents taking screenshots) used to ship processed
image tensors (pixel_values) through every layer. The offload redesign replaced that
with lightweight refs:

1. Ingress (verifiers): every image content part in a message is rewritten to a
   file:// URL pointing into a shared run-assets directory. Files are content-addressed
   (sha256(bytes) as the filename), written once, deduped by construction.
   See deps/verifiers/verifiers/utils/multimodal.py::prepare_images_inplace.
2. Rendering (renderers): the renderer computes the image's token layout (grid /
   placeholder length) without running an image processor — pure math from the image
   dimensions plus baked layout constants (see renderers/qwen3_vl.py::describe_qwen_image_layout).
   It emits a JSON-safe descriptor per image.
3. Wire (renderers client): renderers/client.py::_build_vllm_mm_features serializes
   each image slot as a raw ref string: mmraw:<urlsafe-b64-json> containing
   {family, fingerprint, modality, mm_hash, payload, raw_image_uri}. Every request
   carries a ref for every image in the prompt — current and prior turns.
   See renderers/mm_store.py::raw_mm_ref / split_raw_mm_ref.
4. Inference (prime-rl, THIS SPEC'S TARGET): our custom vLLM serving handler
   PrimeRlServingTokens (src/prime_rl/inference/vllm/serving_tokens.py) receives the
   refs and materializes each one: read the file, verify its sha256 against
   mm_hash, PIL-decode, run the HF image processor via a family adapter
   (src/prime_rl/multimodal/adapters/), and wrap the result as a vLLM
   MultiModalKwargsItem. Only then does the request enter vLLM proper via mm_input().
5. Training (prime-rl trainer): independently materializes from the same file://
   assets (src/prime_rl/utils/mm.py::RawImageMaterializer). Not affected by this spec.

Key vocabulary:

• mm_hash — sha256(image file bytes)[:32]. Content identity of an image.
• raw_ref — the mmraw:... string. Content-addressed: it embeds the mm_hash,
  the layout fingerprint (a hash of the image-processor config the renderer assumed),
  the adapter family (e.g. qwen_vl, kimi_k25), the adapter payload (e.g.
  image_grid_thw), and the raw_image_uri. Two identical ref strings describe
  byte-identical materialization work.
• materialization — read + hash-verify + decode + HF processor + wrap. The expensive
  CPU pipeline this spec caches.

2. The problem

Read _decode_raw_mm_kwargs (serving_tokens.py:274) and
_materialize_raw_image_ref_sync (serving_tokens.py:218). Current behavior:

• Materialization runs for every image slot on every request. Since the client sends
  refs for prior-turn images too, a 20-turn rollout whose prompt accumulates 5 images
  pays ~100 materializations where 5 would do. Multiply by group size
  (rollouts_per_example) hammering the same prompt images.
• The per-image cost is real CPU work: file read + full sha256 + PIL decode + an HF
  image-processor forward (tens of ms and a large float32 pixel_values allocation per
  image; a max-resolution Qwen3-VL image is ~16M pixels).
• The images in one request are materialized sequentially — the await asyncio.to_thread(...) sits inside a for loop (serving_tokens.py:286), so request
  latency is sum-of-images instead of max-of-images.

None of vLLM's own caches can fix this, because the work happens in our handler
before vLLM sees the request. See §3.

3. The three-cache model (why this cache is not redundant)

There are three caches in the pipeline. Each one's hit avoids a different expensive step.
Content of the first two is identical; that is intentional, not redundant:

Cache	Stores	A hit avoids	Lives	Knob
Ours (this spec)	MultiModalKwargsItem (pixel_values + grids)	Producing the tensors (read/verify/decode/HF-process)	our handler, API frontend process, host RAM	new (this spec)
vLLM processor cache	same content	Moving the tensors across the P0(frontend)→P1(engine core) process boundary	vLLM-internal; lru = mirrored per-process pair, shm = shared ring buffer	--mm-processor-cache-gb, --mm-processor-cache-type
vLLM encoder cache	vision-encoder output embeddings	Encoding (vision-tower GPU forward)	per engine core (= per DP rank), GPU memory	encoder budget (separate)

Produce → move → encode. In the old processed-tensor world, vLLM's processor cache
covered both produce and move because vLLM did the producing. The raw-ref design
moved production into PrimeRlServingTokens, out of that cache's reach. This spec
rebuilds the produce guard at the layer where production now happens.

Composition facts your implementation can rely on (verified against vLLM 0.24 in
.venv/lib/python3.12/site-packages/vllm/):

• We always hand vLLM full kwargs + hashes. vLLM's input processor has an explicit
  injection path for externally-processed kwargs
  (vllm/v1/engine/input_processor.py, the method whose docstring begins "Inject
  pre-processed mm_kwargs into the processor cache") — it inserts our items into its own
  cache by mm_hash and dedupes the P0→P1 hop on repeats. Works identically whether the
  engine is deployed with lru, shm, or processor cache disabled. Our cache makes
  no assumption about, and needs no coordination with, any of it.
• DP topology: one API frontend process serves all DP ranks of an engine (prime-rl pins
  rollouts to DP ranks via the X-data-parallel-rank header — see
  src/prime_rl/utils/client.py:246). Our cache lives in that frontend, so it is
  naturally shared across all DP ranks of the engine. With multiple engines behind
  the vllm-router, each engine's frontend has its own independent cache; router
  stickiness concentrates a rollout's turns on one engine, so per-engine caches capture
  the locality. No routing changes are needed or wanted.

4. Rejected alternatives (do not build these)

• Engine-native "pointer" scheme: send None for a kwargs slot and let vLLM resolve
  the mm_hash from its internal cache. Rejected: vLLM's P0/P1 caches stay consistent
  through a strict internal lockstep-eviction protocol; we sit outside it and cannot
  query membership without racing evictions. A dangling pointer fails the request
  (a rollout dies because of cache weather), and the recovery path requires keeping all
  the materialization machinery anyway. Strictly more complexity and fragility.
• Shared-memory cache across processes: solves a problem we don't have (we run one
  API frontend per engine). Revisit only if --api-server-count > 1 ever ships here.
• Disk tier for processed tensors: the content-addressed raw file already is the
  disk tier; the processor run is the only cost worth caching and RAM is its right home.

---

5. Design

An in-process, byte-bounded, single-flight LRU in serving_tokens.py, keyed by
content-addressed identity, holding finished MultiModalKwargsItems, with concurrent
per-request materialization.

5.1 Key

(raw_ref: str, expected_placeholder_length: int, processor_model_name: str)

• raw_ref is content-addressed by construction (§1 vocabulary), so identical keys ⇒
  identical work ⇒ hits are sound with zero extra bookkeeping.
• expected_placeholder_length is included so a ref arriving with a different claimed
  placeholder length misses, re-materializes, and hits the existing validation error
  in the adapter (materialize_for_vllm checks it) — preserving fail-loud instead of
  serving a value that skipped that check. In the healthy pipeline it is a deterministic
  function of the ref, so it never fragments the cache.
• processor_model_name guards hypothetical multi-model serving; today it's constant
  per process.

Note the key is deliberately stricter than mm_hash: identical pixels at two
different file:// paths are two entries. In the normal pipeline this never happens
(the offload store is content-addressed, same bytes → same path → same ref).

5.2 Value and byte accounting

The finished MultiModalKwargsItem returned by
adapter.materialize_for_vllm(...), plus its size in bytes.

For sizing, reuse vLLM's own helper rather than writing tensor-walking code:

from vllm.multimodal.cache import MultiModalCache
nbytes = MultiModalCache.get_item_size(item)

(Verify this import against the installed vLLM at implementation time; it's present in
0.24 at vllm/multimodal/cache.py:120 and is exactly what vLLM uses to size its own
mm LRU. If it moves in a future vLLM bump, summing .nbytes over the item's tensor
leaves is an acceptable fallback.)

5.3 Structure and concurrency model

class _MaterializedRefCache:
    """Byte-bounded LRU of materialized raw image refs, with single-flight misses.

    All mutation happens on the event loop thread (materialization itself runs in a
    worker thread via asyncio.to_thread, but insertion/eviction/lookup happen in the
    async caller) — so no lock is needed. Do not touch this cache from sync code.
    """
    def __init__(self, max_bytes: int) -> None:
        self.max_bytes = max_bytes
        self._items: OrderedDict[Key, tuple[Any, int]] = OrderedDict()  # key -> (item, nbytes)
        self._inflight: dict[Key, asyncio.Future] = {}
        self._bytes = 0
        self.hits = 0
        self.misses = 0
        self.evictions = 0

get_or_materialize(key, materialize_fn) flow:

1. Hit: key in _items → move_to_end, hits += 1, return the item.
2. In-flight: key in _inflight → await that future (single-flight: N concurrent
   requests for one new image trigger one materialization).
3. Miss: create a future in _inflight; run
   await asyncio.to_thread(materialize_fn);
   • on success: resolve the future, insert (item, nbytes), add to _bytes, then
     evict from the front (popitem(last=False)) until _bytes <= max_bytes
     (evictions += n); if a single item exceeds max_bytes, return it without
     caching (don't churn the whole cache for one giant entry);
   • on failure: set the exception on the future and remove the in-flight entry so
     the next request retries cleanly — never cache a failure. The exception propagates
     to all awaiters (they all get the same _MMImageRefError → HTTP 400, matching
     current behavior).
4. max_bytes == 0 disables the cache entirely: skip all bookkeeping, just call
   materialize_fn via to_thread (keep single-flight off too — zero behavior delta
   from today).

Eviction semantics (the part reviewers should verify): eviction is invisible. Keys
are content-addressed so entries can only go cold, never stale; a post-eviction request
simply misses and re-materializes from the durable file:// asset. Worst case anywhere
= today's behavior for one request. Requests currently holding an evicted item are safe:
eviction only drops the cache's reference; Python refcounting keeps the tensors alive
for in-flight users.

5.4 Aliasing guard (subtle, important)

vLLM's injection path replaces elements of the kwargs lists we pass to mm_input()
(e.g. swapping an item for a shared-memory address when
--mm-processor-cache-type=shm). Therefore:

• Build a fresh list per request in _decode_raw_mm_kwargs; never hand vLLM a list
  object the cache owns (the current code already builds a fresh decoded list — keep
  that).
• Treat the cached MultiModalKwargsItem as immutable after insertion. vLLM
  replaces list slots rather than mutating items in place, so sharing the item object
  across requests is safe; do not add any code that mutates it.

5.5 Concurrency within a request

Replace the sequential loop in _decode_raw_mm_kwargs with asyncio.gather across the
images of a request (across modalities too — flatten, gather, regroup). Combined with
single-flight, the first turn of a group of rollouts materializes each unique image
exactly once, concurrently.

5.6 Configuration knob

One knob: max cache bytes, expressed in GiB.

• Env var: PRIME_RL_MM_MATERIALIZE_CACHE_GB (float, default 2.0, 0 disables).
• Read once at handler init (e.g. a cached_property on PrimeRlServingTokens, or
  module level). Note: PrimeRlServingTokens is grafted via object.__new__ +
  __dict__.update in custom_init_app_state (see the _max_tokens_defaults
cached_property comment at serving_tokens.py:304 for why __init__ never runs) —
  a cached_property is the established pattern; use it.
• Why an env var and not a config field: the vLLM server process is configured through
  the inference config's env_vars passthrough (see [inference.env_vars] /
  DEFAULT_INFERENCE_ENV_VARS in src/prime_rl/entrypoints/rl.py), and the serving
  subclass has no clean path to prime-rl pydantic config. Precedent:
  PRIME_NO_MOE_LORA in src/prime_rl/inference/vllm/server.py:224.
• Document the memory math in the docstring: host RAM on an inference node also carries
  vLLM's processor cache (mm_processor_cache_gb × (api_server_count + data_parallel_size) per vLLM's own config comment in
  vllm/config/multimodal.py:129). Our budget is additive to that.

5.7 Observability

• Counters on the cache: hits, misses, evictions, _bytes.
• Log one INFO summary line periodically — simplest robust rule: every 1000 misses+hits,
  log mm materialize cache: hits=X misses=Y hit_rate=Z% bytes=A/B evictions=C. Use the
  module's existing logger. Do not build a metrics endpoint; logs are how this layer is
  monitored today (see skills/training/monitor-run/SKILL.md, "Multimodal image offload
  checks" section — update that section's bullet list with one line about the new log
  signature as part of this change; AGENTS.md makes skill maintenance part of code
  changes).

---

6. Implementation plan

All in src/prime_rl/inference/vllm/serving_tokens.py unless noted.

1. Add _MaterializedRefCache (§5.3) near the other module helpers.
2. Give PrimeRlServingTokens a cached_property _mm_materialize_cache that reads
   PRIME_RL_MM_MATERIALIZE_CACHE_GB and constructs the cache.
3. Rework _decode_raw_mm_kwargs:
   • it needs the cache (pass it in; the function is currently module-level called from
     serve_tokens, serving_tokens.py:362 — keep it module-level and pass the cache
     as an argument, matching the existing style of explicit parameters),
   • flatten (modality, index, raw_ref, mm_hash, placeholder) across modalities,
   • asyncio.gather over cache.get_or_materialize(key, fn) for each,
   • regroup into dict[modality, list] preserving order.
4. _materialize_raw_image_ref_sync stays exactly as is — it becomes the materialize_fn.
5. Periodic log line (§5.7).
6. Update the monitor-run skill bullet (§5.7).

Do not touch: the adapters, _load_image_processor (already lru_cached,
serving_tokens.py:171), the error-response handling in serve_tokens (the
_MMImageRefError → create_error_response path must behave identically), or anything
outside the inference layer.

7. Testing

Repo conventions (from AGENTS.md): plain pytest functions, no test classes, run with
uv run pytest, be conservative — targeted tests only. Extend
tests/unit/inference/test_serving_tokens.py; the existing test
test_materialize_raw_image_ref_uses_generic_family_payload shows the established
harness (monkeypatch _load_image_processor and get_multimodal_adapter, build a real
ref with renderers.mm_store.raw_mm_ref, write a real tiny image with PIL).

Add, in rough priority order:

1. Hit skips work: two _decode_raw_mm_kwargs calls with the same ref; assert the
   (mocked) adapter's materialize_for_vllm ran once and both calls returned the same
   item object.
2. Byte-budget eviction: cache with a tiny budget; insert two items whose mocked
   sizes exceed it; assert the older key is gone, re-requesting it re-materializes
   (adapter called again), and nothing errors.
3. Failure is not cached: adapter raises on first call (e.g. missing file →
   _MMImageRefError), succeeds on second; assert the first request errors and the
   second succeeds — i.e. the in-flight future was cleaned up.
4. (Optional, if cheap with the harness) single-flight: two concurrent calls for the
   same fresh ref → one adapter invocation.

For sizing in tests, monkeypatch MultiModalCache.get_item_size (or the cache's size
function) so tests don't depend on real tensor layouts.

Run: uv run pytest tests/unit/inference/test_serving_tokens.py -q and the broader
uv run pytest tests/unit -q --deselect tests/unit/train (train tests need GPUs; the
rest pass on CPU — 407 passing as of 1f290dc32).

8. Validation and rollout

• Behavioral invariant: with PRIME_RL_MM_MATERIALIZE_CACHE_GB=0 the request path
  must be byte-identical to today (modulo the gather reordering of thread work). This is
  the safety story — the knob is also the kill switch.
• Live smoke (needs a GPU box): start the inference server for a cached VLM
  (Qwen/Qwen3-VL-4B-Instruct is the one used across the test suites), drive a
  multi-turn multimodal rollout (the mini-browse env or the monitor-run skill's smoke
  notes describe the setup), and confirm: (a) hit-rate log lines climb after turn 1,
  (b) invalid_mm_image_ref count stays zero, (c) tokens/s on later turns improves or
  request latency for image-heavy turns drops.
• Memory check: watch the frontend process RSS against the configured budget.

9. Conventions checklist (from AGENTS.md / CLAUDE.md)

• uv run for everything; never raw python.
• Minimal try/except — the only intentional catch here is the single-flight failure
  propagation (§5.3), which re-raises to the awaiters.
• Targeted comments only; no work-narration comments.
• Branch: work directly on feat/v1-raw-mm-offload, or branch
  feat/mm-materialize-cache off it and open a draft PR (gh pr create --draft)
  based against feat/v1-raw-mm-offload — Eli's call; ask before opening a new PR.
• If you push to Support v1 raw multimodal image offload #2836, update its PR description in the same breath (gh pr edit 2836 --body-file ...), preserving the <!-- CURSOR_SUMMARY --> block.

10. Reference index

What	Where
Target handler + current per-request materialization	src/prime_rl/inference/vllm/serving_tokens.py:218,274,362
Family adapters (materialize_for_vllm, placeholder-length validation)	src/prime_rl/multimodal/adapters/{qwen_vl,kimi_k25}.py
Ref format (mmraw: build/parse)	deps/renderers/renderers/mm_store.py::raw_mm_ref,split_raw_mm_ref
Client that emits refs every request	deps/renderers/renderers/client.py::_build_vllm_mm_features
Ingress offload (content-addressed file store)	deps/verifiers/verifiers/utils/multimodal.py
DP-rank pinning (why one frontend cache serves all ranks)	src/prime_rl/utils/client.py:22-29,246
vLLM item sizing helper	.venv/.../vllm/multimodal/cache.py::MultiModalCache.get_item_size
vLLM external-kwargs cache injection (P0→P1 dedup we compose with)	.venv/.../vllm/v1/engine/input_processor.py ("Inject pre-processed mm_kwargs…")
vLLM mm cache knobs + memory math	.venv/.../vllm/config/multimodal.py:123-137
Existing test harness to extend	tests/unit/inference/test_serving_tokens.py::test_materialize_raw_image_ref_uses_generic_family_payload
Monitoring doc to update	skills/training/monitor-run/SKILL.md ("Multimodal image offload checks")

[cursor]

Missing image zeros packed batch

Medium Severity

placeholder_zero_loss clears the entire microbatch loss_mask and advantages when any raw image file is missing. Raw-ref samples now pack with text and same-family peers, so one missing asset also drops loss for every other sample sharing that pack, including ones whose images are present.

Additional Locations (1)

• src/prime_rl/trainer/batch.py#L355-L381

 

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

[cursor]

Unknown family returns server error

Medium Severity

get_multimodal_adapter raises NotImplementedError for an unknown family, but _materialize_raw_image_ref_sync only maps TypeError and ValueError to invalid_mm_image_ref 400s. An unrecognized family therefore escapes as a 500 instead of a client validation error.

Additional Locations (1)

• src/prime_rl/multimodal/registry.py#L12-L17

 

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

[cursor]

Detached cache task unreferenced

Medium Severity

_MaterializedRefCache fires materialization with create_task and only keeps the shared Future in _inflight, not the Task. Under asyncio’s weak task refs, that task can be garbage-collected before it resolves, leaving awaiters blocked on a future that never completes.

 

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

[cursor]

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

There are 4 total unresolved issues (including 3 from previous reviews).

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

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

[cursor]

Kimi modality key mismatch

High Severity

_decode_raw_mm_kwargs always inserts materialized items under the request feature_modality key, while KimiK25Adapter.materialize_for_vllm returns a vision_chunk-tagged MultiModalKwargsItem. vllm_modality is parsed on RawMMItem but never used to rekey mm_kwargs / hashes / placeholders, so Kimi requests either fail MultiModalKwargsItems construction or never reach the vision path.

Additional Locations (2)

• src/prime_rl/multimodal/adapters/kimi_k25.py#L151-L156
• src/prime_rl/multimodal/schema.py#L17-L18

 

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