From fea4778d3b97a22afb4b93f2067565cc1e1b74a7 Mon Sep 17 00:00:00 2001 From: jadewang21 Date: Fri, 11 Apr 2025 12:57:17 +0800 Subject: [PATCH] Fix append_paged_kv_cache call in FlashinferAttentionWrapper by adding batch_indices and positions Signed-off-by: jadewang21 --- .../flashinfer_attention_wrapper copy.py | 258 ++++++++++++++++++ .../attention/flashinfer_attention_wrapper.py | 66 ++--- 2 files changed, 286 insertions(+), 38 deletions(-) create mode 100644 sarathi/model_executor/attention/flashinfer_attention_wrapper copy.py diff --git a/sarathi/model_executor/attention/flashinfer_attention_wrapper copy.py b/sarathi/model_executor/attention/flashinfer_attention_wrapper copy.py new file mode 100644 index 0000000..80ae842 --- /dev/null +++ b/sarathi/model_executor/attention/flashinfer_attention_wrapper copy.py @@ -0,0 +1,258 @@ +from typing import List, Optional + +import torch +from flashinfer import BatchPrefillWithPagedKVCacheWrapper, append_paged_kv_cache + +from sarathi.config import ModelConfig, ParallelConfig +from sarathi.core.datatypes.sequence import SequenceMetadata +from sarathi.metrics.constants import OperationMetrics +from sarathi.model_executor.attention.base_attention_wrapper import BaseAttentionWrapper + + +class FlashinferAttentionWrapper(BaseAttentionWrapper): + _inst = None + + def init( + self, + model_config: ModelConfig, + parallel_config: ParallelConfig, + block_size: int, + device: torch.device, + ): + super().init(model_config, parallel_config, block_size, device) + + prefill_workspace_buffer = torch.empty( + 128 * 1024 * 1024, dtype=torch.uint8, device=device + ) + self.prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( + prefill_workspace_buffer, "NHD" + ) + + decode_workspace_buffer = torch.empty( + 128 * 1024 * 1024, dtype=torch.uint8, device=device + ) + self.decode_wrapper = BatchPrefillWithPagedKVCacheWrapper( + decode_workspace_buffer, "NHD" + ) + + self.is_metadata_initialized = False + self.is_profiling_iteration = False + self.contains_prefill = False + self.contains_decode = False + self.num_prefill_tokens = 0 + self.num_total_tokens = 0 + + self.append_qo_indptr_tensor = None + self.append_kv_page_indices_tensor = None + self.append_kv_page_indptr_tensor = None + self.append_kv_last_page_len_tensor = None + + def to_int_tensor(self, data: List[int]) -> torch.Tensor: + return torch.tensor(data, dtype=torch.int32, device="cuda") + + def get_cache_block(self, num_blocks: int, **kwargs) -> torch.Tensor: + return torch.randn( + num_blocks, + 2, + self.block_size, + self.num_kv_heads, + self.head_dim, + **kwargs, + ) + + def begin_forward( + self, + seq_metadata_list: List[SequenceMetadata], + ) -> None: + # The indptr tensor captures the location query tokens in the input tensor. + # |<---------------------- num_valid_tokens ----------------------------------------------------->| + # |<--------------- num_prompt_tokens -------------->||<------- num_generation_tokens (M) ------->| + # |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->||<--generation_0-->|...|<--generation_M-1-->|<--padding-->| + # + # Flashinfer calls this layout as a raggedtensor. The indptr tensor captures the start of each + # sequence in the ragged tensor. The length of the indptr tensor is the number of sequences + 1. + # We perform both prefill and decode attention in a single call to batched prefill kernel. + # prefill_qo_indptr: [0, prompt_0, prompt_0 + prompt_1, ..., prompt_0 + ... + prompt_N-1, generation_0, generation_0 + 1, ..., generation_0 + ... + M] + prefill_qo_indptr: List[int] = [0] + decode_qo_indptr: List[int] = [0] + # The kv_page_indices tensor captures the pages of the key-value cache that + # are assigned to each token in the input tensor. Since there is a variable number + # of pages assigned to each sequence, a ragged tensor to represent this. + prefill_kv_page_indices: List[int] = [] + decode_kv_page_indices: List[int] = [] + # the last page might not be full, so we need to keep track of the length of the last page + prefill_kv_last_page_len: List[int] = [] + decode_kv_last_page_len: List[int] = [] + # Since the prefill_kv_page_indices tensor is a ragged tensor, we also need to keep track of the + # indptr tensor for the prefill_kv_page_indices tensor. This tensor captures the start of each sequence + # in the ragged tensor. + prefill_kv_page_indptr: List[int] = [0] + decode_kv_page_indptr: List[int] = [0] + + self.is_profiling_iteration = False + self.is_metadata_initialized = True + + self.contains_prefill = False + self.contains_decode = False + + for seq_metadata in seq_metadata_list: + if not seq_metadata.is_prompt: + continue + + # ONLY used for profiling + if seq_metadata.block_table is None: + self.is_profiling_iteration = True + # During memory profiling, the block tables are not initialized yet. + # We will just skip the attention computation for now. + return + + self.contains_prefill = True + + prompt_chunk_len = seq_metadata.prompt_chunk_len + processed_prompt_len = seq_metadata.seq.get_num_prompt_tokens_processed() + current_total_len = processed_prompt_len + prompt_chunk_len + + # indptr for the prompt tokens in q/o tensor + prefill_qo_indptr.append(prefill_qo_indptr[-1] + prompt_chunk_len) + # Compute the kv page indices for the prompt tokens. + num_blocks_in_use = ( + current_total_len + self.block_size - 1 + ) // self.block_size + prefill_kv_page_indices.extend(seq_metadata.block_table[:num_blocks_in_use]) + prefill_kv_page_indptr.append( + prefill_kv_page_indptr[-1] + num_blocks_in_use + ) + prefill_kv_last_page_len.append( + current_total_len % self.block_size or self.block_size + ) + + for seq_metadata in seq_metadata_list: + if seq_metadata.is_prompt: + continue + + if seq_metadata.block_table is None: + self.is_profiling_iteration = True + return + + self.contains_decode = True + + context_len = seq_metadata.seq.get_len() + # indptr for the prompt tokens in q/o tensor + decode_qo_indptr.append(decode_qo_indptr[-1] + 1) + # Compute the kv page indices for the prompt tokens. + num_blocks_in_use = (context_len + self.block_size - 1) // self.block_size + decode_kv_page_indices.extend(seq_metadata.block_table[:num_blocks_in_use]) + decode_kv_page_indptr.append(decode_kv_page_indptr[-1] + num_blocks_in_use) + decode_kv_last_page_len.append( + context_len % self.block_size or self.block_size + ) + + if self.contains_prefill: + self.prefill_wrapper.begin_forward( + self.to_int_tensor(prefill_qo_indptr), + self.to_int_tensor(prefill_kv_page_indptr), + self.to_int_tensor(prefill_kv_page_indices), + self.to_int_tensor(prefill_kv_last_page_len), + self.num_q_heads, + self.num_kv_heads, + self.head_dim, + self.block_size, + ) + + if self.contains_decode: + self.decode_wrapper.begin_forward( + self.to_int_tensor(decode_qo_indptr), + self.to_int_tensor(decode_kv_page_indptr), + self.to_int_tensor(decode_kv_page_indices), + self.to_int_tensor(decode_kv_last_page_len), + self.num_q_heads, + self.num_kv_heads, + self.head_dim, + self.block_size, + ) + + self.num_prefill_tokens = prefill_qo_indptr[-1] + self.num_total_tokens = self.num_prefill_tokens + len(decode_qo_indptr) - 1 + + self.append_qo_indptr_tensor = self.to_int_tensor( + prefill_qo_indptr[:-1] + + [x + prefill_qo_indptr[-1] for x in decode_qo_indptr] + ) + self.append_kv_page_indices_tensor = self.to_int_tensor( + prefill_kv_page_indices + decode_kv_page_indices + ) + self.append_kv_page_indptr_tensor = self.to_int_tensor( + prefill_kv_page_indptr[:-1] + + [x + prefill_kv_page_indptr[-1] for x in decode_kv_page_indptr] + ) + self.append_kv_last_page_len_tensor = self.to_int_tensor( + prefill_kv_last_page_len + decode_kv_last_page_len + ) + + def end_forward(self): + if self.contains_prefill: + self.prefill_wrapper.end_forward() + + if self.contains_decode: + self.decode_wrapper.end_forward() + + self.is_metadata_initialized = False + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + softmax_scale: float = 1.0, + layer_id: Optional[int] = None, + ) -> torch.Tensor: + assert self.is_metadata_initialized, "Metadata is not initialized." + + if self.is_profiling_iteration: + # there is no need to call attention in profiling mode + return torch.zeros_like(query) + + with self.get_timer(OperationMetrics.ATTN_INPUT_RESHAPE, layer_id): + query = query.contiguous().reshape(-1, self.num_q_heads, self.head_dim) + key = key.contiguous().reshape(-1, self.num_kv_heads, self.head_dim) + value = value.contiguous().reshape(-1, self.num_kv_heads, self.head_dim) + + output = torch.empty_like(query) + + with self.get_timer(OperationMetrics.ATTN_KV_CACHE_SAVE, layer_id): + append_paged_kv_cache( + key, + value, + self.append_qo_indptr_tensor, + kv_cache, + self.append_kv_page_indices_tensor, + self.append_kv_page_indptr_tensor, + self.append_kv_last_page_len_tensor, + kv_layout="NHD", + ) + + with self.get_timer(OperationMetrics.ATTN_PREFILL, layer_id): + if self.contains_prefill: + output[: self.num_prefill_tokens] = self.prefill_wrapper.forward( + query[: self.num_prefill_tokens], + kv_cache, + pos_encoding_mode="NONE", + sm_scale=softmax_scale, + ) + + with self.get_timer(OperationMetrics.ATTN_DECODE, layer_id): + if self.contains_decode: + output[self.num_prefill_tokens : self.num_total_tokens] = ( + self.decode_wrapper.forward( + query[self.num_prefill_tokens : self.num_total_tokens], + kv_cache, + pos_encoding_mode="NONE", + sm_scale=softmax_scale, + ) + ) + + with self.get_timer(OperationMetrics.ATTN_OUTPUT_RESHAPE, layer_id): + output = output.reshape(-1, self.num_q_heads * self.head_dim) + + return output diff --git a/sarathi/model_executor/attention/flashinfer_attention_wrapper.py b/sarathi/model_executor/attention/flashinfer_attention_wrapper.py index 80ae842..0270a46 100644 --- a/sarathi/model_executor/attention/flashinfer_attention_wrapper.py +++ b/sarathi/model_executor/attention/flashinfer_attention_wrapper.py @@ -46,6 +46,8 @@ def init( self.append_kv_page_indices_tensor = None self.append_kv_page_indptr_tensor = None self.append_kv_last_page_len_tensor = None + self.append_batch_indices_tensor = None # newly increased + self.append_positions_tensor = None # newly increased def to_int_tensor(self, data: List[int]) -> torch.Tensor: return torch.tensor(data, dtype=torch.int32, device="cuda") @@ -64,30 +66,19 @@ def begin_forward( self, seq_metadata_list: List[SequenceMetadata], ) -> None: - # The indptr tensor captures the location query tokens in the input tensor. - # |<---------------------- num_valid_tokens ----------------------------------------------------->| - # |<--------------- num_prompt_tokens -------------->||<------- num_generation_tokens (M) ------->| - # |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->||<--generation_0-->|...|<--generation_M-1-->|<--padding-->| - # - # Flashinfer calls this layout as a raggedtensor. The indptr tensor captures the start of each - # sequence in the ragged tensor. The length of the indptr tensor is the number of sequences + 1. - # We perform both prefill and decode attention in a single call to batched prefill kernel. - # prefill_qo_indptr: [0, prompt_0, prompt_0 + prompt_1, ..., prompt_0 + ... + prompt_N-1, generation_0, generation_0 + 1, ..., generation_0 + ... + M] prefill_qo_indptr: List[int] = [0] decode_qo_indptr: List[int] = [0] - # The kv_page_indices tensor captures the pages of the key-value cache that - # are assigned to each token in the input tensor. Since there is a variable number - # of pages assigned to each sequence, a ragged tensor to represent this. prefill_kv_page_indices: List[int] = [] decode_kv_page_indices: List[int] = [] - # the last page might not be full, so we need to keep track of the length of the last page prefill_kv_last_page_len: List[int] = [] decode_kv_last_page_len: List[int] = [] - # Since the prefill_kv_page_indices tensor is a ragged tensor, we also need to keep track of the - # indptr tensor for the prefill_kv_page_indices tensor. This tensor captures the start of each sequence - # in the ragged tensor. prefill_kv_page_indptr: List[int] = [0] decode_kv_page_indptr: List[int] = [0] + # newly increase batch_indices and positions + prefill_batch_indices: List[int] = [] + decode_batch_indices: List[int] = [] + prefill_positions: List[int] = [] + decode_positions: List[int] = [] self.is_profiling_iteration = False self.is_metadata_initialized = True @@ -95,15 +86,12 @@ def begin_forward( self.contains_prefill = False self.contains_decode = False - for seq_metadata in seq_metadata_list: + for seq_idx, seq_metadata in enumerate(seq_metadata_list): if not seq_metadata.is_prompt: continue - # ONLY used for profiling if seq_metadata.block_table is None: self.is_profiling_iteration = True - # During memory profiling, the block tables are not initialized yet. - # We will just skip the attention computation for now. return self.contains_prefill = True @@ -112,21 +100,18 @@ def begin_forward( processed_prompt_len = seq_metadata.seq.get_num_prompt_tokens_processed() current_total_len = processed_prompt_len + prompt_chunk_len - # indptr for the prompt tokens in q/o tensor prefill_qo_indptr.append(prefill_qo_indptr[-1] + prompt_chunk_len) - # Compute the kv page indices for the prompt tokens. - num_blocks_in_use = ( - current_total_len + self.block_size - 1 - ) // self.block_size + num_blocks_in_use = (current_total_len + self.block_size - 1) // self.block_size prefill_kv_page_indices.extend(seq_metadata.block_table[:num_blocks_in_use]) - prefill_kv_page_indptr.append( - prefill_kv_page_indptr[-1] + num_blocks_in_use - ) + prefill_kv_page_indptr.append(prefill_kv_page_indptr[-1] + num_blocks_in_use) prefill_kv_last_page_len.append( current_total_len % self.block_size or self.block_size ) + # computing batch_indices and positions + prefill_batch_indices.extend([seq_idx] * prompt_chunk_len) + prefill_positions.extend(range(processed_prompt_len, current_total_len)) - for seq_metadata in seq_metadata_list: + for seq_idx, seq_metadata in enumerate(seq_metadata_list): if seq_metadata.is_prompt: continue @@ -137,15 +122,16 @@ def begin_forward( self.contains_decode = True context_len = seq_metadata.seq.get_len() - # indptr for the prompt tokens in q/o tensor decode_qo_indptr.append(decode_qo_indptr[-1] + 1) - # Compute the kv page indices for the prompt tokens. num_blocks_in_use = (context_len + self.block_size - 1) // self.block_size decode_kv_page_indices.extend(seq_metadata.block_table[:num_blocks_in_use]) decode_kv_page_indptr.append(decode_kv_page_indptr[-1] + num_blocks_in_use) decode_kv_last_page_len.append( context_len % self.block_size or self.block_size ) + # computing batch_indices and positions + decode_batch_indices.append(seq_idx) + decode_positions.append(context_len - 1) # Only the latest token is appended when decoding if self.contains_prefill: self.prefill_wrapper.begin_forward( @@ -175,19 +161,23 @@ def begin_forward( self.num_total_tokens = self.num_prefill_tokens + len(decode_qo_indptr) - 1 self.append_qo_indptr_tensor = self.to_int_tensor( - prefill_qo_indptr[:-1] - + [x + prefill_qo_indptr[-1] for x in decode_qo_indptr] + prefill_qo_indptr[:-1] + [x + prefill_qo_indptr[-1] for x in decode_qo_indptr] ) self.append_kv_page_indices_tensor = self.to_int_tensor( prefill_kv_page_indices + decode_kv_page_indices ) self.append_kv_page_indptr_tensor = self.to_int_tensor( - prefill_kv_page_indptr[:-1] - + [x + prefill_kv_page_indptr[-1] for x in decode_kv_page_indptr] + prefill_kv_page_indptr[:-1] + [x + prefill_kv_page_indptr[-1] for x in decode_kv_page_indptr] ) self.append_kv_last_page_len_tensor = self.to_int_tensor( prefill_kv_last_page_len + decode_kv_last_page_len ) + self.append_batch_indices_tensor = self.to_int_tensor( + prefill_batch_indices + decode_batch_indices + ) + self.append_positions_tensor = self.to_int_tensor( + prefill_positions + decode_positions + ) def end_forward(self): if self.contains_prefill: @@ -210,7 +200,6 @@ def forward( assert self.is_metadata_initialized, "Metadata is not initialized." if self.is_profiling_iteration: - # there is no need to call attention in profiling mode return torch.zeros_like(query) with self.get_timer(OperationMetrics.ATTN_INPUT_RESHAPE, layer_id): @@ -224,7 +213,8 @@ def forward( append_paged_kv_cache( key, value, - self.append_qo_indptr_tensor, + self.append_batch_indices_tensor, # change to batch_indices + self.append_positions_tensor, # change to positions kv_cache, self.append_kv_page_indices_tensor, self.append_kv_page_indptr_tensor, @@ -255,4 +245,4 @@ def forward( with self.get_timer(OperationMetrics.ATTN_OUTPUT_RESHAPE, layer_id): output = output.reshape(-1, self.num_q_heads * self.head_dim) - return output + return output \ No newline at end of file