A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 02:45 UTC
[Model Runner V2] Support FlashInfer backend & Fix CUDA Graph bug [1/2] (#32348)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 15, 2026 at 08:59 UTC
6218034dd7f9a56596e4fd8c8c8fc1d8011ed9c2
2 files changed
+17
-7
vllm/v1/worker/gpu/cudagraph_utils.py
+9
-5
@@ -195,15 +195,19 @@ def get_cudagraph_size(
195
cudagraph_sizes: dict[int, int],
196
cudagraph_mode: CUDAGraphMode,
197
) -> int | None:
198
+ if not cudagraph_mode.has_full_cudagraphs():
199
+ # No full CUDA graph is used.
200
+ return None
201
+
202
size = cudagraph_sizes.get(num_tokens_after_dp_padding)
203
if size is None:
204
# No CUDA graph for this size.
205
return None
202
- if cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY:
203
- all_decode = all(x == 1 for x in num_tokens_per_request)
204
- if not all_decode:
205
- # Prefill is included.
206
- return None
206
+
207
+ is_mixed = any(x > 1 for x in num_tokens_per_request)
208
+ if is_mixed and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL:
209
+ # Prefill is included, and this mode doesn't use CUDA graph for it.
210
+ return None
211
return size
212
213
vllm/v1/worker/gpu/model_runner.py
+8
-2
@@ -230,8 +230,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
230
)
231
232
# TODO(woosuk): Support other backends.
233
- if not all(b.get_name() == "FLASH_ATTN" for b in self.attn_backends.values()):
234
- raise NotImplementedError("Only FLASH_ATTN backend is supported currently.")
233
+ supported_backends = ("FLASH_ATTN", "FLASHINFER")
234
+ for backend in self.attn_backends.values():
235
+ backend_name = backend.get_name()
236
+ if backend_name not in supported_backends:
237
+ raise NotImplementedError(
238
+ f"The {backend_name} attention backend is not supported yet. "
239
+ f"Supported backends are: {supported_backends}."
240
+ )
241
242
self.kv_caches: list[torch.Tensor] = []
243
init_kv_cache(