[Optimization] Return image_grid_thw in render response for disaggregated mRoPE

tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py

@@ -123,6 +123,16 @@ async def test_render_to_generate_roundtrip(client, test_image):
     assert "image" in features["kwargs_data"]
     assert len(features["kwargs_data"]["image"]) > 0
 
+    assert "image_grid_thw" in features
+    assert "image" in features["image_grid_thw"]
+    image_grids = features["image_grid_thw"]["image"]
+    assert isinstance(image_grids, list)
+    assert len(image_grids) == len(features["kwargs_data"]["image"])
+    for grid in image_grids:
+        assert isinstance(grid, list)
+        assert len(grid) == 3
+        assert all(isinstance(v, int) and v > 0 for v in grid)
+
     # Build generate request from render output
     generate_payload = render_data
     generate_payload["sampling_params"] = {

vllm/entrypoints/serve/disagg/protocol.py

@@ -58,6 +58,15 @@ class MultiModalFeatures(BaseModel):
     ``None`` for metadata-only (cache-hit) responses.
     """
 
+    image_grid_thw: dict[str, list[list[int] | None]] | None = None
+    """Per-modality grid dimensions for mRoPE position computation.
+
+    Each value is a list parallel to ``mm_hashes[modality]``.  Each entry
+    is a ``[grid_t, grid_h, grid_w]`` list (or ``None`` for cache hits).
+    This allows the prefill worker to compute mRoPE positions without
+    deserializing the full ``kwargs_data`` blobs.
+    """
+
 
 class GenerateRequest(BaseModel):
     request_id: str = Field(

vllm/entrypoints/serve/disagg/serving.py

@@ -11,6 +11,7 @@
 import msgspec
 import numpy as np
 import pybase64 as base64
+import torch
 from fastapi import Request
 
 from vllm.engine.protocol import EngineClient
@@ -43,6 +44,8 @@
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
 from vllm.multimodal.inputs import (
+    MultiModalBatchedField,
+    MultiModalFieldElem,
     MultiModalKwargsItem,
     MultiModalKwargsItems,
     PlaceholderRange,
@@ -156,6 +159,23 @@ async def serve_tokens(
                         decode_mm_kwargs_item(item) if item is not None else None
                         for item in items
                     ]
+            elif features.image_grid_thw is not None:
+                # Lightweight path: construct minimal items containing
+                # only grid metadata for mRoPE position computation.
+                for modality, grids in features.image_grid_thw.items():
+                    items_list: list[MultiModalKwargsItem | None] = []
+                    thw_key = f"{modality}_grid_thw"
+                    for grid in grids:
+                        if grid is not None:
+                            tensor = torch.tensor(grid, dtype=torch.int64)
+                            elem = MultiModalFieldElem(
+                                data=tensor,
+                                field=MultiModalBatchedField(keep_on_cpu=True),
+                            )
+                            items_list.append(MultiModalKwargsItem({thw_key: elem}))
+                        else:
+                            items_list.append(None)
+                    mm_kwargs[modality] = items_list
             else:
                 for modality, hashes in features.mm_hashes.items():
                     mm_kwargs[modality] = [None] * len(hashes)

vllm/entrypoints/serve/render/serving.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Sequence
 from http import HTTPStatus
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from openai_harmony import Message as OpenAIMessage
 
+if TYPE_CHECKING:
+    import torch
+
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import (
     ChatTemplateContentFormatOption,
@@ -380,18 +383,33 @@ def _extract_mm_features(
 
         # Serialize tensor data per modality.
         kwargs_data: dict[str, list[str | None]] | None = None
+        image_grid_thw: dict[str, list[list[int] | None]] | None = None
         if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
             kwargs_data = {}
+            image_grid_thw = {}
             for modality, items in raw_mm_kwargs.items():
                 kwargs_data[modality] = [
                     encode_mm_kwargs_item(item) if item is not None else None
                     for item in items
                 ]
+                thw_key = f"{modality}_grid_thw"
+                grids: list[list[int] | None] = []
+                for item in items:
+                    if item is not None and thw_key in item:
+                        thw_tensor = cast("torch.Tensor", item[thw_key].data)
+                        grids.append(thw_tensor.tolist())
+                    else:
+                        grids.append(None)
+                if any(g is not None for g in grids):
+                    image_grid_thw[modality] = grids
+            if not image_grid_thw:
+                image_grid_thw = None
 
         return MultiModalFeatures(
             mm_hashes=mm_hashes,
             mm_placeholders=mm_placeholders,
             kwargs_data=kwargs_data,
+            image_grid_thw=image_grid_thw,
         )
 
     def _make_request_with_harmony(