[Core] Cache policy framework

vllm/attention/backends/flash_attn.py

@@ -400,6 +400,7 @@ def _add_seq_group(
         """
         is_prompt = inter_data.is_prompt
         block_tables = inter_data.block_tables
+        slot_mappings = inter_data.slot_mappings
 
         for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
              curr_sliding_window_block) in zip(
@@ -444,11 +445,11 @@ def _add_seq_group(
             # Compute slot mapping.
             is_profile_run = is_block_tables_empty(block_tables)
             start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
-                                                       context_len,
+                                                       self.block_size,
                                                        self.sliding_window)
             compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
                                  seq_len, context_len, start_idx,
-                                 self.block_size, inter_data.block_tables)
+                                 self.block_size, slot_mappings)
 
     def _get_graph_runner_block_tables(
             self, num_seqs: int,

vllm/attention/backends/flashinfer.py

@@ -515,6 +515,7 @@ def _add_seq_group(
         """
         is_prompt = inter_data.is_prompt
         block_tables = inter_data.block_tables
+        slot_mappings = inter_data.slot_mappings
         computed_block_nums = inter_data.computed_block_nums
 
         for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
@@ -556,11 +557,11 @@ def _add_seq_group(
 
             # Compute slot mapping.
             start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
-                                                       context_len,
+                                                       self.block_size,
                                                        self.sliding_window)
             compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
                                  seq_len, context_len, start_idx,
-                                 self.block_size, inter_data.block_tables)
+                                 self.block_size, slot_mappings)
 
             # It is not necessary to add paged_kv_indices, paged_kv_indptr,
             # and paged_kv_last_page_len for profile run because we will

vllm/attention/backends/utils.py

@@ -4,7 +4,6 @@
 from itertools import accumulate
 from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, TypeVar, Union
 
-import numpy as np
 import torch
 
 from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder,
@@ -42,43 +41,21 @@ def is_block_tables_empty(block_tables: Union[None, Dict]):
 
 
 def compute_slot_mapping_start_idx(is_prompt: bool, query_len: int,
-                                   context_len: int, sliding_window: int):
+                                   block_size: int, sliding_window: int):
     """
     Compute the start index of slot mapping.
     """
     start_idx = 0
     if is_prompt and sliding_window is not None:
-        start_idx = max(0, query_len - sliding_window)
+        num_blocks = (sliding_window + block_size - 1) // block_size
+        start_idx = max(0, query_len - num_blocks * block_size)
     return start_idx
 
 
-def _compute_slot_mapping_python(slot_mapping: List[int],
-                                 block_table: List[int], range_start: int,
-                                 range_end: int, block_size: int):
-    for i in range(range_start, range_end):
-        block_number = block_table[i // block_size]
-        block_offset = i % block_size
-        slot = block_number * block_size + block_offset
-        slot_mapping.append(slot)
-
-
-def _compute_slot_mapping_numpy(slot_mapping: List[int],
-                                block_table: List[int], range_start: int,
-                                range_end: int, block_size: int):
-    block_table_array = np.array(block_table)
-    idx = np.arange(range_start, range_end)
-    block_offset = idx % block_size
-    idx //= block_size
-    seq_slot_mapping_array = block_table_array[idx]
-    seq_slot_mapping_array *= block_size
-    seq_slot_mapping_array += block_offset
-    slot_mapping.extend(seq_slot_mapping_array)
-
-
 def compute_slot_mapping(is_profile_run: bool, slot_mapping: List[int],
                          seq_id: int, seq_len: int, context_len: int,
                          start_idx: int, block_size: int,
-                         block_tables: Dict[int, List[int]]):
+                         slot_mappings: Dict[int, List[int]]):
     """
     Compute slot mapping.
     """
@@ -95,23 +72,14 @@ def compute_slot_mapping(is_profile_run: bool, slot_mapping: List[int],
     # sliding_window). For example, if the prompt len is 10,
     # sliding window is 8, and block size is 4, the first two
     # tokens are masked and the slot mapping will be
-    # [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].
+    # [-1, -1, 2, 3, 4, 5, 6, 7, 8, 9].
     padding_mask_len = max(0, start_idx - context_len)
     slot_mapping.extend([PAD_SLOT_ID] * padding_mask_len)
 
     range_start = max(start_idx, context_len)
     range_end = seq_len
-    numel = range_end - range_start
-    block_table = block_tables[seq_id]
-
-    # numpy implementation will be faster than python if we have
-    # many elements, otherwise it will be slower.
-    if numel < _COMPUTE_SLOT_MAPPING_NUMPY_NUMEL:
-        _compute_slot_mapping_python(slot_mapping, block_table, range_start,
-                                     range_end, block_size)
-    else:
-        _compute_slot_mapping_numpy(slot_mapping, block_table, range_start,
-                                    range_end, block_size)
+    seq_slot_mapping = slot_mappings[seq_id]
+    slot_mapping.extend(seq_slot_mapping[range_start:range_end])
 
 
 TAttentionMetadata = TypeVar("TAttentionMetadata", bound='AttentionMetadata')
@@ -145,6 +113,7 @@ def _add_seq_group(
             chunked_prefill_enabled: bool):
         is_prompt = inter_data.is_prompt
         block_tables = inter_data.block_tables
+        slot_mappings = inter_data.slot_mappings
 
         for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
              curr_sliding_window_block) in zip(
@@ -189,11 +158,11 @@ def _add_seq_group(
             # Compute slot mapping.
             is_profile_run = is_block_tables_empty(block_tables)
             start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
-                                                       context_len,
+                                                       self.block_size,
                                                        self.sliding_window)
             compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
                                  seq_len, context_len, start_idx,
-                                 self.block_size, inter_data.block_tables)
+                                 self.block_size, slot_mappings)
 
     def build(self, seq_lens: List[int], query_lens: List[int],
               cuda_graph_pad_size: int, batch_size: int):