xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-03 06:15:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] multi-step + flashinfer: ensure cuda graph compatible (#8427)

Browse Source
This commit is contained in:
Alexander Matveev 2024-09-12 17:01:50 -04:00 committed by GitHub
parent 551ce01078
commit 019877253b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
vllm/attention/backends/flashinfer.py
View File

@ -597,9 +597,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
# The shape of graph_block_tables is
# [max batch size, max context len // block size].
input_block_tables = self.runner.graph_block_tables[:batch_size]
max_blocks = input_block_tables.shape[1]
for i, block_table in enumerate(self.block_tables):
if block_table:
input_block_tables[i, :len(block_table)] = block_table
num_blocks = len(block_table)
if num_blocks <= max_blocks:
input_block_tables[i, :num_blocks] = block_table
else:
# It may be possible to have more blocks allocated due
# to lookahead slots of multi-step, however, they are
# not used anyway, so can be safely ignored.
input_block_tables[
i, :max_blocks] = block_table[:max_blocks]
block_tables = torch.from_numpy(input_block_tables).to(
device, non_blocking=True)