[Feature][Multimodal] Support KV-backed continuation through Generate inputs

docs/features/disagg_prefill.md

@@ -79,6 +79,14 @@ decode = client.chat.completions.create(
 
 If `messages` has non-text content or `echo` is set, the ids are ignored and `messages` is rendered instead, so it must match the prefill request. Otherwise `kv_transfer_params["prompt_token_ids"]` must be a non-empty list of non-negative integers, or the request fails with HTTP 400, as it always does on `/v1/chat/completions/batch`.
 
+## Continuing a multimodal request from existing KV
+
+A decode instance can continue a multimodal request without the media and without encoder capacity when the prompt's KV already exists, locally or through a KV connector. The [Generate API](../serving/online_serving/token_in_token_out.md) input carries the prefill-side `token_ids` and a `features` object with the item hashes and placeholder ranges, and `features.requires_kv: true` states that every item must be covered by decoder KV. The same hashes and positions make the decode instance compute the same block hashes as prefill, so prefix caching and hash-keyed KV stores see the same blocks on both sides. `--limit-mm-per-prompt` may be 0 on the decode instance: such inputs are exempt from the encoder-cache capacity check.
+
+The flag is a requirement, not proof that the KV exists. Before scheduling any computation, the engine checks that the computed tokens cover every required span, including after a preemption or a failed KV load, and fails the request (`finish_reason="error"`) rather than recompute an image from placeholder token ids. Items must carry a nonempty hash and no media payload; the prompt must end with at least one token after the last multimodal span, because the final prompt token is always recomputed; prompt logprobs, pooling, M-RoPE models and encoder-decoder models are not supported. Local prefix-cache hits are whole blocks and never include the block holding the final prompt token, so without a connector the spans must end within the first `floor((N - 1) / block_size) * block_size` tokens of an `N`-token prompt.
+
+Python callers submit the same input directly to `AsyncLLM.generate()` or `LLMEngine.add_request()` as a `MultiModalInput` with `mm_kwargs` set to `None` per item and `mm_requires_kv=True`. `LLM.generate()` and the renderers reject such inputs, because rendering them again would drop the requirement and the media identities.
+
 ## Generate API output modes
 
 When the prefill and decode stages use the [Generate API](../serving/online_serving/token_in_token_out.md) (`/inference/v1/generate`), only the decode response reaches the client, so set `output_mode` on the decode request only. A proxy that reuses the client's request body for the prefill request must reset `output_mode` to `tokens` there. A prefill instance started with `--tokens-only` has no tokenizer and rejects `output_mode: "text"` with a 400.

tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py

@@ -1878,3 +1878,50 @@ async def mock_generate(*args, **kwargs):
         .delta.content
         == "a"
     )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("remote_kv", [False, True])
