review comments

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

Author: LucasWilkinson

vllm/v1/attention/backends/flex_attention.py

@@ -3,7 +3,7 @@
 """Attention layer with FlashAttention."""
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
 import torch
 from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
@@ -23,9 +23,6 @@
 
 logger = init_logger(__name__)
 
-if TYPE_CHECKING:
-    pass
-
 create_block_mask_compiled = torch.compile(create_block_mask,
                                            fullgraph=True,
                                            mode="reduce-overhead")

vllm/v1/worker/gpu_model_runner.py

@@ -685,6 +685,9 @@ def _prepare_inputs(
         for kv_cache_group_id, kv_cache_group_spec in enumerate(
                 self.kv_cache_config.kv_cache_groups):
 
+            blk_table = self.input_batch.block_table[kv_cache_group_id]
+            blk_table_tensor = blk_table.get_device_tensor()[:num_reqs]
+            slot_mapping = blk_table.slot_mapping[:total_num_scheduled_tokens]
             common_attn_metadata = CommonAttentionMetadata(
                 query_start_loc=self.query_start_loc[:num_reqs + 1],
                 query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
@@ -695,10 +698,8 @@ def _prepare_inputs(
                 num_reqs=num_reqs,
                 num_actual_tokens=total_num_scheduled_tokens,
                 max_query_len=max_num_scheduled_tokens,
-                block_table_tensor=self.input_batch.
-                block_table[kv_cache_group_id].get_device_tensor()[:num_reqs],
-                slot_mapping=self.input_batch.block_table[kv_cache_group_id].
-                slot_mapping[:total_num_scheduled_tokens],
+                block_table_tensor=blk_table_tensor,
+                slot_mapping=slot_mapping,
             )
 
             if self.speculative_config and \