From 02d79f7f140cad34c38e00c147ff8991075b3fbd Mon Sep 17 00:00:00 2001
From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
Date: Tue, 22 Sep 2026 12:00:23 -0700
Subject: [PATCH 1/4] fix(pd): process worker-side media once, on the prefill
 leg

With prefill/decode disaggregation and worker-side media processing
both legs fetched and processed the same media: the decode leg kept
media_refs, and the servicer ran the processor on any request that
carried them.

The prefill servicer now returns the processed media's identity
(expanded ids, hashes, placeholders, grid tensors) on its completion,
and the router builds the decode leg from it, references cleared, the
way router-side processing already serves decode. n>1 keeps the
references, since decode recomputes locally; an older servicer's
missing identity falls back to today's behaviour with a warning. A
media_refs request is no longer fanned out into n PD pairs.

Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
---
 .github/workflows/pr-test-rust.yml            |   2 +-
 crates/grpc_client/proto/vllm_engine.proto    |  14 ++
 crates/grpc_client/python/pyproject.toml      |   2 +-
 e2e_test/infra/worker.py                      |  10 ++
 e2e_test/router/test_pd_multimodal.py         |  40 ++++-
 .../smg_grpc_servicer/vllm/media_identity.py  | 121 +++++++++++++
 .../smg_grpc_servicer/vllm/servicer.py        |  17 ++
 .../tests/test_vllm_media_identity.py         | 144 +++++++++++++++
 .../tests/test_vllm_pd_media_identity.py      | 156 +++++++++++++++++
 .../grpc/common/response_formatting.rs        |  10 +-
 .../grpc/common/stages/request_execution.rs   | 164 +++++++++++++++++-
 .../src/routers/grpc/proto_wrapper.rs         |  39 ++++-
 .../src/routers/grpc/regular/streaming.rs     |  12 +-
 13 files changed, 704 insertions(+), 27 deletions(-)
 create mode 100644 grpc_servicer/smg_grpc_servicer/vllm/media_identity.py
 create mode 100644 grpc_servicer/tests/test_vllm_media_identity.py
 create mode 100644 grpc_servicer/tests/test_vllm_pd_media_identity.py

diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml
index cfcad09ed1..7f43408d32 100644
--- a/.github/workflows/pr-test-rust.yml
+++ b/.github/workflows/pr-test-rust.yml
@@ -1014,7 +1014,7 @@ jobs:
       test_timeout: 34
       test_dirs: e2e_test/router/test_pd_multimodal.py
       mm_processing: worker
-      min_selected: 3 # the 3 PD multimodal cases
+      min_selected: 4 # the 4 PD multimodal cases
     secrets: inherit
 
   # === 4 GPU ===
diff --git a/crates/grpc_client/proto/vllm_engine.proto b/crates/grpc_client/proto/vllm_engine.proto
index badbbf6678..c759dd3f45 100644
--- a/crates/grpc_client/proto/vllm_engine.proto
+++ b/crates/grpc_client/proto/vllm_engine.proto
@@ -319,6 +319,20 @@ message GenerateComplete {
 
   uint32 spec_accepted_tokens = 13;
   uint32 spec_draft_tokens = 14;
+
+  // Identity of worker-processed media, returned by the prefill leg so the
+  // decode leg can be served without pixels: the expanded prompt ids plus one
+  // pixel-less MultimodalInputs per modality (hashes, placeholders, grid
+  // tensors). Absent on non-PD requests and from older servicers.
+  MediaIdentity media_identity = 15;
+}
+
+// What the prefill worker learned about the media it processed, in the shape
+// the decode leg already accepts from a router-preprocessed request.
+message MediaIdentity {
+  repeated uint32 prompt_token_ids = 1;
+  MultimodalInputs mm_inputs = 2;
+  repeated MultimodalInputs extra_mm_inputs = 3;
 }
 
 // =====================
