optimize

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

vllm/v1/attention/backends/flashinfer.py

@@ -18,6 +18,7 @@
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm.utils import cdiv
 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
 from vllm.v1.attention.backends.utils import (
     AttentionMetadataBuilder, CommonAttentionMetadata, PerLayerParameters,
@@ -241,6 +242,12 @@ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
         self.vllm_config = vllm_config
         self.cache_config = vllm_config.cache_config
         self.kv_cache_spec = kv_cache_spec
+        max_num_blocks_per_request = cdiv(
+            vllm_config.model_config.max_model_len,
+            self.kv_cache_spec.block_size)
+        self.block_table_arange = torch.arange(max_num_blocks_per_request,
+                                               dtype=torch.int32,
+                                               device=self.device)
 
     def reorder_batch(self, input_batch: InputBatch,
                       scheduler_output: SchedulerOutput) -> bool:
@@ -432,19 +439,19 @@ def build(self,
             shared_kv_page_indices_cpu = None
             shared_kv_last_page_len_cpu = None
 
-        # Build CPU versions directly from CPU data
-        # paged_kv_indices_cpu: extract from block_table on CPU
-        mask_cpu = (torch.arange(block_table_tensor.size(1),
-                                 dtype=torch.int32,
-                                 device='cpu').unsqueeze(0)
-                    < block_table_bounds_cpu.unsqueeze(1))
-        paged_kv_indices = block_table_tensor[mask_cpu]
+        max_num_blocks = block_table_bounds_cpu.max()
+        block_table_bounds = block_table_bounds_cpu.to(self.device,
+                                                       non_blocking=True)
+        mask = (self.block_table_arange[:max_num_blocks].unsqueeze(0)
+                < block_table_bounds.unsqueeze(1))
+        paged_kv_indices = block_table_tensor[:, :max_num_blocks][mask]
 
         # paged_kv_indptr_cpu: cumulative sum of block_table_bounds_cpu
-        paged_kv_indptr_cpu = torch.cat([
-            torch.zeros(1, dtype=torch.int32, device='cpu'),
-            block_table_bounds_cpu.cumsum(dim=0, dtype=torch.int32)
-        ])
+        paged_kv_indptr_cpu = torch.zeros(len(block_table_bounds_cpu) + 1,
+                                          dtype=torch.int32,
+                                          device='cpu')
+        paged_kv_indptr_cpu[1:] = block_table_bounds_cpu.cumsum(
+            dim=0, dtype=torch.int32)
 
         # paged_kv_last_page_len_cpu: from seq_lens_cpu
         paged_kv_last_page_len_cpu = seq_lens_cpu % page_size