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 9, 06:15 UTC
[BugFix] Fix whisper FA2 + full cudagraphs (#33360)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Lucas Wilkinson committed
Jan 30, 2026 at 21:15 UTC
0a3c71e7e5f06507f64677f77d5fca84d9c7cc65
2 files changed
+12
-12
vllm/v1/attention/backends/flash_attn.py
-12
@@ -263,18 +263,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
263
vllm_config: "VllmConfig",
264
kv_cache_spec: "AttentionSpec",
265
) -> AttentionCGSupport:
266
- # FA2 does not support CUDA graphs with encoder-decoder models due to
267
- # accuracy issues reported in https://github.com/vllm-project/vllm/issues/33091
268
- if (
269
- vllm_config.model_config.is_encoder_decoder
270
- and get_flash_attn_version() == 2
271
- ):
272
- logger.warning_once(
273
- "FlashAttention2 does not support CUDA graphs with "
274
- "encoder-decoder models due to accuracy issues reported in #33091. "
275
- "Disabling CUDA graph."
276
- )
277
- return AttentionCGSupport.NEVER
266
return cls._cudagraph_support
267
268
def __init__(
vllm/v1/worker/gpu_model_runner.py
+12
@@ -1395,12 +1395,14 @@ class GPUModelRunner(
1395
num_scheduled_tokens: dict[str, int],
1396
kv_cache_spec: KVCacheSpec,
1397
num_reqs: int,
1398
+ for_cudagraph_capture: bool = False,
1399
) -> tuple[torch.Tensor | None, np.ndarray | None]:
1400
if not isinstance(kv_cache_spec, CrossAttentionSpec):
1401
return None, None
1402
1403
# Zero out buffer for padding requests that are not actually scheduled (CGs)
1404
self.encoder_seq_lens.np[:num_reqs] = 0
1405
+
1406
# Build encoder_seq_lens array mapping request indices to
1407
# encoder lengths for inputs scheduled in this batch
1408
for req_id in num_scheduled_tokens:
@@ -1417,6 +1419,15 @@ class GPUModelRunner(
1419
feature.mm_position.length for feature in req_state.mm_features
1420
)
1421
self.encoder_seq_lens.np[req_index] = encoder_input_tokens
1422
+ if for_cudagraph_capture:
1423
+ # During CUDA graph capture, we need to use realistic encoder lengths
1424
+ # so that max_seqlen_k is captured with the correct value.
1425
+ max_encoder_len = getattr(
1426
+ self.model_config.hf_config,
1427
+ "max_source_positions",
1428
+ self.max_encoder_len,
1429
+ )
1430
+ self.encoder_seq_lens.np[:num_reqs] = max_encoder_len
1431
1432
self.encoder_seq_lens.copy_to_gpu(num_reqs)
1433
encoder_seq_lens = self.encoder_seq_lens.gpu[:num_reqs]
@@ -1834,6 +1845,7 @@ class GPUModelRunner(
1845
num_scheduled_tokens or {},
1846
kv_cache_group.kv_cache_spec,
1847
num_reqs_padded,
1848
+ for_cudagraph_capture=for_cudagraph_capture,
1849
)
1850
if kv_cache_gid > 0:
1851
cm.block_table_tensor = _get_block_table(kv_cache_gid)