[Bugfix][Spec Decode] Disable EAGLE prefill FULL CUDA graph for live multimodal batches

tests/v1/cudagraph/test_gpu_cudagraph_utils.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+import vllm.v1.worker.gpu.dp_utils as dp_utils
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    BatchExecutionDescriptor,
+    CudaGraphManager,
+)
+from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
+from vllm.v1.worker.gpu.spec_decode.eagle import speculator as eagle_speculator
+
+
+def _make_manager(
+    candidates: list[BatchExecutionDescriptor],
+) -> CudaGraphManager:
+    manager = object.__new__(CudaGraphManager)
+    manager._graphs_captured = True
+    max_num_tokens = max(desc.num_tokens for desc in candidates)
+    manager._candidates = [[] for _ in range(max_num_tokens + 1)]
+    manager._candidates[max_num_tokens] = candidates
+    return manager
+
+
+def _live_mm_inputs() -> tuple[list[torch.Tensor], torch.Tensor]:
+    return [torch.empty(1, 1)], torch.ones(1, dtype=torch.bool)
+
+
+def test_eagle_prefill_text_only_batch_keeps_full_cudagraph_enabled():
+    full_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.FULL,
+        num_tokens=4,
+        num_reqs=1,
+        uniform_token_count=4,
+    )
+    piecewise_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.PIECEWISE,
+        num_tokens=4,
+        num_reqs=None,
+    )
+    manager = _make_manager([full_desc, piecewise_desc])
+
+    invalid_modes = eagle_speculator._get_prefill_invalid_cudagraph_modes(None)
+
+    assert invalid_modes is None
+    assert manager.dispatch(1, 4, 4, invalid_modes=invalid_modes) == full_desc
+
+
+def test_eagle_prefill_live_mm_batch_skips_full_and_uses_piecewise():
+    full_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.FULL,
+        num_tokens=4,
+        num_reqs=1,
+        uniform_token_count=4,
+    )
+    piecewise_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.PIECEWISE,
+        num_tokens=4,
+        num_reqs=None,
+    )
+    manager = _make_manager([full_desc, piecewise_desc])
+
+    invalid_modes = eagle_speculator._get_prefill_invalid_cudagraph_modes(
+        _live_mm_inputs()
+    )
+
+    assert invalid_modes == {CUDAGraphMode.FULL}
+    assert (
+        manager.dispatch(
+            1,
+            4,
+            4,
+            invalid_modes=invalid_modes,
+        )
+        == piecewise_desc
+    )
+
+
+def test_eagle_prefill_live_mm_batch_falls_back_to_none_without_piecewise():
+    full_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.FULL,
+        num_tokens=4,
+        num_reqs=1,
+        uniform_token_count=4,
+    )
+    manager = _make_manager([full_desc])
+
+    invalid_modes = eagle_speculator._get_prefill_invalid_cudagraph_modes(
+        _live_mm_inputs()
+    )
+    desc = manager.dispatch(
+        1,
+        4,
+        4,
+        invalid_modes=invalid_modes,
+    )
+
+    assert invalid_modes == {CUDAGraphMode.FULL}
+    assert desc == BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.NONE,
+        num_tokens=4,
+        num_reqs=1,
+    )
+
+
+def test_eagle_prefill_invalid_modes_only_for_live_mm_inputs():
+    assert eagle_speculator._get_prefill_invalid_cudagraph_modes(None) is None
+    assert (
+        eagle_speculator._get_prefill_invalid_cudagraph_modes(
+            ([], torch.zeros(0, dtype=torch.bool))
+        )
+        is None
+    )
+
+    assert eagle_speculator._get_prefill_invalid_cudagraph_modes(_live_mm_inputs()) == {
+        CUDAGraphMode.FULL
+    }
+
+
+def test_dispatch_cg_and_sync_dp_forwards_invalid_modes():
+    class RecordingManager:
+        invalid_modes = None
+
+        def dispatch(
+            self,
+            num_reqs: int,
+            num_tokens: int,
+            uniform_token_count: int | None,
+            invalid_modes: set[CUDAGraphMode] | None = None,
+        ) -> BatchExecutionDescriptor:
+            self.invalid_modes = invalid_modes
+            return BatchExecutionDescriptor(
+                cg_mode=CUDAGraphMode.NONE,
+                num_tokens=num_tokens,
+                num_reqs=num_reqs,
+            )
+
+    manager = RecordingManager()
+    desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
+        manager,  # type: ignore[arg-type]
+        num_reqs=1,
+        num_tokens=4,
+        uniform_token_count=4,
+        dp_size=1,
+        dp_rank=0,
+        invalid_modes={CUDAGraphMode.FULL},
+    )
+
+    assert manager.invalid_modes == {CUDAGraphMode.FULL}
+    assert desc.cg_mode == CUDAGraphMode.NONE
+    assert num_tokens_across_dp is None
+
+
+def test_dp_sync_does_not_redispatch_higher_mode(monkeypatch):
+    full_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.FULL,
+        num_tokens=4,
+        num_reqs=1,
+        uniform_token_count=4,
+    )
+    piecewise_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.PIECEWISE,
+        num_tokens=4,
+        num_reqs=None,
+    )
+    manager = _make_manager([full_desc, piecewise_desc])
+
+    class FakeDpGroup:
+        cpu_group = object()
+
+    def fake_all_reduce(tensor: torch.Tensor, group: object) -> None:
+        del group
+        tensor[:, :] = torch.tensor(
+            [
+                [4, 4],
+                [CUDAGraphMode.PIECEWISE.value, CUDAGraphMode.FULL.value],
+                [4, 4],
+            ],
+            dtype=tensor.dtype,
+            device=tensor.device,
+        )
+
+    monkeypatch.setattr(dp_utils, "get_dp_group", lambda: FakeDpGroup())
+    monkeypatch.setattr(dp_utils.dist, "all_reduce", fake_all_reduce)
+
+    desc, _ = dp_utils.sync_cudagraph_and_dp_padding(
+        manager,
+        full_desc,
+        num_tokens=4,
+        num_reqs=1,
+        uniform_token_count=4,
+        dp_size=2,
+        dp_rank=1,
+    )
+
+    assert desc == piecewise_desc

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import defaultdict
 from collections.abc import Callable
