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 7, 11:30 UTC
[Model Runner V2] Use block table apis for capture inputs (#35671)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 1, 2026 at 10:31 UTC
72f4d162623854786d29e1d9c6e232cfdf81d3cc
2 files changed
+13
-2
vllm/v1/worker/gpu/block_table.py
+11
@@ -119,6 +119,10 @@ class BlockTables:
119
return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
120
121
def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
122
+ # NOTE(woosuk): The output may be used for CUDA graph capture.
123
+ # Therefore, this method must return the persistent tensor
124
+ # with the same memory address as that used during the model's forward pass,
125
+ # rather than allocating a new tensor.
126
return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
127
128
def compute_slot_mappings(
@@ -150,7 +154,14 @@ class BlockTables:
154
return self.slot_mappings[:, :num_tokens]
155
156
def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
157
+ # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.
158
+ # This is because the padding logic is complex and kernels may access beyond
159
+ # the requested range.
160
self.slot_mappings.fill_(PAD_SLOT_ID)
161
+ # NOTE(woosuk): The output may be used for CUDA graph capture.
162
+ # Therefore, this method must return the persistent tensor
163
+ # with the same memory address as that used during the model's forward pass,
164
+ # rather than allocating a new tensor.
165
return self.slot_mappings[:, :num_tokens]
166
167
vllm/v1/worker/gpu/cudagraph_utils.py
+2
-2
@@ -420,8 +420,8 @@ def prepare_inputs_to_capture(
420
input_buffers.dcp_local_seq_lens[:num_reqs] = num_tokens
421
input_buffers.dcp_local_seq_lens[num_reqs:] = 0
422
423
- input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
424
- slot_mappings = block_tables.slot_mappings[:, :num_tokens]
423
+ input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
424
+ slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
425
slot_mappings_by_layer = build_slot_mappings_by_layer(
426
slot_mappings, kv_cache_config
427
)