[MM][V2] Enable encoder-only ViT CUDA graph capture (#56922)

Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Co-authored-by: Codex <noreply@openai.com>

.buildkite/test_areas/disaggregated_mooncake.yaml

@@ -40,7 +40,14 @@ steps:
     - vllm/v1/worker/gpu_worker.py
     - vllm/v1/worker/gpu/ec_connector.py
     - vllm/v1/worker/gpu/model_runner.py
+    - vllm/v1/worker/mm_encoder_model_runner.py
+    - vllm/v1/worker/encoder_cudagraph.py
+    - vllm/v1/worker/encoder_cudagraph_defs.py
+    - vllm/v1/attention/ops/vit_attn_wrappers.py
+    - vllm/v1/worker/gpu/model_states/
     - vllm/v1/worker/gpu/mm/
+    - tests/v1/cudagraph/test_encoder_cudagraph.py
+    - tests/models/multimodal/generation/test_vit_cudagraph.py
     - examples/disaggregated/disaggregated_encoder/
     - tests/v1/ec_connector/
     - requirements/kv_connectors.txt
@@ -49,6 +56,9 @@ steps:
     PYTHONUNBUFFERED: "1"
     VLLM_HOST_IP: "127.0.0.1"
     VLLM_USE_V2_MODEL_RUNNER: "1"
+    E_CUDAGRAPH_MM_ENCODER: "1"
+    # FA padding NaNs are tracked separately in #57136.
+    E_MM_ENCODER_ATTN_BACKEND: "TRITON_ATTN"
     MOONCAKE_EC_PROTOCOL: "tcp"
     USE_MM_PROMPTS: "1"
     SKIP_BASELINE: "0"

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -499,3 +499,118 @@ def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
 
         # Ensure the output is a string
         assert isinstance(output_text, str)
+
+
+def _check_eonly_encoder_outputs(worker, batches):
+    """Compare real E-only replay with eager, retaining outputs across replay."""
+    import torch
+
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+    runner = worker.model_runner
+    assert isinstance(runner, MMEncoderModelRunner)
+    manager = runner.model_state.encoder_runner.cudagraph_manager
+    assert manager is not None and manager.is_captured()
+    retained: list[tuple[torch.Tensor, torch.Tensor]] = []
+    max_error = 0.0
+    stats = []
+    with torch.inference_mode():
+        for batch in batches:
+            kwargs = {key: value.to(runner.device) for key, value in batch.items()}
+            actual = manager.execute(kwargs)
+            expected = runner.model.embed_multimodal(**kwargs)
+            assert len(actual) == len(expected)
+            for output, reference in zip(actual, expected):
+                assert torch.isfinite(output).all()
+                # Two BF16 ulps at unit scale, fixed before model validation.
+                torch.testing.assert_close(output, reference, rtol=0.016, atol=0.016)
+                max_error = max(max_error, (output - reference).abs().max().item())
+            for output, snapshot in retained:
+                torch.testing.assert_close(output, snapshot, rtol=0, atol=0)
+            retained.extend((output, output.clone()) for output in actual)
+            stats.append(manager.get_cumulative_stats())
+    assert stats[0]["graph_hits"] > 0
+    assert stats[1]["graph_hits"] > stats[0]["graph_hits"]
+    assert stats[2]["graph_misses"] > stats[1]["graph_misses"]
+    assert stats[3]["graph_hits"] > stats[2]["graph_hits"]
+    assert stats[4]["graph_hits"] > stats[3]["graph_hits"]
+    assert stats[5]["graph_hits"] > stats[4]["graph_hits"]
+    assert stats[6]["graph_hits"] > stats[5]["graph_hits"]
+    return {"max_abs_error": max_error, "stats": stats}
+
+
+@pytest.mark.parametrize(
+    "model_id",
+    params_with_marks({key: MODEL_CONFIGS[key] for key in ("qwen3_vl", "qwen3_5")}),
+)
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA graphs")
+def test_eonly_vit_cudagraph_outputs(
+    model_id, vllm_runner, image_assets, video_assets, monkeypatch
+):
+    """Exercise startup capture, image/video replay, fallback and output ownership."""
+    import numpy as np
+    from PIL import Image
+    from transformers import AutoProcessor
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    config = MODEL_CONFIGS[model_id]
+    token_budgets = [64, 256, 2048]
+    processor = AutoProcessor.from_pretrained(config.model)
+    first, second = [asset.pil_image for asset in image_assets]
+    batches = [
+        [first.resize((224, 224))],
+        [second.resize((448, 224)), first.resize((224, 224))],
+        [second.resize((1792, 1792))],
+        [second.resize((224, 224))],
+        [first.resize((1280, 720))],
+        [first.resize((224, 224))],
+    ]
+    frames = sample_frames_from_video(video_assets[0].np_ndarrays, 2)
+    video = np.stack(
+        [np.asarray(Image.fromarray(frame).resize((224, 224))) for frame in frames]
+    )
+    with vllm_runner(
+        config.model,
+        dtype=config.dtype,
+        mm_encoder_only=True,
+        # FA padding NaNs are tracked separately in #57136.
+        mm_encoder_attn_backend="FLASHINFER",
+        enable_prefix_caching=False,
+        max_model_len=4096,
+        max_num_seqs=2,
+        limit_mm_per_prompt={"image": 2, "video": 1},
+        compilation_config={
+            "cudagraph_mode": "NONE",
+            "cudagraph_mm_encoder": True,
+            "encoder_cudagraph_token_budgets": token_budgets,
+            "encoder_cudagraph_max_vision_items_per_batch": 2,
+            "encoder_cudagraph_max_frames_per_batch": 2,
+        },
+    ) as model:
+        mm_config = model.llm.llm_engine.vllm_config.model_config.multimodal_config
+        processor_kwargs = mm_config.merge_mm_processor_kwargs({})
+        inputs = [
+            dict(
+                processor.image_processor(
+                    images=images, return_tensors="pt", **processor_kwargs
+                )
+            )
+            for images in batches
+        ]
+        inputs.append(
+            dict(
+                processor.video_processor(
+                    videos=[video],
+                    do_sample_frames=False,
+                    return_tensors="pt",
+                    **processor_kwargs,
+                )
+            )
+        )
+        results = model.llm.collective_rpc(
+            partial(_check_eonly_encoder_outputs, batches=inputs)
+        )
+        for result in results:
+            assert result["stats"][0]["token_budgets"] == token_budgets
+        print(f"E-only {model_id}: {results}")

tests/v1/cudagraph/test_encoder_cudagraph.py

@@ -13,7 +13,9 @@ Test organization:
 """
 
 from collections.abc import Hashable
+from types import SimpleNamespace
 from typing import Any
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -575,6 +577,99 @@ class TestEncoderCudaGraphCaptureReplay:
             assert out.shape == (4, _HIDDEN)
 
 
+# ---------------------------------------------------------------------------
+# E-only capture and output lifecycle
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm"
+)
+@pytest.mark.usefixtures("dist_init", "workspace_init")
+@pytest.mark.parametrize("profile_only", [False, True])
+@torch.inference_mode()
+def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only):
+    """The E-only entry captures only the encoder and preserves cached outputs."""
+    from vllm.distributed.ec_transfer.ec_connector.base import (
+        ECConnectorBase,
+        ECConnectorMetadata,
+    )
+    from vllm.v1.worker.gpu.ec_connector import ActiveECConnector
+    from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+    from vllm.v1.worker.workspace import lock_workspace
+
+    device = torch.device("cuda:0")
+    dtype = torch.float16
+    model = SimpleMockViTModel().to(device).half()
+    manager = _make_manager_for_gpu(model, _BUDGETS, _MAX_BATCH, device, dtype)
+    encoder = object.__new__(EncoderRunner)
+    encoder.device = device
+    encoder.cudagraph_manager = manager
+    runner = object.__new__(MMEncoderModelRunner)
+    runner.model_state = SimpleNamespace(encoder_runner=encoder)
+    # No decoder manager is installed: capture must be encoder-only.
+    with patch(
+        "vllm.v1.worker.mm_encoder_model_runner.lock_workspace", wraps=lock_workspace
+    ) as lock:
+        runner.capture_model(profile_only=profile_only)
+        assert lock.call_count == int(not profile_only)
+    assert len(manager.budget_graphs["default"]) == len(_BUDGETS)
+
+    cache: dict[str, torch.Tensor] = {}
+    pending_sends: dict[str, torch.Tensor] = {}
+    connector = MagicMock(spec=ECConnectorBase)
+    connector.is_producer = True
+    connector.is_consumer = False
+    connector.get_finished.return_value = (None, None)
+    connector.save_caches.side_effect = lambda *, encoder_cache, mm_hash: (
+        pending_sends.update({mm_hash: encoder_cache[mm_hash]})
+    )
+    with patch(
+        "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector
+    ):
+        ec = ActiveECConnector(SimpleNamespace(), cache)
+    scheduled = SimpleNamespace(
+        ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()
+    )
+    saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = []
+    # Mixed sizes exercise packing order; 64 is the boundary, 81 falls back.
+    for grids in ([[1, 8, 8], [1, 4, 4]], [[1, 16, 16]], [[1, 18, 18]], [[1, 4, 4]]):
+        inputs = _make_mm_kwargs(grids, device, dtype)
+        expected = model.encoder_eager_forward(inputs).split(
+            [t * (h // 2) * (w // 2) for t, h, w in grids]
+        )
+        with ec.maybe_get_output(scheduled) as ec_output:
+            outputs = manager.execute(inputs)
+            assert outputs is not None
+            cache[str(len(cache))] = outputs[0]
+        assert ec_output.finished_sending is None
+        assert outputs is not None
+        for actual, eager in zip(outputs, expected):
+            torch.testing.assert_close(actual, eager)
+        for previous, snapshot in saved_outputs:
+            torch.testing.assert_close(previous, snapshot, rtol=0, atol=0)
+        saved_outputs.extend((output, output.clone()) for output in outputs)
+    assert manager.graph_hits == 4
+    assert manager.graph_misses == 1
+    assert pending_sends.keys() == cache.keys()
+    connector.get_finished.return_value = (set(cache), None)
+    with ec.maybe_get_output(scheduled) as ec_output:
+        pass
+    assert ec_output.finished_sending == set(cache)
+
+
+def test_eonly_without_encoder_graph_skips_capture():
+    from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+    encoder = object.__new__(EncoderRunner)
+    encoder.cudagraph_manager = None
+    runner = object.__new__(MMEncoderModelRunner)
+    runner.model_state = SimpleNamespace(encoder_runner=encoder)
+    assert runner.capture_model() == 0
+
+
 # ---------------------------------------------------------------------------
 # SimpleMockViTVideoModel — extends SimpleMockViTModel with video support
 # ---------------------------------------------------------------------------

tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh

@@ -20,6 +20,8 @@
 #   SKIP_BASELINE               set to 1 to reuse existing BASELINE_FILE
 #   CONCURRENCY / REPEAT        concurrent requests / rounds (defaults 3 / 2)
 #   MAX_MODEL_LEN               context length (default 16384)
+#   E_CUDAGRAPH_MM_ENCODER       enable encoder CUDA graphs (default 0)
+#   E_MM_ENCODER_ATTN_BACKEND    optional encoder attention backend override
 
 set -euo pipefail
 
@@ -35,6 +37,15 @@ if [[ "$USE_MM_PROMPTS" == "1" ]]; then
   TEST_ARGS+=(--use_mm_prompts --mm_smoke_test)
 fi
 MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
+E_ATTN_ARGS=()
+if [[ -n "${E_MM_ENCODER_ATTN_BACKEND:-}" ]]; then
+  E_ATTN_ARGS+=(--mm-encoder-attn-backend "$E_MM_ENCODER_ATTN_BACKEND")
+fi
+case "${E_CUDAGRAPH_MM_ENCODER:-0}" in
+  0) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":false}' ;;
+  1) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":true}' ;;
+  *) echo "E_CUDAGRAPH_MM_ENCODER must be 0 or 1"; exit 1 ;;
+esac
 
 GPU_SINGLE="${GPU_SINGLE:-0}"
 GPU_E="${GPU_E:-0}"
@@ -203,7 +214,7 @@ run_epd_mooncake() {
   cleanup_instances
 
   echo "Starting ENCODER on GPU $GPU_E port $ENCODE_PORT"
-  CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \
+  VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \
     --port "$ENCODE_PORT" \
     --gpu-memory-utilization 0.35 \
     --mm-tensor-ipc torch_shm \
@@ -215,6 +226,9 @@ run_epd_mooncake() {
     --limit-mm-per-prompt '{"image":2,"video":0}' \
     --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
     --ec-transfer-config "$ENC_EC_JSON" \
+    --compilation-config "$E_COMPILATION_CONFIG" \
+    "${E_ATTN_ARGS[@]}" \
+    --worker-extension-cls tests.v1.ec_connector.integration.test_epd_correctness.EncoderGraphTestWorkerExtension \
     >"${LOG_PATH}/mooncake_epd_encoder.log" 2>&1 &
   local ENCODER_PID=$!
   PIDS+=("$ENCODER_PID")
@@ -266,6 +280,26 @@ run_epd_mooncake() {
     --baseline_file "$BASELINE_FILE" \
     "${TEST_ARGS[@]}"
 
+  curl -fsS "http://localhost:${ENCODE_PORT}/collective_rpc" \
+    -H 'Content-Type: application/json' \
+    -d '{"method":"encoder_graph_stats"}' \
+    >"${LOG_PATH}/encoder_graph_stats.json"
+  "$PYTHON_BIN" - "${LOG_PATH}/encoder_graph_stats.json" "${E_CUDAGRAPH_MM_ENCODER:-0}" <<'PY'
+import json
+import sys
+
+with open(sys.argv[1]) as f:
+    results = json.load(f)["results"]
+assert results, "No encoder workers returned graph statistics"
+for stats in results:
+    if sys.argv[2] == "1":
+        assert stats["is_captured"], stats
+        assert stats["graph_hits"] > 0, stats
+    else:
+        assert stats["graph_hits"] == 0, stats
+print("Encoder graph statistics:", results)
+PY
+
   cleanup_instances
 }
 

tests/v1/ec_connector/integration/test_epd_correctness.py

@@ -26,6 +26,7 @@ import json
 import os
 import time
 from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING
 
 import openai
 import requests
@@ -33,6 +34,26 @@ import requests
 from vllm.assets.image import ImageAsset
 from vllm.multimodal.utils import encode_image_url
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+
+class EncoderGraphTestWorkerExtension:
+    """Expose encoder replay evidence through the existing test RPC endpoint."""
+
+    model_runner: "MMEncoderModelRunner"
+
+    def encoder_graph_stats(self):
+        from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+        runner = self.model_runner
+        assert isinstance(runner, MMEncoderModelRunner)
+        manager = runner.model_state.encoder_runner.cudagraph_manager
+        if manager is None:
+            return {"is_captured": False, "graph_hits": 0, "graph_misses": 0}
+        return {"is_captured": manager.is_captured(), **manager.get_cumulative_stats()}
+
+
 MAX_OUTPUT_LEN = 256
 
 # Sample prompts with multimodal content

tests/v1/worker/test_encoder_runner.py

@@ -9,12 +9,15 @@ and tolerated (token-embedding fallback) when it is not, while a miss within
 the processed range still fails loudly.
 """
 
-from unittest.mock import MagicMock
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
 
+from vllm.config.compilation import CUDAGraphMode
+from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalFieldElem,
@@ -31,6 +34,56 @@ pytestmark = pytest.mark.cpu_test
 HIDDEN = 4
 
 
+@pytest.mark.parametrize(
+    "enabled,enforce_eager,supported,expected",
+    [
+        (True, False, True, True),
+        (False, False, True, False),
+        (True, True, True, False),
+        (True, False, False, False),
+    ],
+)
+def test_encoder_graph_configuration(enabled, enforce_eager, supported, expected):
+    """Encoder graph eligibility is independent of decoder graph mode."""
+    config = MagicMock()
+    config.model_config.enforce_eager = enforce_eager
+    config.model_config.dtype = torch.float32
+    config.model_config.get_inputs_embeds_size.return_value = HIDDEN
+    config.scheduler_config.max_num_batched_tokens = 8
+    config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+    config.compilation_config.cudagraph_mm_encoder = enabled
+    model = MagicMock(spec=SupportsEncoderCudaGraph) if supported else torch.nn.Module()
+    state = _model_state(EncoderCache())
+    with patch("vllm.v1.worker.gpu.model_states.interface.EncoderCudaGraphManager"):
+        ModelState.__init__(state, config, model, state.encoder_cache, state.device)
+    assert state.encoder_runner.has_cudagraph() is expected
+
+
+def test_encoder_graph_unsupported_modality_uses_eager_output():
+    """An image graph must not intercept another modality's encoder output."""
+    from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
+
+    runner = _make_runner([], [])
+    output = torch.ones(2, HIDDEN)
+    runner.model = MagicMock()
+    runner.model.embed_multimodal.return_value = [output]
+    manager = MagicMock(spec=EncoderCudaGraphManager)
+    manager.config = SimpleNamespace(modalities=["image"])
+    manager.is_captured.return_value = True
+    manager.supports_modality.side_effect = lambda modality: (
+        EncoderCudaGraphManager.supports_modality(manager, modality)
+    )
+    runner.cudagraph_manager = manager
+    with patch(
+        "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs",
+        return_value=[("audio", 1, {"audio_values": torch.zeros(2)})],
+    ):
+        result = runner.execute_mm_encoder([])
+    assert len(result) == 1
+    assert result[0] is output
+    manager.execute.assert_not_called()
+
+
 def _model_state(cache: EncoderCache) -> MagicMock:
     """A mock ModelState backed by a real EncoderCache."""
     state = MagicMock()

vllm/v1/worker/mm_encoder_model_runner.py

@@ -4,9 +4,10 @@
 
 An encoder-only instance -- `--mm-encoder-only`, or the producer side of
 encoder-cache disaggregation -- encodes the multi-modal items and publishes the
-embeddings. It runs no language model: no KV cache, no sampler, no CUDA graphs.
+embeddings. It runs no language model: no KV cache, sampler, or decoder graphs.
 """
 
+import time
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any
@@ -15,7 +16,9 @@ import torch
 
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.logger import init_logger
 from vllm.sequence import IntermediateTensors
+from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
 from vllm.v1.kv_cache_interface import KVCacheSpec
 from vllm.v1.outputs import (
     ModelRunnerOutput,
@@ -24,10 +27,13 @@ from vllm.v1.outputs import (
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.workspace import lock_workspace
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
 
+logger = init_logger(__name__)
+
 
 class MMEncoderModelRunner(GPUModelRunner):
     """Encoder-only variant of the V2 GPU model runner."""
@@ -67,8 +73,31 @@ class MMEncoderModelRunner(GPUModelRunner):
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         return {}
 
+    @torch.inference_mode()
     def capture_model(self, *, profile_only: bool = False) -> int:
-        return 0
+        encoder_runner = self.model_state.encoder_runner
+        if not encoder_runner.has_cudagraph():
+            return 0
+
+        start_time = time.perf_counter()
+        with freeze_gc_for_cudagraph_capture():
+            torch.accelerator.empty_cache()
+            start_free_memory = torch.accelerator.get_memory_info()[0]
+            encoder_runner.capture()
+            graph_memory = start_free_memory - torch.accelerator.get_memory_info()[0]
+
+        if not profile_only:
+            lock_workspace()
+
+        manager = encoder_runner.cudagraph_manager
+        assert manager is not None
+        logger.info(
+            "Encoder graph capturing finished in %.0f secs, took %.2f GiB (%d graphs)",
+            time.perf_counter() - start_time,
+            graph_memory / (1 << 30),
+            manager.get_num_graphs_to_capture(),
+        )
+        return graph_memory
 
     def needs_cudagraph_capture(self) -> bool:
         return False
@@ -114,7 +143,7 @@ class MMEncoderModelRunner(GPUModelRunner):
 
             batch_req_state, _ = self.gather_batch_req_state(scheduler_output, False)
             assert batch_req_state is not None
-            # No CUDA graph, and no DP peer to agree a padded shape with.
+            # No decoder graph or DP padding; encoder graphs run independently.
             batch_desc = BatchExecutionDescriptor(
                 cg_mode=CUDAGraphMode.NONE,
                 num_tokens=batch_req_state.num_tokens,