diff --git a/crates/grpc_client/python/pyproject.toml b/crates/grpc_client/python/pyproject.toml
index 9df7548807..5ea83acabe 100644
--- a/crates/grpc_client/python/pyproject.toml
+++ b/crates/grpc_client/python/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "smg-grpc-proto"
-version = "0.4.20"
+version = "0.4.21"
 description = "SMG gRPC proto definitions for vLLM, TRT-LLM, MLX, TokenSpeed, and SGLang"
 requires-python = ">=3.10"
 dependencies = [
diff --git a/e2e_test/infra/worker.py b/e2e_test/infra/worker.py
index 9239a3ec33..ccb909d295 100644
--- a/e2e_test/infra/worker.py
+++ b/e2e_test/infra/worker.py
@@ -70,6 +70,8 @@ class Worker:
     # (a deployment-injected SMG_PAIRING_PROTOCOL, for instance).
     extra_env: dict[str, str] | None = None
     process: subprocess.Popen | None = field(default=None, repr=False)
+    # Where the engine's output lands when it is not shown on the terminal.
+    log_path: str | None = field(default=None, repr=False)
     _log_file: IO[Any] | None = field(default=None, repr=False)
     # Used memory per GPU just before launch; ``stop`` waits for it to come back.
     _gpu_mem_baseline: dict[int, int] | None = field(default=None, repr=False)
@@ -628,6 +630,13 @@ def _build_env(self) -> dict[str, str]:
             env.update(self.extra_env)
         return env
 
+    def read_log(self) -> str:
+        """The engine's captured output so far; empty when it goes to the terminal."""
+        if self.log_path is None:
+            return ""
+        with open(self.log_path, encoding="utf-8", errors="replace") as log:
+            return log.read()
+
     def _spawn_process(self, cmd: list[str], env: dict[str, str]) -> subprocess.Popen:
         """Spawn the worker subprocess with output routing."""
         show_output = os.environ.get(ENV_SHOW_WORKER_LOGS, "0") == "1"
@@ -645,6 +654,7 @@ def _spawn_process(self, cmd: list[str], env: dict[str, str]) -> subprocess.Pope
             else:
                 log_path = os.path.join(tempfile.gettempdir(), f"smg-worker-{safe_name}.log")
             self._log_file = open(log_path, "w", encoding="utf-8")
+            self.log_path = log_path
             stdout_target = self._log_file
             stderr_target = subprocess.STDOUT
 
diff --git a/e2e_test/router/test_pd_multimodal.py b/e2e_test/router/test_pd_multimodal.py
index af90d825af..90766be913 100644
--- a/e2e_test/router/test_pd_multimodal.py
+++ b/e2e_test/router/test_pd_multimodal.py
@@ -29,7 +29,9 @@
 
 import httpx
 import pytest
-from infra import assert_mm_processing
+from infra import assert_mm_processing, get_mm_processing
+from infra.constants import MM_PROCESSING_WORKER
+from infra.mm_processing import mm_processing_samples
 from PIL import Image
 
 logger = logging.getLogger(__name__)
@@ -42,6 +44,16 @@
 
 _PROMPT = "Is the dog wrapped in a blanket? Answer with one word: yes or no."
 
+# The servicer's request line for a pixel-less leg built from an identity.
+_DECODE_FROM_IDENTITY = "preprocessed_mm=True, media_refs=0"
+
+
+def _worker_path_count(gateway) -> float:
+    """Requests the gateway has sent down the worker-side media path."""
+    return sum(
+        value for labels, value in mm_processing_samples(gateway) if labels["mode"] == "worker"
+    )
+
 
 def _fixture_image_data_url(name: str) -> str:
     buf = io.BytesIO()
@@ -157,6 +169,32 @@ def test_parallel_sampling_keeps_vision_on_decode(self, setup_backend):
             logger.info("PD multimodal n=2 answer (pug in blanket): %s", answer)
             assert answer == "yes", f"Expected 'yes' from each sample, got: {answer}"
 
+    def test_worker_media_is_processed_once_on_the_prefill_leg(self, setup_backend):
+        # With worker-side processing the prefill leg fetches and processes
+        # the image and answers with its identity; the decode leg is served
+        # from that identity, without references and without pixels.
+        if get_mm_processing() != MM_PROCESSING_WORKER:
+            pytest.skip("router-side lane: the legs receive preprocessed tensors")
+        backend, model, client, gateway = setup_backend
+        prefill, decode = gateway.prefill_workers[0], gateway.decode_workers[0]
+        worker_before = _worker_path_count(gateway)
+        prefill_before = prefill.read_log().count("media_refs=1")
+        decode_refs_before = decode.read_log().count("media_refs=1")
+        decode_identity_before = decode.read_log().count(_DECODE_FROM_IDENTITY)
+
+        assert _ask_blanket(client, model, _PUG_IMAGE) == "yes"
+
+        assert _worker_path_count(gateway) == worker_before + 1
+        assert prefill.read_log().count("media_refs=1") == prefill_before + 1, (
+            "the prefill leg processed the reference"
+        )
+        assert decode.read_log().count("media_refs=1") == decode_refs_before, (
+            "the decode leg received no reference to process"
+        )
+        assert decode.read_log().count(_DECODE_FROM_IDENTITY) == decode_identity_before + 1, (
+            "the decode leg was served from the prefill's media identity"
+        )
+
 
 def _post_blanket(gateway, model: str, image_url: str, **sampling) -> httpx.Response:
     return httpx.post(
diff --git a/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py b/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py
new file mode 100644
index 0000000000..a8a0e797fd
--- /dev/null
+++ b/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py
@@ -0,0 +1,121 @@
+"""The prefill leg's account of the media it processed, for the decode leg.
+
+With prefill/decode disaggregation only the prefill worker runs the media
+processor; its result travels back on the completion in the shape a
+router-preprocessed decode leg already accepts: expanded prompt ids and, per
+modality, hashes, placeholders and the grid tensors, never pixels. Engine-free:
+the processed input is read by shape, not by vLLM's types.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import torch
+from smg_grpc_proto import vllm_engine_pb2
+from smg_grpc_proto.generated import common_pb2
+
+from smg_grpc_servicer.tensor_wire import PROTO_DTYPE_MAP
+
+logger = logging.getLogger(__name__)
+
+# The per-item tensors a decode leg needs besides pixels: the M-RoPE grids and
+# the video timing that goes with them. The router keeps the same set.
+GRID_KEYS = ("image_grid_thw", "video_grid_thw", "second_per_grid_ts")
+
+_MODALITIES = {"image": common_pb2.IMAGE, "video": common_pb2.VIDEO, "audio": common_pb2.AUDIO}
+_WIRE_DTYPES = {dtype: name for name, dtype in PROTO_DTYPE_MAP.items()}
+
+
+def tensor_to_proto(tensor: torch.Tensor) -> vllm_engine_pb2.TensorData:
+    """A grid-sized tensor as inline wire bytes, widened to a dtype the wire
+    names. Bytes go through a Python list: no numpy, and grids are tiny."""
+    data = tensor.detach().cpu()
+    if data.dtype not in _WIRE_DTYPES:
+        data = data.to(torch.float32 if data.is_floating_point() else torch.int64)
+    data = data.contiguous()
+    return vllm_engine_pb2.TensorData(
+        shape=list(data.shape),
+        dtype=_WIRE_DTYPES[data.dtype],
+        inline=bytes(data.flatten().view(torch.uint8).tolist()),
+    )
+
+
+def _first(prompt: Mapping[str, Any], *names: str) -> Any:
+    for name in names:
+        if name in prompt:
+            return prompt[name]
+    return None
+
+
+def _data(elem: Any) -> torch.Tensor:
+    return torch.as_tensor(getattr(elem, "data", elem))
+
+
+def _im_token_id(prompt_token_ids: Sequence[int], ranges: Sequence[Any]) -> int | None:
+    """The token embeddings land on, when a placeholder mixes it with others."""
+    for r in ranges:
+        is_embed = getattr(r, "is_embed", None)
+        if is_embed is None:
+            continue
+        mask = torch.as_tensor(is_embed).flatten().bool()
+        if bool(mask.all()) or not bool(mask.any()):
+            continue
+        first = int(mask.nonzero()[0])
+        return int(prompt_token_ids[r.offset + first])
+    return None
+
+
+def build_media_identity(prompt: Mapping[str, Any]) -> vllm_engine_pb2.MediaIdentity | None:
+    """The identity of a processed engine input, or None when it has no media
+    or a grid tensor cannot be carried (the decode leg then reprocesses)."""
+    hashes = _first(prompt, "mm_hashes", "multi_modal_hashes") or {}
+    placeholders = _first(prompt, "mm_placeholders", "multi_modal_placeholders") or {}
+    kwargs = _first(prompt, "mm_kwargs", "multi_modal_kwargs") or {}
+    prompt_token_ids = list(prompt["prompt_token_ids"])
+    if not hashes:
+        return None
+
+    batches = []
+    for modality, modality_hashes in hashes.items():
+        mm = vllm_engine_pb2.MultimodalInputs(
+            modality=_MODALITIES.get(modality, common_pb2.MODALITY_UNSPECIFIED),
+            mm_hashes=list(modality_hashes),
+        )
+        ranges = list(placeholders.get(modality, []) or [])
+        mm.mm_placeholders.extend(
+            vllm_engine_pb2.PlaceholderRange(offset=int(r.offset), length=int(r.length))
+            for r in ranges
+        )
+        im_token_id = _im_token_id(prompt_token_ids, ranges)
+        if im_token_id is not None:
+            mm.im_token_id = im_token_id
+
+        items = list(kwargs.get(modality, []) or [])
+        for key in GRID_KEYS:
+            elems = [item[key] for item in items if key in item]
+            if not elems:
+                continue
+            tensors = [_data(elem) for elem in elems]
+            if len(tensors) != len(items) or any(t.shape != tensors[0].shape for t in tensors):
+                # Not one row per item: the batched layout the decode leg
+                # reads cannot carry it, so no identity at all.
+                logger.warning(
+                    "media identity: %s grid tensor %s is not one row per item; "
+                    "the decode leg will reprocess the media",
+                    modality,
+                    key,
+                )
+                return None
+            mm.model_specific_tensors[key].CopyFrom(tensor_to_proto(torch.stack(tensors)))
+            mm.batched_keys.append(key)
+            if any(getattr(getattr(elem, "field", None), "keep_on_cpu", False) for elem in elems):
+                mm.keep_on_cpu_keys.append(key)
+        batches.append(mm)
+
+    identity = vllm_engine_pb2.MediaIdentity(prompt_token_ids=prompt_token_ids)
+    identity.mm_inputs.CopyFrom(batches[0])
+    identity.extra_mm_inputs.extend(batches[1:])
+    return identity
diff --git a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py
index 78e1836ed1..ced1cc316c 100755
--- a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py
+++ b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py
@@ -54,6 +54,7 @@
     params_to_response_fields,
     resolve_pd_connector,
 )
+from smg_grpc_servicer.vllm.media_identity import build_media_identity
 from smg_grpc_servicer.vllm.media_refs import parse_media_refs, validate_schemes
 from smg_grpc_servicer.vllm.mm_processor import (
     DEFAULT_MAX_INFLIGHT,
@@ -281,6 +282,8 @@ async def Generate(
         )
 
         kv_transfer_params: dict | None = None
+        # What a PD prefill leg learned about its media, for the decode leg.
+        media_identity = None
         engine_started = False
         try:
             arrival_time = time.time()
@@ -312,6 +315,10 @@ async def Generate(
                     )
                 finally:
                     self._mm_inflight.release()
+                # A PD prefill leg answers with the identity so decode is
+                # served without pixels or references.
+                if kv_transfer_params is not None:
+                    media_identity = build_media_identity(prompt)
             elif has_preprocessed_mm and input_type == "tokenized":
                 # A pixel-less payload (PD decode leg) is only decodable with
                 # remote KV: a local recompute would schedule the vision
@@ -410,6 +417,7 @@ async def Generate(
                                 completion=completion,
                                 num_logprobs=num_logprobs,
                                 num_prompt_logprobs=num_prompt_logprobs,
+                                media_identity=media_identity,
                             )
 
                 # For non-streaming, send complete response when finished
@@ -420,6 +428,7 @@ async def Generate(
                             completion=completion,
                             num_logprobs=num_logprobs,
                             num_prompt_logprobs=num_prompt_logprobs,
+                            media_identity=media_identity,
                         )
 
         except asyncio.CancelledError:
@@ -1184,6 +1193,7 @@ def _complete_response(
         completion: "CompletionOutput | None" = None,
         num_logprobs: int | None = None,
         num_prompt_logprobs: int | None = None,
+        media_identity: "vllm_engine_pb2.MediaIdentity | None" = None,
     ) -> vllm_engine_pb2.GenerateResponse:
         """
         Build a final completion response from vLLM output.
@@ -1197,6 +1207,7 @@ def _complete_response(
                        If None, uses output.outputs[0] for backwards compatibility.
             num_logprobs: Number of top logprobs for output tokens
             num_prompt_logprobs: Number of top logprobs for prompt tokens
+            media_identity: A PD prefill leg's processed media, for the decode leg
 
         Returns:
             GenerateResponse with complete field set
@@ -1239,6 +1250,12 @@ def _complete_response(
 
         # Build matched_stop kwargs from stop_reason (int token ID or str stop sequence)
         stop_kwargs = {}
+        # A proto package predating the field cannot carry the identity.
+        if (
+            media_identity is not None
+            and "media_identity" in vllm_engine_pb2.GenerateComplete.DESCRIPTOR.fields_by_name
+        ):
+            stop_kwargs["media_identity"] = media_identity
         if completion.stop_reason is not None:
             if isinstance(completion.stop_reason, int):
                 stop_kwargs["matched_token_id"] = completion.stop_reason
diff --git a/grpc_servicer/tests/test_vllm_media_identity.py b/grpc_servicer/tests/test_vllm_media_identity.py
new file mode 100644
index 0000000000..86b9208e0e
--- /dev/null
+++ b/grpc_servicer/tests/test_vllm_media_identity.py
@@ -0,0 +1,144 @@
+"""Unit tests for the prefill leg's media identity (engine-free, needs torch).
+
+Run with: pytest grpc_servicer/tests/test_vllm_media_identity.py
+"""
+
+import importlib.util
+import struct
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("smg_grpc_proto")
+from smg_grpc_proto.generated import common_pb2  # noqa: E402
+from smg_grpc_servicer.vllm.mm_tensors import tensor_from_proto  # noqa: E402
+
+# Import the module directly to avoid pulling vllm via the package __init__
+_MODULE_PATH = Path(__file__).parents[1] / "smg_grpc_servicer" / "vllm" / "media_identity.py"
+_spec = importlib.util.spec_from_file_location("media_identity", _MODULE_PATH)
+media_identity = importlib.util.module_from_spec(_spec)
+sys.modules[_spec.name] = media_identity
+_spec.loader.exec_module(media_identity)
+
+
+def _range(offset, length, is_embed=None):
+    return types.SimpleNamespace(offset=offset, length=length, is_embed=is_embed)
+
+
+def _elem(data, keep_on_cpu=False):
+    return types.SimpleNamespace(data=data, field=types.SimpleNamespace(keep_on_cpu=keep_on_cpu))
+
+
+def _image_prompt():
+    """Two images with M-RoPE grids, as a processed engine input."""
+    grids = [torch.tensor([1, 4, 6]), torch.tensor([1, 2, 2])]
+    return {
+        "prompt_token_ids": [7, 8, 100, 100, 100, 9, 100, 5],
+        "mm_hashes": {"image": ["h1", "h2"]},
+        "mm_placeholders": {
+            "image": [
+                # <start> token, three patches, <end> token: embeds land on 100.
+                _range(1, 5, torch.tensor([False, True, True, True, False])),
+                _range(6, 1),
+            ]
+        },
+        "mm_kwargs": {
+            "image": [
+                {
+                    "pixel_values": _elem(torch.zeros(3, 8, 8)),
+                    "image_grid_thw": _elem(grids[0], True),
+                },
+                {
+                    "pixel_values": _elem(torch.zeros(3, 8, 8)),
+                    "image_grid_thw": _elem(grids[1], True),
+                },
+            ]
+        },
+    }
+
+
+class TestBuildMediaIdentity:
+    def test_carries_ids_hashes_placeholders_and_grids_but_no_pixels(self):
+        identity = media_identity.build_media_identity(_image_prompt())
+        assert list(identity.prompt_token_ids) == [7, 8, 100, 100, 100, 9, 100, 5]
+        mm = identity.mm_inputs
+        assert mm.modality == common_pb2.IMAGE
+        assert list(mm.mm_hashes) == ["h1", "h2"]
+        assert [(p.offset, p.length) for p in mm.mm_placeholders] == [(1, 5), (6, 1)]
+        assert mm.im_token_id == 100, "read off the first embed position of the masked range"
+        assert not mm.HasField("pixel_values")
+        assert set(mm.model_specific_tensors) == {"image_grid_thw"}
+        assert list(mm.batched_keys) == ["image_grid_thw"]
+        assert list(mm.keep_on_cpu_keys) == ["image_grid_thw"]
+        grid = tensor_from_proto(mm.model_specific_tensors["image_grid_thw"])
+        assert grid.dtype == torch.int64
+        assert grid.tolist() == [[1, 4, 6], [1, 2, 2]]
+        assert list(identity.extra_mm_inputs) == []
+
+    def test_a_second_modality_goes_to_extra_mm_inputs(self):
+        prompt = _image_prompt()
+        prompt["mm_hashes"]["video"] = ["v1"]
+        prompt["mm_placeholders"]["video"] = [_range(7, 1)]
+        prompt["mm_kwargs"]["video"] = [
+            {
+                "video_grid_thw": _elem(torch.tensor([2, 2, 2])),
+                "second_per_grid_ts": _elem(torch.tensor(0.5)),
+            }
+        ]
+        identity = media_identity.build_media_identity(prompt)
+        assert identity.mm_inputs.modality == common_pb2.IMAGE
+        (video,) = identity.extra_mm_inputs
+        assert video.modality == common_pb2.VIDEO
+        assert list(video.mm_hashes) == ["v1"]
+        assert set(video.model_specific_tensors) == {"video_grid_thw", "second_per_grid_ts"}
+        seconds = tensor_from_proto(video.model_specific_tensors["second_per_grid_ts"])
+        assert seconds.dtype == torch.float32 and seconds.tolist() == [0.5]
+        assert not video.HasField("im_token_id")
+
+    def test_no_media_means_no_identity(self):
+        assert media_identity.build_media_identity({"prompt_token_ids": [1, 2]}) is None
+        assert (
+            media_identity.build_media_identity({"prompt_token_ids": [1], "mm_hashes": {}}) is None
+        )
+
+    def test_a_grid_that_is_not_one_row_per_item_yields_no_identity(self, caplog):
+        prompt = _image_prompt()
+        prompt["mm_kwargs"]["image"][1]["image_grid_thw"] = _elem(torch.tensor([[1, 2, 2]]))
+        with caplog.at_level("WARNING", logger="media_identity"):
+            assert media_identity.build_media_identity(prompt) is None
+        assert any("image_grid_thw" in r.getMessage() for r in caplog.records)
+
+    def test_older_field_names_are_read_too(self):
+        prompt = _image_prompt()
+        prompt["multi_modal_hashes"] = prompt.pop("mm_hashes")
+        prompt["multi_modal_placeholders"] = prompt.pop("mm_placeholders")
+        prompt["multi_modal_kwargs"] = prompt.pop("mm_kwargs")
+        assert media_identity.build_media_identity(prompt).mm_inputs.mm_hashes == ["h1", "h2"]
+
+
+class TestTensorToProto:
+    def test_round_trips_and_widens_unnamed_dtypes(self):
+        for tensor in (
+            torch.tensor([[1, 2, 3]], dtype=torch.int32),
+            torch.tensor([1.5, -2.0], dtype=torch.float64),
+            torch.tensor([[1, 2], [3, 4]], dtype=torch.int64),
+        ):
+            proto = media_identity.tensor_to_proto(tensor)
+            back = tensor_from_proto(proto)
+            assert back.tolist() == tensor.tolist()
+            assert list(proto.shape) == list(tensor.shape)
+        assert media_identity.tensor_to_proto(torch.tensor([1], dtype=torch.int32)).dtype == "int64"
+        assert (
+            media_identity.tensor_to_proto(torch.tensor([1.0], dtype=torch.float64)).dtype
+            == "float32"
+        )
+
+    def test_matches_the_routers_wire_form(self):
+        # The router writes little-endian element bytes; a decode leg reads
+        # either sender the same way.
+        proto = media_identity.tensor_to_proto(torch.tensor([1, 4, 6], dtype=torch.int64))
+        assert proto.inline == struct.pack("<3q", 1, 4, 6)
+        assert proto.dtype == "int64" and list(proto.shape) == [3]
diff --git a/grpc_servicer/tests/test_vllm_pd_media_identity.py b/grpc_servicer/tests/test_vllm_pd_media_identity.py
new file mode 100644
index 0000000000..6dc8019786
--- /dev/null
+++ b/grpc_servicer/tests/test_vllm_pd_media_identity.py
@@ -0,0 +1,156 @@
+"""PD legs and worker-side media (needs vLLM: runs on the vLLM e2e lane).
+
+The prefill leg processes the references once and answers with their
+identity; a regular request answers with none; the decode leg, given that
+identity, never touches the media processor.
+"""
+
+import json
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+pytest.importorskip("vllm")
+torch = pytest.importorskip("torch")
+from smg_grpc_proto import vllm_engine_pb2  # noqa: E402
+from smg_grpc_proto.generated import common_pb2  # noqa: E402
+from smg_grpc_servicer.vllm.servicer import VllmEngineServicer  # noqa: E402
+from vllm.outputs import CompletionOutput, RequestOutput  # noqa: E402
+
+PREFILL_KV = json.dumps({"do_remote_decode": True, "do_remote_prefill": False})
+
+
+class _Processor:
+    """Stands in for the worker-side media processor; records its calls."""
+
+    name = "inprocess"
+    max_inflight = 4
+    schemes = "http,https,data"
+    accepted_schemes = {"http", "https", "data"}
+
+    def __init__(self):
+        self.calls = []
+
+    async def probe(self):
+        return True
+
+    async def process(self, prompt_token_ids, prompt_text, items, arrival_time, *, request_id=""):
+        self.calls.append(request_id)
+        grid = SimpleNamespace(
+            data=torch.tensor([1, 4, 6]), field=SimpleNamespace(keep_on_cpu=True)
+        )
+        return {
+            "prompt_token_ids": [7, 100, 100, 100, 9],
+            "mm_hashes": {"image": ["h1"]},
+            "mm_placeholders": {"image": [SimpleNamespace(offset=1, length=3, is_embed=None)]},
+            "mm_kwargs": {
+                "image": [
+                    {
+                        "pixel_values": SimpleNamespace(data=torch.zeros(3, 2, 2)),
+                        "image_grid_thw": grid,
+                    }
+                ]
+            },
+            "arrival_time": arrival_time,
+        }
+
+
+def _engine(seen_prompts):
+    async def generate(*, prompt, sampling_params, request_id, **kwargs):
+        seen_prompts.append(prompt)
+        yield RequestOutput(
+            request_id=request_id,
+            prompt=None,
+            prompt_token_ids=[7, 100, 100, 100, 9],
+            prompt_logprobs=None,
+            outputs=[
+                CompletionOutput(
+                    index=0,
+                    text="",
+                    token_ids=[42],
+                    cumulative_logprob=None,
+                    logprobs=None,
+                    finish_reason="length",
+                )
+            ],
+            finished=True,
+            kv_transfer_params={"remote_block_ids": [1]},
+        )
+
+    return SimpleNamespace(
+        generate=generate,
+        renderer=SimpleNamespace(process_for_engine=lambda prompt, **kwargs: prompt),
+        model_config=SimpleNamespace(
+            is_multimodal_model=True, dtype=torch.float32, uses_mrope=True
+        ),
+        vllm_config=SimpleNamespace(kv_events_config=None),
+    )
+
+
+def _servicer(monkeypatch, processor):
+    monkeypatch.delenv("SMG_VLLM_MM_PROCESSOR", raising=False)
+    servicer = VllmEngineServicer(_engine([]), start_time=0.0)
+    servicer._mm_processor = processor
+    return servicer
+
+
+def _request(*, refs=False, kv=False, identity=None):
+    request = vllm_engine_pb2.GenerateRequest(
+        request_id="req-pd",
+        tokenized=vllm_engine_pb2.TokenizedInput(input_ids=[7, 8, 9]),
+        sampling_params=vllm_engine_pb2.SamplingParams(max_tokens=1),
+    )
+    if refs:
+        request.media_refs.items.append(
+            vllm_engine_pb2.MediaRef(modality=common_pb2.IMAGE, url="https://a/1.png")
+        )
+    if kv:
+        request.kv_transfer_params_json = PREFILL_KV
+    if identity is not None:
+        request.tokenized.input_ids[:] = identity.prompt_token_ids
+        request.mm_inputs.CopyFrom(identity.mm_inputs)
+    return request
+
+
+async def _complete(servicer, request):
+    context = SimpleNamespace(abort=AsyncMock())
+    responses = [r async for r in servicer.Generate(request, context)]
+    context.abort.assert_not_awaited()
+    (final,) = [r for r in responses if r.HasField("complete")]
+    return final.complete
+
+
+@pytest.mark.asyncio
+async def test_prefill_leg_answers_with_the_media_identity(monkeypatch):
+    processor = _Processor()
+    servicer = _servicer(monkeypatch, processor)
+    complete = await _complete(servicer, _request(refs=True, kv=True))
+    assert processor.calls == ["req-pd"], "the media was processed once"
+    assert complete.HasField("media_identity")
+    identity = complete.media_identity
+    assert list(identity.prompt_token_ids) == [7, 100, 100, 100, 9]
+    assert list(identity.mm_inputs.mm_hashes) == ["h1"]
+    assert [(p.offset, p.length) for p in identity.mm_inputs.mm_placeholders] == [(1, 3)]
+    assert "image_grid_thw" in identity.mm_inputs.model_specific_tensors
+    assert not identity.mm_inputs.HasField("pixel_values")
+
+
+@pytest.mark.asyncio
+async def test_regular_leg_answers_with_no_identity(monkeypatch):
+    processor = _Processor()
+    servicer = _servicer(monkeypatch, processor)
+    complete = await _complete(servicer, _request(refs=True, kv=False))
+    assert processor.calls == ["req-pd"]
+    assert not complete.HasField("media_identity")
+
+
+@pytest.mark.asyncio
+async def test_decode_leg_from_identity_never_touches_the_processor(monkeypatch):
+    processor = _Processor()
+    servicer = _servicer(monkeypatch, processor)
+    prefill = await _complete(servicer, _request(refs=True, kv=True))
+    processor.calls.clear()
+    decode = await _complete(servicer, _request(kv=True, identity=prefill.media_identity))
+    assert processor.calls == [], "the decode leg carried identity, not references"
+    assert not decode.HasField("media_identity")
diff --git a/model_gateway/src/routers/grpc/common/response_formatting.rs b/model_gateway/src/routers/grpc/common/response_formatting.rs
index 25648e7556..174f7f0073 100644
--- a/model_gateway/src/routers/grpc/common/response_formatting.rs
+++ b/model_gateway/src/routers/grpc/common/response_formatting.rs
@@ -151,10 +151,12 @@ mod tests {
             token_ids: vec![1, 2, 3],
             ..Default::default()
         }));
-        tracker.record_complete(&ProtoGenerateComplete::Vllm(vllm::GenerateComplete {
-            completion_tokens: 99,
-            ..Default::default()
-        }));
+        tracker.record_complete(&ProtoGenerateComplete::Vllm(Box::new(
+            vllm::GenerateComplete {
+                completion_tokens: 99,
+                ..Default::default()
+            },
+        )));
         assert_eq!(
             tracker.total(),
             3,
diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs
index 47efc6bc4d..67c451867b 100644
--- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs
+++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs
@@ -4,6 +4,7 @@ use std::{future::Future, pin::Pin, sync::Arc, time::Instant};
 
 use axum::response::Response;
 use futures::future::{join_all, try_join_all};
+use smg_grpc_client::vllm_proto as vllm;
 use tracing::{debug, error, info_span, warn, Instrument};
 
 use super::{
@@ -144,13 +145,38 @@ fn plan_sub_requests(plan: &ExecutionPlan, workers: Option<&WorkerSelection>) ->
 /// read, so it cannot be handed to n prefills; those requests keep the
 /// single dispatch.
 fn pd_fanout_width(request: &ProtoGenerateRequest, protocol: PdProtocol) -> Option<u32> {
-    if protocol.dispatch != PdDispatch::Parallel || request.has_mm_inputs() {
+    if protocol.dispatch != PdDispatch::Parallel
+        || request.has_mm_inputs()
+        || request.has_vllm_media_refs()
+    {
         return None;
     }
     let n = request.sampling_n();
     (n > 1).then_some(n)
 }
 
+/// Give the decode leg the media identity the prefill leg produced, so it
+/// is served without pixels or references. Only with a KV handoff: without
+/// one decode recomputes the prompt and needs the references itself. A
+/// prefill that returned no identity (an older servicer) leaves the leg as
+/// it is, and decode reprocesses the media.
+fn apply_prefill_media_identity(
+    decode_request: &mut ProtoGenerateRequest,
+    relay_kv_params: bool,
+    identity: Option<&vllm::MediaIdentity>,
+) {
+    if !relay_kv_params || !decode_request.has_vllm_media_refs() {
+        return;
+    }
+    match identity {
+        Some(identity) => decode_request.apply_media_identity(identity),
+        None => warn!(
+            request_id = %decode_request.request_id(),
+            "prefill worker returned no media identity; decode leg will reprocess media"
+        ),
+    }
+}
+
 /// Split an n>1 request into `n` single-sample sub-requests. Sub `i` carries
 /// engine id `{id}-{i}`, `n = 1`, a seed offset of `i` when the request
 /// pinned a seed, and whatever rendezvous `remint` stamps on it.
@@ -997,6 +1023,8 @@ async fn execute_sequential_pd(
     // 0.13. The pixel-free leg is the clone, so pixel tensors are never
     // duplicated and die with the prefill send; the per-image mm identity
     // and grid tensors survive for decode-side hashing and positions.
+    // Media references are resolved by the prefill leg and relayed to
+    // decode as that same identity once prefill completes.
     // Without a KV handoff (n>1) decode recomputes the prompt locally and
     // must run the vision encoder, so that leg keeps the full multimodal
     // payload (SHM-backed tensors cannot serve both legs and fail loudly
@@ -1074,8 +1102,10 @@ async fn execute_sequential_pd(
             )
         })?;
 
-    // Drain prefill response, harvesting connector params from the Complete frame
+    // Drain prefill response, harvesting connector params and the processed
+    // media identity from the Complete frame
     let mut prefill_kv_params: Option<String> = None;
+    let mut prefill_media_identity: Option<vllm::MediaIdentity> = None;
     while let Some(result) = prefill_stream.next().await {
         match result {
             Ok(response) => {
@@ -1083,6 +1113,9 @@ async fn execute_sequential_pd(
                     if let Some(json) = complete.kv_transfer_params_json() {
                         prefill_kv_params = Some(json.to_owned());
                     }
+                    if let Some(identity) = complete.media_identity() {
+                        prefill_media_identity = Some(identity.clone());
+                    }
                 }
             }
             Err(e) => {
@@ -1110,6 +1143,11 @@ async fn execute_sequential_pd(
 
     debug!("vLLM PD: prefill completed, sending decode request");
 
+    apply_prefill_media_identity(
+        &mut decode_request,
+        relay_kv_params,
+        prefill_media_identity.as_ref(),
+    );
     if let Some(rank) = workers.decode_worker().and_then(|w| w.dp_rank()) {
         decode_request.set_data_parallel_rank(rank as i32);
     }
@@ -1613,7 +1651,8 @@ mod tests {
 
     #[test]
     fn clone_without_mm_pixels_keeps_vllm_media_refs() {
-        // Both PD legs process the references themselves.
+        // The clone keeps the references; on the sequential path they are
+        // replaced by the prefill leg's identity before decode is sent.
         let refs = vllm::MediaRefs {
             items: vec![vllm::MediaRef {
                 modality: smg_grpc_client::common_proto::Modality::Image as i32,
@@ -1636,6 +1675,115 @@ mod tests {
         assert!(decode.mm_inputs.is_none());
     }
 
+    fn media_refs_request(id: &str, n: u32) -> ProtoGenerateRequest {
+        let mut request = ProtoGenerateRequest::Vllm(Box::new(vllm::GenerateRequest {
+            request_id: id.to_string(),
+            input: Some(vllm::generate_request::Input::Tokenized(
+                vllm::TokenizedInput {
+                    original_text: "describe <|image|>".to_string(),
+                    input_ids: vec![7, 8, 9],
+                },
+            )),
+            sampling_params: Some(vllm::SamplingParams {
+                n,
+                ..Default::default()
+            }),
+            ..Default::default()
+        }));
+        request
+            .set_vllm_media_refs(vllm::MediaRefs {
+                items: vec![vllm::MediaRef {
+                    modality: smg_grpc_client::common_proto::Modality::Image as i32,
+                    url: "https://a/1.png".to_string(),
+                }],
+            })
+            .expect("vLLM request accepts media refs");
+        request
+    }
+
+    fn identity() -> vllm::MediaIdentity {
+        vllm::MediaIdentity {
+            prompt_token_ids: vec![7, 8, 100, 100, 100, 9],
+            mm_inputs: Some(vllm::MultimodalInputs {
+                mm_hashes: vec!["h1".to_string(), "h2".to_string()],
+                mm_placeholders: vec![
+                    vllm::PlaceholderRange {
+                        offset: 2,
+                        length: 2,
+                    },
+                    vllm::PlaceholderRange {
+                        offset: 4,
+                        length: 1,
+                    },
+                ],
+                model_specific_tensors: std::collections::HashMap::from([(
+                    "image_grid_thw".to_string(),
+                    vllm::TensorData {
+                        shape: vec![2, 3],
+                        dtype: "int64".to_string(),
+                        payload: Some(vllm::tensor_data::Payload::Inline(vec![0; 48])),
+                    },
+                )]),
+                batched_keys: vec!["image_grid_thw".to_string()],
+                modality: smg_grpc_client::common_proto::Modality::Image as i32,
+                ..Default::default()
+            }),
+            extra_mm_inputs: vec![],
+        }
+    }
+
+    /// A media_refs request under parallel PD dispatch would otherwise become
+    /// n pairs, each processing the media on both legs.
+    #[test]
+    fn media_refs_request_never_fans_out() {
+        let sglang = PdProtocol::for_runtime(RuntimeType::Sglang).unwrap();
+        assert_eq!(sglang.dispatch, PdDispatch::Parallel);
+        assert_eq!(pd_fanout_width(&media_refs_request("fan", 3), sglang), None);
+    }
+
+    /// Without a KV handoff decode recomputes the prompt locally, so it needs
+    /// the media itself and the identity must not replace its references.
+    #[test]
+    fn n_greater_than_one_keeps_media_refs_on_decode() {
+        let mut decode = media_refs_request("n2", 2).clone_without_mm_pixels();
+        apply_prefill_media_identity(&mut decode, false, Some(&identity()));
+        assert!(decode.has_vllm_media_refs());
+        let ProtoGenerateRequest::Vllm(request) = &decode else {
+            panic!("expected vLLM request");
+        };
+        assert!(request.mm_inputs.is_none());
+    }
+
+    #[test]
+    fn decode_leg_from_media_identity_has_no_refs_and_no_pixels() {
+        let mut decode = media_refs_request("relay", 1).clone_without_mm_pixels();
+        apply_prefill_media_identity(&mut decode, true, Some(&identity()));
+        assert!(!decode.has_vllm_media_refs());
+        let ProtoGenerateRequest::Vllm(request) = &decode else {
+            panic!("expected vLLM request");
+        };
+        assert!(request.media_refs.is_none());
+        let mm = request.mm_inputs.as_ref().expect("identity mm_inputs");
+        assert!(mm.pixel_values.is_none());
+        assert_eq!(mm.mm_hashes.len(), 2);
+        assert_eq!(mm.mm_placeholders.len(), 2);
+        assert!(mm.model_specific_tensors.contains_key("image_grid_thw"));
+        let Some(vllm::generate_request::Input::Tokenized(tokenized)) = &request.input else {
+            panic!("expected tokenized input");
+        };
+        assert_eq!(tokenized.input_ids, vec![7, 8, 100, 100, 100, 9]);
+        assert_eq!(tokenized.original_text, "describe <|image|>");
+    }
+
+    /// An older servicer returns no identity: the decode leg keeps its
+    /// references and reprocesses, as before.
+    #[test]
+    fn missing_identity_leaves_the_decode_leg_untouched() {
+        let mut decode = media_refs_request("old", 1).clone_without_mm_pixels();
+        apply_prefill_media_identity(&mut decode, true, None);
+        assert!(decode.has_vllm_media_refs());
+    }
+
     #[test]
     fn clone_without_mm_pixels_keeps_vllm_identity_and_grid_tensors() {
         // The decode leg drops the pixel tensors but keeps the per-image
@@ -2070,22 +2218,22 @@ mod tests {
     fn kv_transfer_params_json_complete_accessor_filters_empty() {
         use crate::routers::grpc::proto_wrapper::ProtoGenerateComplete;
 
-        let complete = ProtoGenerateComplete::Vllm(vllm::GenerateComplete {
+        let complete = ProtoGenerateComplete::Vllm(Box::new(vllm::GenerateComplete {
             kv_transfer_params_json: Some(r#"{"do_remote_prefill":true}"#.to_string()),
             ..Default::default()
-        });
+        }));
         assert_eq!(
             complete.kv_transfer_params_json(),
             Some(r#"{"do_remote_prefill":true}"#)
         );
 
-        let empty = ProtoGenerateComplete::Vllm(vllm::GenerateComplete {
+        let empty = ProtoGenerateComplete::Vllm(Box::new(vllm::GenerateComplete {
             kv_transfer_params_json: Some(String::new()),
             ..Default::default()
-        });
+        }));
         assert_eq!(empty.kv_transfer_params_json(), None);
 
-        let unset = ProtoGenerateComplete::Vllm(vllm::GenerateComplete::default());
+        let unset = ProtoGenerateComplete::Vllm(Box::default());
         assert_eq!(unset.kv_transfer_params_json(), None);
     }
 }
diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs
index 2d71cbb57f..3f6ad1e43d 100644
--- a/model_gateway/src/routers/grpc/proto_wrapper.rs
+++ b/model_gateway/src/routers/grpc/proto_wrapper.rs
@@ -1657,6 +1657,21 @@ impl ProtoGenerateRequest {
         matches!(self, Self::Vllm(req) if req.media_refs.is_some())
     }
 
+    /// Replace this vLLM request's media references with the identity the
+    /// prefill leg produced: the expanded prompt ids and the pixel-less
+    /// per-modality inputs, in the shape a router-preprocessed decode leg has.
+    pub fn apply_media_identity(&mut self, identity: &vllm::MediaIdentity) {
+        let Self::Vllm(req) = self else {
+            return;
+        };
+        if let Some(vllm::generate_request::Input::Tokenized(tokenized)) = req.input.as_mut() {
+            tokenized.input_ids.clone_from(&identity.prompt_token_ids);
+        }
+        req.mm_inputs.clone_from(&identity.mm_inputs);
+        req.extra_mm_inputs.clone_from(&identity.extra_mm_inputs);
+        req.media_refs = None;
+    }
+
     /// Number of parallel samples requested (1 when unset). vLLM, SGLang
     /// and TokenSpeed carry it in their sampling params; the others have no
     /// per-request sample count.
@@ -1907,7 +1922,7 @@ impl ProtoGenerateResponse {
                     ProtoResponseVariant::Chunk(ProtoGenerateStreamChunk::Vllm(chunk))
                 }
                 Some(vllm::generate_response::Response::Complete(complete)) => {
-                    ProtoResponseVariant::Complete(ProtoGenerateComplete::Vllm(complete))
+                    ProtoResponseVariant::Complete(ProtoGenerateComplete::Vllm(Box::new(complete)))
                 }
                 None => ProtoResponseVariant::None,
             },
@@ -2135,7 +2150,10 @@ impl ProtoGenerateStreamChunk {
 #[derive(Clone)]
 pub enum ProtoGenerateComplete {
     Sglang(sglang::GenerateComplete),
-    Vllm(vllm::GenerateComplete),
+    // Boxed: the media identity a PD prefill leg returns made this variant
+    // several times the size of the others, and completions move through
+    // the same channel as every streamed chunk.
+    Vllm(Box<vllm::GenerateComplete>),
     Trtllm(trtllm::GenerateComplete),
     Mlx(mlx::GenerateComplete),
     TokenSpeed(tokenspeed::GenerateComplete),
@@ -2447,6 +2465,15 @@ impl ProtoGenerateComplete {
         }
     }
 
+    /// The media identity a prefill leg returned for its processed
+    /// references (vLLM only; absent from older servicers).
+    pub fn media_identity(&self) -> Option<&vllm::MediaIdentity> {
+        match self {
+            Self::Vllm(c) => c.media_identity.as_ref(),
+            Self::Sglang(_) | Self::Trtllm(_) | Self::Mlx(_) | Self::TokenSpeed(_) => None,
+        }
+    }
+
     /// Get opaque connector KV-transfer params JSON returned by the engine (vLLM only).
     pub fn kv_transfer_params_json(&self) -> Option<&str> {
         match self {
@@ -3409,11 +3436,9 @@ mod tests {
                 .chunk_semantics()
                 .is_delta()
         );
-        assert!(
-            ProtoGenerateComplete::Vllm(vllm::GenerateComplete::default())
-                .chunk_semantics()
-                .is_delta()
-        );
+        assert!(ProtoGenerateComplete::Vllm(Box::default())
+            .chunk_semantics()
+            .is_delta());
 
         // Every other shape reports running totals.
         for chunk in [
diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs
index 9915e2fead..bb83760ac1 100644
--- a/model_gateway/src/routers/grpc/regular/streaming.rs
+++ b/model_gateway/src/routers/grpc/regular/streaming.rs
@@ -3492,11 +3492,13 @@ mod tests {
             ));
         }
         assert_eq!(tracker.snapshot().completion_tokens, 3);
-        tracker.record_complete(&ProtoGenerateComplete::Vllm(proto::GenerateComplete {
-            prompt_tokens: 10,
-            completion_tokens: 3,
-            ..Default::default()
-        }));
+        tracker.record_complete(&ProtoGenerateComplete::Vllm(Box::new(
+            proto::GenerateComplete {
+                prompt_tokens: 10,
+                completion_tokens: 3,
+                ..Default::default()
+            },
+        )));
         assert_eq!(tracker.snapshot().completion_tokens, 3);
     }
 

From 6a44b65d8b1c8b67e11ac3337d56a667096c3758 Mon Sep 17 00:00:00 2001
From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
Date: Tue, 22 Sep 2026 12:54:26 -0700
Subject: [PATCH 2/4] fix(pd): apply the media identity only once decode holds
 a KV handoff

A NIXL prefill that returned an identity but no KV params had its
decode leg rebuilt pixel-less, which the servicer refuses without KV;
the references are now replaced only after the handoff is injected or
relayed, the missing-identity warning fires only when one was asked
for, a servicer with an older proto package returns no identity, and
the e2e case skips its log assertions when worker output is shown.

Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
---
 e2e_test/router/test_pd_multimodal.py         |  2 +
 .../smg_grpc_servicer/vllm/media_identity.py  |  5 ++
 .../smg_grpc_servicer/vllm/servicer.py        | 11 ++-
 .../tests/test_vllm_media_identity.py         |  6 ++
 .../grpc/common/stages/request_execution.rs   | 67 +++++++++++++------
 5 files changed, 70 insertions(+), 21 deletions(-)

diff --git a/e2e_test/router/test_pd_multimodal.py b/e2e_test/router/test_pd_multimodal.py
index 90766be913..1bd000501c 100644
--- a/e2e_test/router/test_pd_multimodal.py
+++ b/e2e_test/router/test_pd_multimodal.py
@@ -177,6 +177,8 @@ def test_worker_media_is_processed_once_on_the_prefill_leg(self, setup_backend):
             pytest.skip("router-side lane: the legs receive preprocessed tensors")
         backend, model, client, gateway = setup_backend
         prefill, decode = gateway.prefill_workers[0], gateway.decode_workers[0]
+        if prefill.log_path is None or decode.log_path is None:
+            pytest.skip("worker output goes to the terminal; the log assertions need captured logs")
         worker_before = _worker_path_count(gateway)
         prefill_before = prefill.read_log().count("media_refs=1")
         decode_refs_before = decode.read_log().count("media_refs=1")
diff --git a/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py b/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py
index a8a0e797fd..3e0226e489 100644
--- a/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py
+++ b/grpc_servicer/smg_grpc_servicer/vllm/media_identity.py
@@ -29,6 +29,11 @@
 _WIRE_DTYPES = {dtype: name for name, dtype in PROTO_DTYPE_MAP.items()}
 
 
+def media_identity_supported() -> bool:
+    """Whether the installed proto package can carry an identity at all."""
+    return "media_identity" in vllm_engine_pb2.GenerateComplete.DESCRIPTOR.fields_by_name
+
+
 def tensor_to_proto(tensor: torch.Tensor) -> vllm_engine_pb2.TensorData:
     """A grid-sized tensor as inline wire bytes, widened to a dtype the wire
     names. Bytes go through a Python list: no numpy, and grids are tiny."""
diff --git a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py
index ced1cc316c..93fd367087 100755
--- a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py
+++ b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py
@@ -54,7 +54,7 @@
     params_to_response_fields,
     resolve_pd_connector,
 )
-from smg_grpc_servicer.vllm.media_identity import build_media_identity
+from smg_grpc_servicer.vllm.media_identity import build_media_identity, media_identity_supported
 from smg_grpc_servicer.vllm.media_refs import parse_media_refs, validate_schemes
 from smg_grpc_servicer.vllm.mm_processor import (
     DEFAULT_MAX_INFLIGHT,
@@ -318,7 +318,14 @@ async def Generate(
                 # A PD prefill leg answers with the identity so decode is
                 # served without pixels or references.
                 if kv_transfer_params is not None:
-                    media_identity = build_media_identity(prompt)
+                    if media_identity_supported():
+                        media_identity = build_media_identity(prompt)
+                    else:
+                        logger.warning(
+                            "Request %s: the installed smg-grpc-proto has no media_identity; "
+                            "the decode leg will reprocess the media",
+                            request_id,
+                        )
             elif has_preprocessed_mm and input_type == "tokenized":
                 # A pixel-less payload (PD decode leg) is only decodable with
                 # remote KV: a local recompute would schedule the vision
diff --git a/grpc_servicer/tests/test_vllm_media_identity.py b/grpc_servicer/tests/test_vllm_media_identity.py
index 86b9208e0e..7a0b553ffa 100644
--- a/grpc_servicer/tests/test_vllm_media_identity.py
+++ b/grpc_servicer/tests/test_vllm_media_identity.py
@@ -119,6 +119,12 @@ def test_older_field_names_are_read_too(self):
         assert media_identity.build_media_identity(prompt).mm_inputs.mm_hashes == ["h1", "h2"]
 
 
+class TestSupport:
+    def test_the_installed_proto_carries_the_identity(self):
+        # The proto package built from this tree (as CI does) has the field.
+        assert media_identity.media_identity_supported()
+
+
 class TestTensorToProto:
     def test_round_trips_and_widens_unnamed_dtypes(self):
         for tensor in (
diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs
index 67c451867b..025a113853 100644
--- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs
+++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs
@@ -156,24 +156,28 @@ fn pd_fanout_width(request: &ProtoGenerateRequest, protocol: PdProtocol) -> Opti
 }
 
 /// Give the decode leg the media identity the prefill leg produced, so it
-/// is served without pixels or references. Only with a KV handoff: without
-/// one decode recomputes the prompt and needs the references itself. A
-/// prefill that returned no identity (an older servicer) leaves the leg as
-/// it is, and decode reprocesses the media.
+/// is served without pixels or references. Only once decode actually holds
+/// a KV handoff (`handed_off`): without one it recomputes the prompt and
+/// needs the references itself, and a pixel-less leg without KV is refused
+/// by the servicer. A prefill that was asked for an identity (`solicited`:
+/// its request carried KV params) but returned none is an older servicer;
+/// the leg is left as it is, and decode reprocesses the media.
 fn apply_prefill_media_identity(
     decode_request: &mut ProtoGenerateRequest,
-    relay_kv_params: bool,
+    handed_off: bool,
+    solicited: bool,
     identity: Option<&vllm::MediaIdentity>,
 ) {
-    if !relay_kv_params || !decode_request.has_vllm_media_refs() {
+    if !handed_off || !decode_request.has_vllm_media_refs() {
         return;
     }
     match identity {
         Some(identity) => decode_request.apply_media_identity(identity),
-        None => warn!(
+        None if solicited => warn!(
             request_id = %decode_request.request_id(),
             "prefill worker returned no media identity; decode leg will reprocess media"
         ),
+        None => {}
     }
 }
 
@@ -1143,14 +1147,14 @@ async fn execute_sequential_pd(
 
     debug!("vLLM PD: prefill completed, sending decode request");
 
-    apply_prefill_media_identity(
-        &mut decode_request,
-        relay_kv_params,
-        prefill_media_identity.as_ref(),
-    );
     if let Some(rank) = workers.decode_worker().and_then(|w| w.dp_rank()) {
         decode_request.set_data_parallel_rank(rank as i32);
     }
+    // The prefill request carried KV params, so the servicer was asked for
+    // the media identity; and whether decode ends up holding a handoff.
+    let identity_solicited =
+        mooncake_transfer_id.is_some() || (mode == KvConnectorMode::Nixl && relay_kv_params);
+    let mut handed_off = false;
     match (&mode, prefill_kv_params) {
         // Modern Mooncake: synthesized params under the minted transfer_id
         (
@@ -1173,6 +1177,7 @@ async fn execute_sequential_pd(
                 host,
                 *port,
             ));
+            handed_off = true;
         }
         // Legacy Mooncake (no engine_id discovered, or n>1): typed host/port injection
         (KvConnectorMode::Mooncake { host, port, .. }, _) => {
@@ -1182,6 +1187,7 @@ async fn execute_sequential_pd(
                 "vLLM PD: injecting kv_transfer_params into decode request"
             );
             decode_request.set_kv_transfer_params(host.clone(), *port);
+            handed_off = true;
         }
         (KvConnectorMode::Nixl | KvConnectorMode::Passthrough, Some(json)) if relay_kv_params => {
             debug!(
@@ -1190,6 +1196,7 @@ async fn execute_sequential_pd(
                 "vLLM PD: relaying prefill kv_transfer_params to decode request"
             );
             decode_request.set_kv_transfer_params_json(json);
+            handed_off = true;
         }
         (KvConnectorMode::Nixl, None) if relay_kv_params => {
             Metrics::record_pd_kv_transfer_failure();
@@ -1202,6 +1209,12 @@ async fn execute_sequential_pd(
         }
         _ => {}
     }
+    apply_prefill_media_identity(
+        &mut decode_request,
+        handed_off,
+        identity_solicited,
+        prefill_media_identity.as_ref(),
+    );
 
     // A stripped multimodal leg cannot fall back to a local recompute: the
     // image KV exists only on the prefill worker, so without a handoff the
@@ -1742,11 +1755,24 @@ mod tests {
     }
 
     /// Without a KV handoff decode recomputes the prompt locally, so it needs
-    /// the media itself and the identity must not replace its references.
+    /// the media itself and the identity must not replace its references:
+    /// n>1 never hands off, and a NIXL prefill that returned no params does
+    /// not either, even when it did return an identity.
     #[test]
     fn n_greater_than_one_keeps_media_refs_on_decode() {
         let mut decode = media_refs_request("n2", 2).clone_without_mm_pixels();
-        apply_prefill_media_identity(&mut decode, false, Some(&identity()));
+        apply_prefill_media_identity(&mut decode, false, false, Some(&identity()));
+        assert!(decode.has_vllm_media_refs());
+        let ProtoGenerateRequest::Vllm(request) = &decode else {
+            panic!("expected vLLM request");
+        };
+        assert!(request.mm_inputs.is_none());
+    }
+
+    #[test]
+    fn an_identity_without_a_kv_handoff_keeps_the_references() {
+        let mut decode = media_refs_request("no-kv", 1).clone_without_mm_pixels();
+        apply_prefill_media_identity(&mut decode, false, true, Some(&identity()));
         assert!(decode.has_vllm_media_refs());
         let ProtoGenerateRequest::Vllm(request) = &decode else {
             panic!("expected vLLM request");
@@ -1757,7 +1783,7 @@ mod tests {
     #[test]
     fn decode_leg_from_media_identity_has_no_refs_and_no_pixels() {
         let mut decode = media_refs_request("relay", 1).clone_without_mm_pixels();
-        apply_prefill_media_identity(&mut decode, true, Some(&identity()));
+        apply_prefill_media_identity(&mut decode, true, true, Some(&identity()));
         assert!(!decode.has_vllm_media_refs());
         let ProtoGenerateRequest::Vllm(request) = &decode else {
             panic!("expected vLLM request");
@@ -1776,12 +1802,15 @@ mod tests {
     }
 
     /// An older servicer returns no identity: the decode leg keeps its
-    /// references and reprocesses, as before.
+    /// references and reprocesses, as before; the same when no identity was
+    /// asked for (a prefill request without KV params).
     #[test]
     fn missing_identity_leaves_the_decode_leg_untouched() {
-        let mut decode = media_refs_request("old", 1).clone_without_mm_pixels();
-        apply_prefill_media_identity(&mut decode, true, None);
-        assert!(decode.has_vllm_media_refs());
+        for solicited in [true, false] {
+            let mut decode = media_refs_request("old", 1).clone_without_mm_pixels();
+            apply_prefill_media_identity(&mut decode, true, solicited, None);
+            assert!(decode.has_vllm_media_refs(), "solicited={solicited}");
+        }
     }
 
     #[test]

From 512a020ddaa568d298da8fb753f740b440d3a7c2 Mon Sep 17 00:00:00 2001
From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
Date: Tue, 22 Sep 2026 13:47:53 -0700
Subject: [PATCH 3/4] test(e2e): fail in CI when the PD media-once case finds
 no worker logs

The case skipped whenever a leg's log path was unset, which in CI, where
the lane always writes worker log files, would hide a broken capture
instead of reporting it. It now goes through assert_worker_logs_captured
like the other PD log assertions: a skip locally with SHOW_WORKER_LOGS=1,
a failure under CI.

Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
---
 e2e_test/router/test_pd_multimodal.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/e2e_test/router/test_pd_multimodal.py b/e2e_test/router/test_pd_multimodal.py
index 1bd000501c..c492ddb396 100644
--- a/e2e_test/router/test_pd_multimodal.py
+++ b/e2e_test/router/test_pd_multimodal.py
@@ -32,6 +32,7 @@
 from infra import assert_mm_processing, get_mm_processing
 from infra.constants import MM_PROCESSING_WORKER
 from infra.mm_processing import mm_processing_samples
+from infra.pd_logs import assert_worker_logs_captured
 from PIL import Image
 
 logger = logging.getLogger(__name__)
@@ -177,8 +178,11 @@ def test_worker_media_is_processed_once_on_the_prefill_leg(self, setup_backend):
             pytest.skip("router-side lane: the legs receive preprocessed tensors")
         backend, model, client, gateway = setup_backend
         prefill, decode = gateway.prefill_workers[0], gateway.decode_workers[0]
-        if prefill.log_path is None or decode.log_path is None:
-            pytest.skip("worker output goes to the terminal; the log assertions need captured logs")
+        # A serving worker has logged its startup, so an empty log means the
+        # output is not captured: a skip locally (SHOW_WORKER_LOGS=1), a
+        # failure in CI, where the lane always writes log files.
+        for leg, worker in (("prefill", prefill), ("decode", decode)):
+            assert_worker_logs_captured(worker.read_log(), f"the {leg} leg's media placement")
         worker_before = _worker_path_count(gateway)
         prefill_before = prefill.read_log().count("media_refs=1")
         decode_refs_before = decode.read_log().count("media_refs=1")

From b9327680a5f215f29a93b396d02748d9e712ccb0 Mon Sep 17 00:00:00 2001
From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
Date: Tue, 22 Sep 2026 19:16:02 -0700
Subject: [PATCH 4/4] fix(pd): never serve an n>1 decode leg from the media
 identity

Legacy Mooncake injects its bootstrap host and port into every decode
leg, n>1 included, so the handoff flag alone did not say that decode
would pull its prompt KV from prefill: an n>1 leg recomputes the prompt
locally and needs its media references, yet an identity returned by
prefill would have replaced them. The identity now also requires the
single-sample relay, with a test for the legacy Mooncake n>1 shape.

Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com>
---
 .../grpc/common/stages/request_execution.rs   | 46 +++++++++++++------
 1 file changed, 33 insertions(+), 13 deletions(-)

diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs
index 025a113853..17651dfdb8 100644
--- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs
+++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs
@@ -156,19 +156,23 @@ fn pd_fanout_width(request: &ProtoGenerateRequest, protocol: PdProtocol) -> Opti
 }
 
 /// Give the decode leg the media identity the prefill leg produced, so it
-/// is served without pixels or references. Only once decode actually holds
-/// a KV handoff (`handed_off`): without one it recomputes the prompt and
-/// needs the references itself, and a pixel-less leg without KV is refused
-/// by the servicer. A prefill that was asked for an identity (`solicited`:
-/// its request carried KV params) but returned none is an older servicer;
-/// the leg is left as it is, and decode reprocesses the media.
+/// is served without pixels or references. Only on a leg that will pull its
+/// prompt KV from prefill: it must hold a KV handoff (`handed_off`) and be
+/// a single-sample request (`relay_kv_params`); with n>1 decode recomputes
+/// the prompt and needs the references itself, even where legacy Mooncake
+/// still injects its host and port, and a pixel-less leg without KV is
+/// refused by the servicer. A prefill that was asked for an identity
+/// (`solicited`: its request carried KV params) but returned none is an
+/// older servicer; the leg is left as it is, and decode reprocesses the
+/// media.
 fn apply_prefill_media_identity(
     decode_request: &mut ProtoGenerateRequest,
     handed_off: bool,
+    relay_kv_params: bool,
     solicited: bool,
     identity: Option<&vllm::MediaIdentity>,
 ) {
-    if !handed_off || !decode_request.has_vllm_media_refs() {
+    if !handed_off || !relay_kv_params || !decode_request.has_vllm_media_refs() {
         return;
     }
     match identity {
@@ -1212,6 +1216,7 @@ async fn execute_sequential_pd(
     apply_prefill_media_identity(
         &mut decode_request,
         handed_off,
+        relay_kv_params,
         identity_solicited,
         prefill_media_identity.as_ref(),
     );
@@ -1756,12 +1761,27 @@ mod tests {
 
     /// Without a KV handoff decode recomputes the prompt locally, so it needs
     /// the media itself and the identity must not replace its references:
-    /// n>1 never hands off, and a NIXL prefill that returned no params does
-    /// not either, even when it did return an identity.
+    /// n>1 never relays, and a NIXL prefill that returned no params does not
+    /// hand off either, even when it did return an identity.
     #[test]
     fn n_greater_than_one_keeps_media_refs_on_decode() {
         let mut decode = media_refs_request("n2", 2).clone_without_mm_pixels();
-        apply_prefill_media_identity(&mut decode, false, false, Some(&identity()));
+        apply_prefill_media_identity(&mut decode, false, false, false, Some(&identity()));
+        assert!(decode.has_vllm_media_refs());
+        let ProtoGenerateRequest::Vllm(request) = &decode else {
+            panic!("expected vLLM request");
+        };
+        assert!(request.mm_inputs.is_none());
+    }
+
+    /// Legacy Mooncake injects its host and port into every decode leg, n>1
+    /// included, so `handed_off` alone does not say that decode will pull
+    /// its prompt KV from prefill: an n>1 leg recomputes locally and keeps
+    /// its references whatever the prefill returned.
+    #[test]
+    fn n_greater_than_one_keeps_media_refs_despite_a_legacy_mooncake_handoff() {
+        let mut decode = media_refs_request("n2-mooncake", 2).clone_without_mm_pixels();
+        apply_prefill_media_identity(&mut decode, true, false, false, Some(&identity()));
         assert!(decode.has_vllm_media_refs());
         let ProtoGenerateRequest::Vllm(request) = &decode else {
             panic!("expected vLLM request");
@@ -1772,7 +1792,7 @@ mod tests {
     #[test]
     fn an_identity_without_a_kv_handoff_keeps_the_references() {
         let mut decode = media_refs_request("no-kv", 1).clone_without_mm_pixels();
-        apply_prefill_media_identity(&mut decode, false, true, Some(&identity()));
+        apply_prefill_media_identity(&mut decode, false, true, true, Some(&identity()));
         assert!(decode.has_vllm_media_refs());
         let ProtoGenerateRequest::Vllm(request) = &decode else {
             panic!("expected vLLM request");
@@ -1783,7 +1803,7 @@ mod tests {
     #[test]
     fn decode_leg_from_media_identity_has_no_refs_and_no_pixels() {
         let mut decode = media_refs_request("relay", 1).clone_without_mm_pixels();
-        apply_prefill_media_identity(&mut decode, true, true, Some(&identity()));
+        apply_prefill_media_identity(&mut decode, true, true, true, Some(&identity()));
         assert!(!decode.has_vllm_media_refs());
         let ProtoGenerateRequest::Vllm(request) = &decode else {
             panic!("expected vLLM request");
@@ -1808,7 +1828,7 @@ mod tests {
     fn missing_identity_leaves_the_decode_leg_untouched() {
         for solicited in [true, false] {
             let mut decode = media_refs_request("old", 1).clone_without_mm_pixels();
-            apply_prefill_media_identity(&mut decode, true, solicited, None);
+            apply_prefill_media_identity(&mut decode, true, true, solicited, None);
             assert!(decode.has_vllm_media_refs(), "solicited={solicited}");
         }
     }