From 1fe462168c381f604a5ef9d491a230a3dd861d2c Mon Sep 17 00:00:00 2001 From: Huamin Li <3ericli@gmail.com> Date: Fri, 20 Feb 2026 06:21:56 -0800 Subject: [PATCH] [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 --- vllm/v1/attention/backends/utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index eda50155d..1b030eaf1 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -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)