flashinfer

vllm/attention/backends/blocksparse_attn.py

@@ -102,6 +102,10 @@ def get_impl_cls() -> Type["BlocksparseFlashAttentionImpl"]:
     def get_metadata_cls() -> Type["AttentionMetadata"]:
         return BlocksparseFlashAttentionMetadata
 
+    @staticmethod
+    def get_builder_cls() -> Type["BlocksparseFlashAttentionMetadataBuilder"]:
+        return BlocksparseFlashAttentionMetadataBuilder
+
     @staticmethod
     def get_kv_cache_shape(
         num_blocks: int,

vllm/attention/backends/flash_attn.py

@@ -11,7 +11,8 @@
                                               AttentionMetadataBuilder,
                                               AttentionType)
 from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
-                                           compute_slot_mapping_start_idx)
+                                           compute_slot_mapping_start_idx,
+                                           is_block_tables_empty)
 from vllm.sequence import SequenceGroupMetadata
 from vllm.utils import make_tensor_with_pad
 
@@ -258,12 +259,14 @@ def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata,
             self.block_tables.append(block_table)
 
             # 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.sliding_window,
                 self.use_v2_block_manager)
-            compute_slot_mapping(self.slot_mapping, seq_id, seq_len,
-                                 context_len, start_idx, self.block_size,
-                                 block_tables)
+            compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
+                                 seq_len, context_len, start_idx,
+                                 self.block_size,
+                                 seq_group_metadata.block_tables)
 
     def build(self, runner: "GPUModelRunnerBase", seq_lens, query_lens,
               use_captured_graph: bool, cuda_graph_pad_size: int,

vllm/attention/backends/flashinfer.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Set, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Type
 
 try:
     from flashinfer import BatchDecodeWithPagedKVCacheWrapper
@@ -14,7 +14,18 @@
 
 from vllm import _custom_ops as ops
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
-                                              AttentionMetadata, AttentionType)
+                                              AttentionMetadata,
+                                              AttentionMetadataBuilder,
+                                              AttentionType)
+from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
+                                           compute_slot_mapping_start_idx,
+                                           is_block_tables_empty)
+from vllm.sequence import SequenceGroupMetadata
+from vllm.utils import get_kv_cache_torch_dtype, make_tensor_with_pad
+
+if TYPE_CHECKING:
+    from vllm.worker.model_runner import (GPUModelRunnerBase,
+                                          ModelInputForGPUBuilder)
 
 
 class FlashInferBackend(AttentionBackend):
@@ -31,6 +42,10 @@ def get_impl_cls() -> Type["FlashInferImpl"]:
     def get_metadata_cls() -> Type["AttentionMetadata"]:
         return FlashInferMetadata
 
