fix(pd): process worker-side media once, on the prefill leg

.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 ===

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;
 }
 
 // =====================

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 = [

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
 

e2e_test/router/test_pd_multimodal.py

@@ -29,7 +29,10 @@
 
 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 infra.pd_logs import assert_worker_logs_captured
 from PIL import Image
 
 logger = logging.getLogger(__name__)
@@ -42,6 +45,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 +170,37 @@ 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]
+        # 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")
+        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(

grpc_servicer/smg_grpc_servicer/vllm/media_identity.py

@@ -0,0 +1,126 @@
+"""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 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."""
+    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

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, 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,
@@ -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,17 @@ 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:
+                    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
@@ -410,6 +424,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 +435,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 +1200,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 +1214,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 +1257,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