From 6218034dd7f9a56596e4fd8c8c8fc1d8011ed9c2 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 15 Jan 2026 08:59:23 -0800 Subject: [PATCH] [Model Runner V2] Support FlashInfer backend & Fix CUDA Graph bug [1/2] (#32348) Signed-off-by: Woosuk Kwon --- vllm/v1/worker/gpu/cudagraph_utils.py | 14 +++++++++----- vllm/v1/worker/gpu/model_runner.py | 10 ++++++++-- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 51784cdc6..abcdb69e4 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -195,15 +195,19 @@ def get_cudagraph_size( cudagraph_sizes: dict[int, int], cudagraph_mode: CUDAGraphMode, ) -> int | None: + if not cudagraph_mode.has_full_cudagraphs(): + # No full CUDA graph is used. + return None + size = cudagraph_sizes.get(num_tokens_after_dp_padding) if size is None: # No CUDA graph for this size. return None - if cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY: - all_decode = all(x == 1 for x in num_tokens_per_request) - if not all_decode: - # Prefill is included. - return None + + is_mixed = any(x > 1 for x in num_tokens_per_request) + if is_mixed and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL: + # Prefill is included, and this mode doesn't use CUDA graph for it. + return None return size diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f603dc96d..6e3eaad4d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -230,8 +230,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin): ) # TODO(woosuk): Support other backends. - if not all(b.get_name() == "FLASH_ATTN" for b in self.attn_backends.values()): - raise NotImplementedError("Only FLASH_ATTN backend is supported currently.") + supported_backends = ("FLASH_ATTN", "FLASHINFER") + for backend in self.attn_backends.values(): + backend_name = backend.get_name() + if backend_name not in supported_backends: + raise NotImplementedError( + f"The {backend_name} attention backend is not supported yet. " + f"Supported backends are: {supported_backends}." + ) self.kv_caches: list[torch.Tensor] = [] init_kv_cache(