diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml index 8c061a95ad..f6e944d705 100644 --- a/.buildkite/test_areas/disaggregated_mooncake.yaml +++ b/.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" diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index f8136ae8f3..1ed84ebc52 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/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}") diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 096f0e4075..1f145bdc89 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh index 05af1bb6fd..c4e9b0d069 100755 --- a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh +++ b/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 } diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index b8565a0acd..4e37637ea4 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/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 diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index f8bd50c6d9..98aaa568f0 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/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() diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py index e19f1b332a..4dc0bf77dd 100644 --- a/vllm/v1/worker/mm_encoder_model_runner.py +++ b/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,