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 15, 01:00 UTC
[BugFix] Fix potential cuda-graph IMA (#21196)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jul 19, 2025 at 05:18 UTC
59f935300c4818cb10db8a0efadb431a2f169506
2 files changed
+6
-6
vllm/v1/attention/backends/utils.py
-5
@@ -59,11 +59,6 @@ class CommonAttentionMetadata:
59
block_table_tensor: torch.Tensor
60
slot_mapping: torch.Tensor
61
62
- def __post_init__(self):
63
- # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
64
- # mode.
65
- self.slot_mapping[self.num_actual_tokens:].fill_(-1)
66
-
62
63
M = TypeVar("M")
64
vllm/v1/worker/gpu_model_runner.py
+6
-1
@@ -684,7 +684,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
684
self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
685
non_blocking=True)
686
687
- # Fill unused with -1. Needed for reshape_and_cache
687
+ # Fill unused with 0 for full cuda graph mode.
688
self.seq_lens[num_reqs:].fill_(0)
689
# Note: pad query_start_loc to be non-decreasing, as kernels
690
# like FlashAttention requires that
@@ -704,6 +704,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
704
blk_table = self.input_batch.block_table[kv_cache_group_id]
705
blk_table_tensor = blk_table.get_device_tensor()[:num_reqs]
706
slot_mapping = blk_table.slot_mapping[:total_num_scheduled_tokens]
707
+
708
+ # Fill unused with -1. Needed for reshape_and_cache in full cuda
709
+ # graph mode.
710
+ blk_table.slot_mapping[total_num_scheduled_tokens:].fill_(-1)
711
+
712
common_attn_metadata = CommonAttentionMetadata(
713
query_start_loc=self.query_start_loc[:num_reqs + 1],
714
query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],