+from collections.abc import Set as AbstractSet
 from dataclasses import dataclass
 from typing import Any, NamedTuple
 
@@ -258,10 +259,16 @@ def dispatch(
         num_reqs: int,
         num_tokens: int,
         uniform_token_count: int | None,
+        invalid_modes: AbstractSet[CUDAGraphMode] | None = None,
     ) -> BatchExecutionDescriptor:
-        """Find matching cudagraph descriptor from priority-ordered candidates."""
+        """Find matching cudagraph descriptor from priority-ordered candidates.
+
+        invalid_modes optionally excludes runtime modes for the current batch.
+        """
         if self._graphs_captured and 0 < num_tokens < len(self._candidates):
             for desc in self._candidates[num_tokens]:
+                if invalid_modes and desc.cg_mode in invalid_modes:
+                    continue
                 if _is_compatible(desc, num_reqs, num_tokens, uniform_token_count):
                     return desc
         return BatchExecutionDescriptor(

vllm/v1/worker/gpu/dp_utils.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations
 
+from collections.abc import Set as AbstractSet
+
 import torch
 import torch.distributed as dist
 
@@ -21,6 +23,7 @@ def sync_cudagraph_and_dp_padding(
     uniform_token_count: int | None,
     dp_size: int,
     dp_rank: int,
+    invalid_modes: AbstractSet[CUDAGraphMode] | None = None,
 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
     """
     Coordinates the batch descriptor and DP padding across all ranks.
@@ -69,8 +72,16 @@ def sync_cudagraph_and_dp_padding(
 
     # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs
     # so we don't perform request padding for PIECEWISE graphs
+    # Preserve the DP-wide downgrade when redispatching with synced padding.
+    synced_invalid_modes = invalid_modes
+    if synced_cg_mode != CUDAGraphMode.FULL:
+        synced_invalid_modes = set(synced_invalid_modes or ())
+        synced_invalid_modes.add(CUDAGraphMode.FULL)
     synced_desc = cudagraph_manager.dispatch(
-        num_reqs, synced_num_tokens, synced_uniform_token_count
+        num_reqs,
+        synced_num_tokens,
+        synced_uniform_token_count,
+        invalid_modes=synced_invalid_modes,
     )
 
     # Update num_tokens_across_dp to reflect padded size.
@@ -87,6 +98,7 @@ def dispatch_cg_and_sync_dp(
     dp_size: int,
     dp_rank: int,
     need_eager: bool = False,
+    invalid_modes: AbstractSet[CUDAGraphMode] | None = None,
 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
     if need_eager:
         batch_desc = BatchExecutionDescriptor(
@@ -100,7 +112,10 @@ def dispatch_cg_and_sync_dp(
             "where need_eager must be True"
         )
         batch_desc = cudagraph_manager.dispatch(
-            num_reqs, num_tokens, uniform_token_count
+            num_reqs,
+            num_tokens,
+            uniform_token_count,
+            invalid_modes=invalid_modes,
         )
 
     if dp_size == 1:
@@ -114,4 +129,5 @@ def dispatch_cg_and_sync_dp(
         uniform_token_count,
         dp_size,
         dp_rank,
+        invalid_modes=invalid_modes,
     )

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

@@ -533,6 +533,12 @@ def propose(
             input_batch.num_tokens,
             max_query_len,
         )
+        # EAGLE draft prefill FULL CUDA graph replay does not update per-batch
+        # multimodal embeddings/masks. For live multimodal batches, disable
+        # only FULL and fall back to PIECEWISE/eager, where mm_inputs are
+        # passed through self.prefill(...). Do not use self.supports_mm_inputs:
+        # MM-capable models may still serve text-only batches.
+        prefill_invalid_modes = _get_prefill_invalid_cudagraph_modes(mm_inputs)
         prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
             self.prefill_cudagraph_manager,
             num_reqs,
@@ -541,6 +547,7 @@ def propose(
             dp_size=self.dp_size,
             dp_rank=self.dp_rank,
             need_eager=is_profile,
+            invalid_modes=prefill_invalid_modes,
         )
 
         if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
@@ -598,6 +605,18 @@ def propose(
         return self.draft_tokens[:num_reqs]
 
 
+def _has_live_mm_inputs(
+    mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,
+) -> bool:
+    return mm_inputs is not None and len(mm_inputs[0]) > 0
+
+
+def _get_prefill_invalid_cudagraph_modes(
+    mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,
+) -> set[CUDAGraphMode] | None:
+    return {CUDAGraphMode.FULL} if _has_live_mm_inputs(mm_inputs) else None
+
+
 @triton.jit
 def _prepare_eagle_inputs_kernel(
     last_token_indices_ptr,