+    @staticmethod
+    def get_builder_cls() -> Type["FlashInferMetadataBuilder"]:
+        return FlashInferMetadataBuilder
+
     @staticmethod
     def get_kv_cache_shape(
         num_blocks: int,
@@ -188,6 +203,216 @@ def decode_metadata(self) -> Optional["FlashInferMetadata"]:
         return self
 
 
+class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
+
+    def __init__(self, input_builder: "ModelInputForGPUBuilder"):
+        self.slot_mapping: List[int] = []
+        self.prefill_seq_lens: List[int] = []
+        self.context_lens: List[int] = []
+        self.block_tables: List[List[int]] = []
+        self.decode_seq_lens: List[int] = []
+        self.num_prefills = 0
+        self.num_prefill_tokens = 0
+        self.num_decode_tokens = 0
+
+        self.sliding_window = input_builder.sliding_window
+        self.block_size = input_builder.block_size
+        self.use_v2_block_manager = (
+            input_builder.scheduler_config.use_v2_block_manager)
+
+        # Please follow https://docs.flashinfer.ai/tutorials/kv_layout.html#page-layout
+        # for the precise definition of the following fields.
+        # An example:
+        # request 1, page indices [0, 5, 8]
+        # request 2, page indices [1, 6, 7]
+        # request 3, page indices [3, 4]
+        # paged_kv_indices is a concatenation of page indices of all requests:
+        # [0, 5, 8, 1, 6, 7, 3, 4]
+        # paged_kv_indptr is used to index into paged_kv_indices:
+        # [0, 3, 6, 8]
+        self.paged_kv_indices: List[int] = []
+        # 0 at the beginning of paged_kv_indptr indicates the start of the
+        # first request’s page indices in the paged_kv_indices list.
+        self.paged_kv_indptr: List[int] = [0]
+        # paged_kv_last_page_len is the length of the last page of each request
+        self.paged_kv_last_page_len: List[int] = []
+
+    def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata,
+                      token_lens: List[int], seq_lens: List[int],
+                      sliding_seq_lens: List[int], query_lens: List[int],
+                      context_lens: List[int],
+                      curr_sliding_window_blocks: List[int], prefix_cache_hit,
+                      chunked_prefill_enabled):
+        is_prompt = seq_group_metadata.is_prompt
+        block_tables = seq_group_metadata.block_tables
+        computed_block_nums = seq_group_metadata.computed_block_nums
+
+        for (seq_id, token_len, seq_len, sliding_seq_len, query_len,
+             context_len, curr_sliding_window_block) in zip(
+                 seq_group_metadata.seq_data.keys(), token_lens, seq_lens,
+                 sliding_seq_lens, query_lens, context_lens,
+                 curr_sliding_window_blocks):
+            self.context_lens.append(context_len)
+            if is_prompt:
+                self.num_prefills += 1
+                self.num_prefill_tokens += token_len
+                self.prefill_seq_lens.append(seq_len)
+            else:
+                assert query_len == 1, (
+                    "seq_len: {}, context_len: {}, query_len: {}".format(
+                        seq_len, context_len, query_len))
+                self.num_decode_tokens += query_len
+                self.decode_seq_lens.append(sliding_seq_len)
+
+            # Compute block table.
+            # TODO(sang): Combine chunked prefill and prefix caching by
+            # only allowing multiple of block_size chunk size.
+            # NOTE: This only works for oooooooxxx style attention.
+            block_table = []
+            if prefix_cache_hit:
+                block_table = computed_block_nums
+            elif ((chunked_prefill_enabled or not is_prompt)
+                  and block_tables is not None):
+                block_table = block_tables[seq_id][-curr_sliding_window_block:]
+            self.block_tables.append(block_table)
+
+            is_profile_run = is_block_tables_empty(block_tables)
+
+            # Compute slot mapping.
+            start_idx = compute_slot_mapping_start_idx(
+                is_prompt, query_len, context_len, self.sliding_window,
+                self.use_v2_block_manager)
+            compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
+                                 seq_len, context_len, start_idx,
+                                 self.block_size,
+                                 seq_group_metadata.block_tables)
+            if is_profile_run:
+                return
+
+            # Get the number of valid blocks based on sequence length.
+            # If seq_len = 16, block_size = 16,
+            # block_table_bound is 1 with 1 valid block.
+            # If seq_len = 15, block_size = 16,
+            # block_table_bound is 0 + 1 with 1 valid block.
+            block_table_bound = seq_len // self.block_size + 1 \
+                                if seq_len % self.block_size != 0 \
+                                else seq_len // self.block_size
+            block_table = block_tables[seq_id]
+            self.paged_kv_indices.extend(block_table[:block_table_bound])
+            self.paged_kv_indptr.append(self.paged_kv_indptr[-1] +
+                                        block_table_bound)
+
+            last_page_len = seq_len % self.block_size
+            if last_page_len == 0:
+                last_page_len = self.block_size
+            self.paged_kv_last_page_len.append(last_page_len)
+
+    def build(self, runner: "GPUModelRunnerBase", seq_lens, query_lens,
+              use_captured_graph: bool, cuda_graph_pad_size: int,
+              batch_size: int):
+        device = runner.device
+
+        max_query_len = max(query_lens)
+        max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
+        num_decode_tokens = self.num_decode_tokens
+
+        if use_captured_graph:
+            self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
+            self.block_tables.extend([] * cuda_graph_pad_size)
+            num_decode_tokens = batch_size + cuda_graph_pad_size
+
+            # The shape of graph_block_tables is
+            # [max batch size, max context len // block size].
+            input_block_tables = runner.graph_block_tables[:batch_size]
+            for i, block_table in enumerate(self.block_tables):
+                if block_table:
+                    input_block_tables[i, :len(block_table)] = block_table
+            block_tables = torch.tensor(input_block_tables, device=device)
+
+            last_paged_kv_indptr = self.paged_kv_indptr[-1]
+            self.paged_kv_indptr.extend([last_paged_kv_indptr] *
+                                        cuda_graph_pad_size)
+            self.paged_kv_last_page_len.extend([0] * cuda_graph_pad_size)
+        else:
+            max_block_table_len = max(
+                len(block_table) for block_table in self.block_tables)
+            block_tables = make_tensor_with_pad(
+                self.block_tables,
+                max_len=max_block_table_len,
+                pad=0,
+                dtype=torch.int,
+                device=device,
+            )
+        assert max_query_len > 0, ("query_lens: {}".format(query_lens))
+
+        seq_lens_tensor = torch.tensor(seq_lens,
+                                       dtype=torch.int,
+                                       device=device)
+        query_lens_tensor = torch.tensor(query_lens,
+                                         dtype=torch.long,
+                                         device=device)
+        query_start_loc = torch.zeros(query_lens_tensor.shape[0] + 1,
+                                      dtype=torch.int32,
+                                      device=device)
+        seq_start_loc = torch.zeros(seq_lens_tensor.shape[0] + 1,
+                                    dtype=torch.int32,
+                                    device=device)
+        torch.cumsum(seq_lens_tensor,
+                     dim=0,
+                     dtype=seq_start_loc.dtype,
+                     out=seq_start_loc[1:])
+        torch.cumsum(query_lens_tensor,
+                     dim=0,
+                     dtype=query_start_loc.dtype,
+                     out=query_start_loc[1:])
+
+        slot_mapping_tensor = torch.tensor(self.slot_mapping,
+                                           dtype=torch.long,
+                                           device=device)
+
+        logits_soft_cap = getattr(runner.model_config.hf_config,
+                                  "attn_logit_softcapping", None)
+
+        if len(self.paged_kv_indptr) > 0:
+            paged_kv_indices_tensor = torch.tensor(self.paged_kv_indices,
+                                                   device="cpu",
+                                                   dtype=torch.int)
+            paged_kv_indptr_tensor = torch.tensor(self.paged_kv_indptr,
+                                                  device="cpu",
+                                                  dtype=torch.int)
+            paged_kv_last_page_len_tensor = torch.tensor(
+                self.paged_kv_last_page_len, device="cpu", dtype=torch.int)
+        else:
+            paged_kv_indices_tensor = None
+            paged_kv_indptr_tensor = None
+            paged_kv_last_page_len_tensor = None
+
+        kv_cache_dtype = get_kv_cache_torch_dtype(runner.kv_cache_dtype,
+                                                  runner.model_config.dtype)
+        return FlashInferMetadata(
+            num_prefills=self.num_prefills,
+            slot_mapping=slot_mapping_tensor,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=num_decode_tokens,
+            max_prefill_seq_len=max_prefill_seq_len,
+            block_tables=block_tables,
+            paged_kv_indptr=paged_kv_indptr_tensor,
+            paged_kv_indices=paged_kv_indices_tensor,
+            paged_kv_last_page_len=paged_kv_last_page_len_tensor,
+            num_qo_heads=runner.model_config.get_num_attention_heads(
+                runner.parallel_config),
+            num_kv_heads=runner.model_config.get_num_kv_heads(
+                runner.parallel_config),
+            head_dim=runner.model_config.get_head_size(),
+            page_size=self.block_size,
+            seq_start_loc=seq_start_loc,
+            query_start_loc=query_start_loc,
+            device=device,
+            data_type=kv_cache_dtype,
+            use_cuda_graph=use_captured_graph,
+            logits_soft_cap=logits_soft_cap)
+
+
 class FlashInferImpl(AttentionImpl):
 
     def __init__(

vllm/attention/backends/rocm_flash_attn.py

@@ -37,6 +37,10 @@ def get_impl_cls() -> Type["ROCmFlashAttentionImpl"]:
     def get_metadata_cls() -> Type["AttentionMetadata"]:
         return ROCmFlashAttentionMetadata
 
+    @staticmethod
+    def get_builder_cls() -> Type["ROCmFlashAttentionMetadataBuilder"]:
+        return ROCmFlashAttentionMetadataBuilder
+
     @staticmethod
     def get_kv_cache_shape(
         num_blocks: int,

vllm/attention/backends/utils.py

@@ -44,10 +44,9 @@ def compute_slot_mapping_start_idx(is_prompt, query_len, context_len,
     return start_idx
 
 
-def compute_slot_mapping(slot_mapping, seq_id, seq_len, context_len, start_idx,
-                         block_size, block_tables):
+def compute_slot_mapping(is_profile_run, slot_mapping, seq_id, seq_len,
+                         context_len, start_idx, block_size, block_tables):
     """TBA."""
-    is_profile_run = is_block_tables_empty(block_tables)
     if is_profile_run:
         # During memory profiling, the block tables are not
         # initialized yet. In this case, we just use a dummy
@@ -112,11 +111,13 @@ def metadata_builder_add_seq_group_common(
         attn_metadata_builder.block_tables.append(block_table)
 
         # 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,
             attn_metadata_builder.sliding_window,
             attn_metadata_builder.use_v2_block_manager)
-        compute_slot_mapping(attn_metadata_builder.slot_mapping, seq_id,
+        compute_slot_mapping(is_profile_run,
+                             attn_metadata_builder.slot_mapping, seq_id,
                              seq_len, context_len, start_idx,
                              attn_metadata_builder.block_size,
                              seq_group_metadata.block_tables)