+async def test_serve_tokens_requires_kv_features_skip_rendering_and_media(remote_kv):
+    engine = _mock_engine()
+    payload: dict[str, Any] = {
+        "token_ids": [1, 2, 2, 2, 3],
+        "features": {
+            "mm_hashes": {"image": ["image-a"]},
+            "mm_placeholders": {"image": [{"offset": 1, "length": 3}]},
+            "requires_kv": True,
+        },
+        "cache_salt": "tenant-a",
+        "sampling_params": {"max_tokens": 1},
+        "return_token_ids": True,
+    }
+    if remote_kv:
+        payload["kv_transfer_params"] = {"do_remote_prefill": True}
+    request = GenerateRequest.model_validate_json(json.dumps(payload))
+
+    async def mock_generate(prompt, sampling_params, *args, **kwargs):
+        assert prompt["prompt_token_ids"] == payload["token_ids"]
+        assert prompt["mm_hashes"] == {"image": ["image-a"]}
+        assert prompt["mm_kwargs"] == {"image": [None]}
+        assert prompt["mm_requires_kv"] is True
+        assert prompt["cache_salt"] == "tenant-a"
+        (placeholder,) = prompt["mm_placeholders"]["image"]
+        assert (placeholder.offset, placeholder.length) == (1, 3)
+        if remote_kv:
+            assert sampling_params.extra_args["kv_transfer_params"] == {
+                "do_remote_prefill": True
+            }
+        yield _make_request_output(
+            "req-1", token_ids=[10], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    engine.renderer.render_cmpl_async = AsyncMock(
+        side_effect=AssertionError("processed features must not be rendered again")
+    )
+    response = await serving.serve_tokens(request)
+    assert isinstance(response, GenerateTokensResponse)
+    engine.generate.assert_called_once()
+    serving.online_renderer.preprocess_completion.assert_not_awaited()
+    engine.renderer.render_cmpl_async.assert_not_awaited()

tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py

@@ -13,6 +13,7 @@
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.entrypoints.scale_out.token_in_token_out.mm_features import (
+    engine_input_from_features,
     extract_mm_features,
     merge_mm_kwargs_items,
     mm_kwargs_from_features,
@@ -26,7 +27,7 @@
     MultiModalFeatures,
     PlaceholderRangeInfo,
 )
-from vllm.inputs import mm_input
+from vllm.inputs import MultiModalInput, mm_input
 from vllm.multimodal.inputs import (
     MultiModalBatchedField,
     MultiModalFieldElem,
@@ -131,7 +132,7 @@ def _image_engine_input(
     *,
     pixel_keep_on_cpu: bool = False,
     grid_keep_on_cpu: bool = True,
-) -> tuple[object, MultiModalFieldElem, MultiModalFieldElem]:
+) -> tuple[MultiModalInput, MultiModalFieldElem, MultiModalFieldElem]:
     pixel_values = MultiModalFieldElem(
         data=torch.randn(5, 3, dtype=torch.float32),
         field=MultiModalBatchedField(keep_on_cpu=pixel_keep_on_cpu),
@@ -422,3 +423,74 @@ def test_metadata_can_replace_or_extend_full_mm_data():
     merged = merge_mm_kwargs_items(full_item, metadata_item)
     assert merged is not None
     assert set(merged) == {"pixel_values", "image_grid_thw"}
+
+
+def test_requires_kv_is_omitted_from_wire_format_when_false():
+    engine_input, _, _ = _image_engine_input()
+    features = extract_mm_features(engine_input)
+    assert features is not None
+    assert features.requires_kv is False
+    assert "requires_kv" not in features.model_dump()
+    assert "requires_kv" not in features.model_dump_json()
+    restored = MultiModalFeatures.model_validate_json(features.model_dump_json())
+    assert restored.requires_kv is False
+
+
+@pytest.mark.parametrize("has_transfer", [False, True])
+def test_requires_kv_restores_media_free_engine_input(has_transfer):
+    engine_input, _, _ = _image_engine_input()
+    features = extract_mm_features(engine_input).model_dump()
+    features["kwargs_data"] = None
+    features["mm_metadata"] = None
+    features["requires_kv"] = True
+    request = GenerateRequest.model_validate(
+        {
+            "token_ids": engine_input["prompt_token_ids"],
+            "features": features,
+            "sampling_params": {},
+            "cache_salt": "tenant-a",
+            "kv_transfer_params": {"remote": True} if has_transfer else None,
+        }
+    )
+    assert request.features is not None
+    assert request.features.requires_kv is True
+    restored = engine_input_from_features(
+        request.token_ids, request.features, cache_salt=request.cache_salt
+    )
+    assert restored["mm_requires_kv"] is True
+    assert restored["cache_salt"] == "tenant-a"
+    assert restored["mm_hashes"] == engine_input["mm_hashes"]
+    assert restored["mm_kwargs"]["image"][0] is None
+    assert [(p.offset, p.length) for p in restored["mm_placeholders"]["image"]] == [
+        (p.offset, p.length) for p in engine_input["mm_placeholders"]["image"]
+    ]
+
+
+def test_features_without_kv_requirement_restore_their_payloads():
+    engine_input, _, _ = _image_engine_input()
+    features = extract_mm_features(engine_input)
+    restored = engine_input_from_features(engine_input["prompt_token_ids"], features)
+    assert restored.get("mm_requires_kv", False) is False
+    assert "pixel_values" in restored["mm_kwargs"]["image"][0]
+
+
+@pytest.mark.parametrize(
+    ("change", "error"),
+    [
+        ("payload", "serialized media payloads"),
+        ("metadata", "serialized media payloads"),
+        ("hash", "nonempty hash"),
+    ],
+)
+def test_requires_kv_rejects_media_payloads_and_empty_hashes(change, error):
+    engine_input, _, _ = _image_engine_input()
+    features = extract_mm_features(engine_input).model_dump()
+    if change == "metadata":
+        features["mm_metadata"] = features["kwargs_data"]
+    if change != "payload":
+        features["kwargs_data"] = None
+    if change == "hash":
+        features["mm_hashes"]["image"][0] = ""
+    features["requires_kv"] = True
+    with pytest.raises(ValidationError, match=error):
+        MultiModalFeatures.model_validate(features)