biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[perf] Avoid dtype promotion sync in mamba_get_block_table_tensor (#34870)

Browse Source 
Signed-off-by: Huamin Li <3ericli@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
This commit is contained in:
Huamin Li
2026-02-20 06:21:56 -08:00
committed by GitHub
parent ed31a020ee
commit 1fe462168c
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/v1/attention/backends/utils.py
View File

@@ -855,8 +855,12 @@ def mamba_get_block_table_tensor(
(seq_lens - 1) // kv_cache_spec.block_size,
min=0,
)
# Use int32 for arithmetic to avoid dtype promotion overhead,
# then convert to int64 for gather (which requires Long indices)
offsets = torch.arange(
1 + kv_cache_spec.num_speculative_blocks, device=block_table.device
1 + kv_cache_spec.num_speculative_blocks,
device=block_table.device,
dtype=torch.int32,
)
indices_to_gather = start_indices.unsqueeze(1) + offsets
indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
return torch.gather(block_table, 1, indices_to_gather)