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 8, 01:15 UTC
[perf] Avoid dtype promotion sync in mamba_get_block_table_tensor (#34870)
Signed-off-by: Huamin Li <3ericli@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Huamin Li committed
Feb 20, 2026 at 06:21 UTC
1fe462168c381f604a5ef9d491a230a3dd861d2c
1 file changed
+6
-2
vllm/v1/attention/backends/utils.py
+6
-2
@@ -855,8 +855,12 @@ def mamba_get_block_table_tensor(
855
(seq_lens - 1) // kv_cache_spec.block_size,
856
min=0,
857
)
858
+ # Use int32 for arithmetic to avoid dtype promotion overhead,
859
+ # then convert to int64 for gather (which requires Long indices)
860
offsets = torch.arange(
859
- 1 + kv_cache_spec.num_speculative_blocks, device=block_table.device
861
+ 1 + kv_cache_spec.num_speculative_blocks,
862
+ device=block_table.device,
863
+ dtype=torch.int32,
864
)
861
- indices_to_gather = start_indices.unsqueeze(1) + offsets
865
+ indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
866
return torch.gather(block_table, 1, indices_to_gather)