diff --git a/fastvideo/v1/attention/backends/__init__.py b/fastvideo/v1/attention/backends/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/fastvideo/v1/attention/backends/abstract.py b/fastvideo/v1/attention/backends/abstract.py deleted file mode 100644 index 0cd95e074..000000000 --- a/fastvideo/v1/attention/backends/abstract.py +++ /dev/null @@ -1,300 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from abc import ABC, abstractmethod -from contextlib import contextmanager -from dataclasses import dataclass, fields -from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, - Protocol, Set, Tuple, Type, TypeVar) - -import torch - -from vllm.multimodal import MultiModalPlaceholderMap - -if TYPE_CHECKING: - from vllm.worker.model_runner_base import (ModelRunnerBase, - ModelRunnerInputBase, - ModelRunnerInputBuilderBase) - - -class AttentionType: - """ - Attention type. - Use string to be compatible with `torch.compile`. - """ - # Decoder attention between previous layer Q/K/V - DECODER = "decoder" - # Encoder attention between previous layer Q/K/V for encoder-decoder - ENCODER = "encoder" - # Encoder attention between previous layer Q/K/V - ENCODER_ONLY = "encoder_only" - # Attention between dec. Q and enc. K/V for encoder-decoder - ENCODER_DECODER = "encoder_decoder" - - -class AttentionBackend(ABC): - """Abstract class for attention backends.""" - # For some attention backends, we allocate an output tensor before - # calling the custom op. When piecewise cudagraph is enabled, this - # makes sure the output tensor is allocated inside the cudagraph. - accept_output_buffer: bool = False - - @staticmethod - @abstractmethod - def get_name() -> str: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_impl_cls() -> Type["AttentionImpl"]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_metadata_cls() -> Type["AttentionMetadata"]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_state_cls() -> Type["AttentionState"]: - raise NotImplementedError - - @classmethod - def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": - return cls.get_metadata_cls()(*args, **kwargs) - - @staticmethod - @abstractmethod - def get_builder_cls() -> Type["AttentionMetadataBuilder"]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - ) -> Tuple[int, ...]: - raise NotImplementedError - - @staticmethod - @abstractmethod - def swap_blocks( - src_kv_cache: torch.Tensor, - dst_kv_cache: torch.Tensor, - src_to_dst: torch.Tensor, - ) -> None: - raise NotImplementedError - - @staticmethod - @abstractmethod - def copy_blocks( - kv_caches: List[torch.Tensor], - src_to_dists: torch.Tensor, - ) -> None: - raise NotImplementedError - - def advance_step(self, model_input: "ModelRunnerInputBase", - sampled_token_ids: Optional[torch.Tensor], - block_size: int, num_seqs: int, num_queries: int) -> None: - raise NotImplementedError - - -@dataclass -class AttentionMetadata: - """Attention metadata for prefill and decode batched together.""" - # Total number of prefill requests. - num_prefills: int - # Number of prefill tokens. - num_prefill_tokens: int - # Number of decode tokens. Note that it is equivalent to the number of - # decode requests. - num_decode_tokens: int - # (num_tokens,). The indices of the token slots that input tokens will be - # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size - # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot - # in block 0, and 1st slot in block 1, respectively. - slot_mapping: torch.Tensor - - # The index maps that relate multi-modal embeddings to the corresponding - # placeholders. - # - # N.B. These aren't really related to attention and don't belong on this - # type -- this is just a temporary solution to make them available to - # `model_executable`. - multi_modal_placeholder_index_maps: Optional[Dict[ - str, MultiModalPlaceholderMap.IndexMap]] - - # Enable/disable KV scales calculation. This is so that we can disable the - # calculation until after prefill and cuda graph capture. - enable_kv_scales_calculation: bool - - @property - @abstractmethod - def prefill_metadata(self) -> Optional["AttentionMetadata"]: - """Return the attention metadata that's required to run prefill - attention.""" - pass - - @property - @abstractmethod - def decode_metadata(self) -> Optional["AttentionMetadata"]: - """Return the attention metadata that's required to run decode - attention.""" - pass - - def asdict_zerocopy(self, - skip_fields: Optional[Set[str]] = None - ) -> Dict[str, Any]: - """Similar to dataclasses.asdict, but avoids deepcopying.""" - if skip_fields is None: - skip_fields = set() - # Note that if we add dataclasses as fields, they will need - # similar handling. - return { - field.name: getattr(self, field.name) - for field in fields(self) if field.name not in skip_fields - } - - -T = TypeVar("T", bound=AttentionMetadata) - - -class AttentionState(ABC, Generic[T]): - """Holds attention backend-specific objects reused during the - lifetime of the model runner.""" - - @abstractmethod - def __init__(self, runner: "ModelRunnerBase"): - ... - - @abstractmethod - @contextmanager - def graph_capture(self, max_batch_size: int): - """Context manager used when capturing CUDA graphs.""" - yield - - @abstractmethod - def graph_clone(self, batch_size: int) -> "AttentionState[T]": - """Clone attention state to save in CUDA graph metadata.""" - ... - - @abstractmethod - def graph_capture_get_metadata_for_batch( - self, - batch_size: int, - is_encoder_decoder_model: bool = False) -> T: - """Get attention metadata for CUDA graph capture of batch_size.""" - ... - - @abstractmethod - def get_graph_input_buffers( - self, - attn_metadata: T, - is_encoder_decoder_model: bool = False) -> Dict[str, Any]: - """Get attention-specific input buffers for CUDA graph capture.""" - ... - - @abstractmethod - def prepare_graph_input_buffers( - self, - input_buffers: Dict[str, Any], - attn_metadata: T, - is_encoder_decoder_model: bool = False) -> None: - """In-place modify input buffers dict for CUDA graph replay.""" - ... - - @abstractmethod - def begin_forward(self, model_input: "ModelRunnerInputBase") -> None: - """Prepare state for forward pass.""" - ... - - -class AttentionMetadataBuilder(ABC, Generic[T]): - """Abstract class for attention metadata builders.""" - - @abstractmethod - def __init__(self, input_builder: "ModelRunnerInputBuilderBase") -> None: - """Create the builder, remember some configuration and parameters.""" - raise NotImplementedError - - @abstractmethod - def prepare(self) -> None: - """Prepare for one batch.""" - raise NotImplementedError - - @abstractmethod - def build(self, seq_lens: List[int], query_lens: List[int], - cuda_graph_pad_size: int, batch_size: int) -> T: - """Build attention metadata with on-device tensors.""" - raise NotImplementedError - - -class AttentionLayer(Protocol): - - _k_scale: torch.Tensor - _v_scale: torch.Tensor - _k_scale_float: float - _v_scale_float: float - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - ... - - -class AttentionImpl(ABC, Generic[T]): - - @abstractmethod - def __init__( - self, - num_heads: int, - head_size: int, - scale: float, - num_kv_heads: Optional[int] = None, - alibi_slopes: Optional[List[float]] = None, - sliding_window: Optional[int] = None, - kv_cache_dtype: str = "auto", - blocksparse_params: Optional[Dict[str, Any]] = None, - logits_soft_cap: Optional[float] = None, - attn_type: str = AttentionType.DECODER, - ) -> None: - raise NotImplementedError - - @abstractmethod - def forward( - self, - layer: AttentionLayer, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: T, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - raise NotImplementedError - - -class MLAAttentionImpl(AttentionImpl[T], Generic[T]): - - @abstractmethod - def forward( - self, - layer: AttentionLayer, - hidden_states_or_cq: torch.Tensor, - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: T, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - raise NotImplementedError - - -def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: - return kv_cache_dtype != "auto" diff --git a/fastvideo/v1/attention/backends/flash_attn.py b/fastvideo/v1/attention/backends/flash_attn.py deleted file mode 100755 index c9c1e0c35..000000000 --- a/fastvideo/v1/attention/backends/flash_attn.py +++ /dev/null @@ -1,1042 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Attention layer with FlashAttention.""" -from collections import defaultdict -from dataclasses import dataclass -from itertools import accumulate -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type - -import torch -import math - -from vllm import _custom_ops as ops -# yapf conflicts with isort for this block -# yapf: disable -from fastvideo.v1.attention.backends.abstract import (AttentionBackend, AttentionImpl, - AttentionLayer, - AttentionMetadata, - AttentionMetadataBuilder, - AttentionType, - is_quantized_kv_cache) -# yapf: enable -from fastvideo.v1.attention.backends.utils import ( - PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping, - compute_slot_mapping_start_idx, get_flash_attn_version, - get_num_prefill_decode_query_kv_tokens, get_seq_len_block_table_args, - is_all_cross_attn_metadata_set, is_all_encoder_attn_metadata_set, - is_block_tables_empty) -from fastvideo.v1.logger import init_logger -# from vllm.multimodal import MultiModalPlaceholderMap -from vllm.utils import async_tensor_h2d, make_tensor_with_pad -from vllm.vllm_flash_attn import (flash_attn_varlen_func, - flash_attn_with_kvcache) - -if TYPE_CHECKING: - from vllm.worker.model_runner import (ModelInputForGPUBuilder, - ModelInputForGPUWithSamplingMetadata) - -logger = init_logger(__name__) - - -class FlashAttentionBackend(AttentionBackend): - - accept_output_buffer: bool = True - - @staticmethod - def get_supported_head_sizes() -> List[int]: - return [32, 64, 96, 128, 160, 192, 224, 256] - - @staticmethod - def get_name() -> str: - return "FLASH_ATTN" - - @staticmethod - def get_impl_cls() -> Type["FlashAttentionImpl"]: - return FlashAttentionImpl - - @staticmethod - def get_metadata_cls() -> Type["AttentionMetadata"]: - return FlashAttentionMetadata - - @staticmethod - def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]: - return FlashAttentionMetadataBuilder - - @staticmethod - def get_state_cls() -> Type["CommonAttentionState"]: - return CommonAttentionState - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - ) -> Tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - return (2, num_blocks, block_size, num_kv_heads, head_size) - - @staticmethod - def swap_blocks( - src_kv_cache: torch.Tensor, - dst_kv_cache: torch.Tensor, - src_to_dst: torch.Tensor, - ) -> None: - src_key_cache = src_kv_cache[0] - dst_key_cache = dst_kv_cache[0] - ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst) - src_value_cache = src_kv_cache[1] - dst_value_cache = dst_kv_cache[1] - ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst) - - @staticmethod - def copy_blocks( - kv_caches: List[torch.Tensor], - src_to_dists: torch.Tensor, - ) -> None: - key_caches = [kv_cache[0] for kv_cache in kv_caches] - value_caches = [kv_cache[1] for kv_cache in kv_caches] - - ops.copy_blocks(key_caches, value_caches, src_to_dists) - - -@dataclass -class FlashAttentionMetadata(AttentionMetadata): - """Metadata for FlashAttentionBackend. - - NOTE: Any python object stored here is not updated when it is - cuda-graph replayed. If you have values that need to be changed - dynamically, it should be stored in tensor. The tensor has to be - updated from `CUDAGraphRunner.forward` API. - """ - # (batch_size,). The sequence length per sequence. Sequence length means - # the computed tokens + new tokens None if it is a decoding. - seq_lens: Optional[List[int]] - # seq_lens stored as a tensor. - seq_lens_tensor: Optional[torch.Tensor] - - # NOTE(sang): Definition of context_len, query_len, and seq_len. - # |---------- N-1 iteration --------| - # |---------------- N iteration ---------------------| - # |- tokenA -|......................|-- newTokens ---| - # |---------- context_len ----------| - # |-------------------- seq_len ---------------------| - # |-- query_len ---| - - # Maximum sequence length among prefill batch. 0 if there are decoding - # requests only. - max_prefill_seq_len: int - # Maximum sequence length among decode batch. 0 if there are prefill - # requests only. - max_decode_seq_len: int - # (batch_size,) A tensor of context lengths (tokens that are computed - # so far). - context_lens_tensor: Optional[torch.Tensor] - - # (batch_size, max_blocks_per_seq). - # Block addresses per sequence. (Seq id -> list of physical block) - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks - # in the kv cache. Each block can contain up to block_size tokens. - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph - # captured. - block_tables: Optional[torch.Tensor] - - # Whether or not if cuda graph is enabled. - # Cuda-graph is currently enabled for decoding only. - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention. - - use_cuda_graph: bool - - # Maximum query length in the batch. - max_query_len: Optional[int] = None - - # Max number of query tokens among request in the batch. - max_decode_query_len: Optional[int] = None - - # (batch_size + 1,). The cumulative subquery lengths of the sequences in - # the batch, used to index into subquery. E.g., if the subquery length - # is [4, 6], it is [0, 4, 10]. - query_start_loc: Optional[torch.Tensor] = None - # (batch_size + 1,). The cumulative sequence lengths of the sequences in - # the batch, used to index into sequence. E.g., if the sequence length is - # [4, 6], it is [0, 4, 10]. - seq_start_loc: Optional[torch.Tensor] = None - - _cached_prefill_metadata: Optional["FlashAttentionMetadata"] = None - _cached_decode_metadata: Optional["FlashAttentionMetadata"] = None - - # Begin encoder attn & enc/dec cross-attn fields... - - # Encoder sequence lengths representation - encoder_seq_lens: Optional[List[int]] = None - encoder_seq_lens_tensor: Optional[torch.Tensor] = None - # (batch_size + 1,). The cumulative sequence lengths of the sequences in - # the batch, used to index into sequence. E.g., if the sequence length is - # [4, 6], it is [0, 4, 10]. - encoder_seq_start_loc: Optional[torch.Tensor] = None - # Maximum sequence length among encoder sequences - max_encoder_seq_len: Optional[int] = None - # Number of tokens input to encoder - num_encoder_tokens: Optional[int] = None - - # Cross-attention memory-mapping data structures: slot mapping - # and block tables - cross_slot_mapping: Optional[torch.Tensor] = None - cross_block_tables: Optional[torch.Tensor] = None - - @property - def is_all_encoder_attn_metadata_set(self): - ''' - All attention metadata required for encoder attention is set. - ''' - return is_all_encoder_attn_metadata_set(self) - - @property - def is_all_cross_attn_metadata_set(self): - ''' - All attention metadata required for enc/dec cross-attention is set. - - Superset of encoder attention required metadata. - ''' - return is_all_cross_attn_metadata_set(self) - - @property - def prefill_metadata(self) -> Optional["FlashAttentionMetadata"]: - if self.num_prefills == 0: - return None - - if self._cached_prefill_metadata is not None: - return self._cached_prefill_metadata - - assert ((self.seq_lens is not None) - or (self.encoder_seq_lens is not None)) - assert ((self.seq_lens_tensor is not None) - or (self.encoder_seq_lens_tensor is not None)) - - # Compute some attn_metadata fields which default to None - query_start_loc = (None if self.query_start_loc is None else - self.query_start_loc[:self.num_prefills + 1]) - slot_mapping = (None if self.slot_mapping is None else - self.slot_mapping[:self.num_prefill_tokens]) - seq_lens = (None if self.seq_lens is None else - self.seq_lens[:self.num_prefills]) - seq_lens_tensor = (None if self.seq_lens_tensor is None else - self.seq_lens_tensor[:self.num_prefills]) - seq_start_loc = (None if self.seq_start_loc is None else - self.seq_start_loc[:self.num_prefills + 1]) - context_lens_tensor = (None if self.context_lens_tensor is None else - self.context_lens_tensor[:self.num_prefills]) - block_tables = (None if self.block_tables is None else - self.block_tables[:self.num_prefills]) - - self._cached_prefill_metadata = FlashAttentionMetadata( - num_prefills=self.num_prefills, - num_prefill_tokens=self.num_prefill_tokens, - num_decode_tokens=0, - slot_mapping=slot_mapping, - multi_modal_placeholder_index_maps=self. - multi_modal_placeholder_index_maps, - enable_kv_scales_calculation=self.enable_kv_scales_calculation, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - max_query_len=self.max_query_len, - max_prefill_seq_len=self.max_prefill_seq_len, - max_decode_query_len=0, - max_decode_seq_len=0, - query_start_loc=query_start_loc, - seq_start_loc=seq_start_loc, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=False, - # Begin encoder & cross attn fields below... - encoder_seq_lens=self.encoder_seq_lens, - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor, - encoder_seq_start_loc=self.encoder_seq_start_loc, - max_encoder_seq_len=self.max_encoder_seq_len, - cross_slot_mapping=self.cross_slot_mapping, - cross_block_tables=self.cross_block_tables) - return self._cached_prefill_metadata - - @property - def decode_metadata(self) -> Optional["FlashAttentionMetadata"]: - if self.num_decode_tokens == 0: - return None - - if self._cached_decode_metadata is not None: - return self._cached_decode_metadata - assert ((self.seq_lens_tensor is not None) - or (self.encoder_seq_lens_tensor is not None)) - - # Compute some attn_metadata fields which default to None - slot_mapping = (None if self.slot_mapping is None else - self.slot_mapping[self.num_prefill_tokens:]) - seq_lens_tensor = (None if self.seq_lens_tensor is None else - self.seq_lens_tensor[self.num_prefills:]) - block_tables = (None if self.block_tables is None else - self.block_tables[self.num_prefills:]) - - self._cached_decode_metadata = FlashAttentionMetadata( - num_prefills=0, - num_prefill_tokens=0, - num_decode_tokens=self.num_decode_tokens, - slot_mapping=slot_mapping, - multi_modal_placeholder_index_maps=None, - enable_kv_scales_calculation=True, - seq_lens=None, - seq_lens_tensor=seq_lens_tensor, - max_decode_query_len=self.max_decode_query_len, - max_query_len=self.max_query_len, - max_prefill_seq_len=0, - max_decode_seq_len=self.max_decode_seq_len, - # Batch may be composed of prefill|decodes, adjust query start - # indices to refer to the start of decodes. E.g. - # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6]. - query_start_loc=(self.query_start_loc[self.num_prefills:] - - self.query_start_loc[self.num_prefills]) - if self.query_start_loc is not None else None, - seq_start_loc=self.seq_start_loc[self.num_prefills:] - if self.seq_start_loc is not None else None, - context_lens_tensor=None, - block_tables=block_tables, - use_cuda_graph=self.use_cuda_graph, - # Begin encoder & cross attn fields below... - encoder_seq_lens=self.encoder_seq_lens, - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor, - encoder_seq_start_loc=self.encoder_seq_start_loc, - max_encoder_seq_len=self.max_encoder_seq_len, - cross_slot_mapping=self.cross_slot_mapping, - cross_block_tables=self.cross_block_tables) - return self._cached_decode_metadata - - def advance_step(self, - model_input: "ModelInputForGPUWithSamplingMetadata", - sampled_token_ids: Optional[torch.Tensor], - block_size: int, - num_seqs: int, - num_queries: int, - turn_prefills_into_decodes: bool = False): - """ - Update metadata in-place to advance one decode step. - """ - # When using cudagraph, the num_seqs is padded to the next captured - # batch sized, but num_queries tracks the actual number of requests in - # the batch. For --enforce-eager mode, num_seqs == num_queries - if num_seqs != num_queries: - assert num_seqs > num_queries - assert self.use_cuda_graph - - if turn_prefills_into_decodes: - # When Mutli-Step is enabled with Chunked-Prefill, prefills and - # decodes are scheduled together. In the first step, all the - # prefills turn into decodes. This update reflects that - # conversion. - assert self.num_decode_tokens + self.num_prefills == num_seqs - self.num_decode_tokens += self.num_prefills - self.num_prefills = 0 - self.num_prefill_tokens = 0 - self.max_prefill_seq_len = 0 - self.max_query_len = 1 - - self.slot_mapping = self.slot_mapping[:num_seqs] - else: - assert self.seq_lens is not None - assert self.max_decode_seq_len == max(self.seq_lens) - - assert self.num_prefills == 0 - assert self.num_prefill_tokens == 0 - assert self.num_decode_tokens == num_seqs - assert self.slot_mapping.shape == (num_seqs, ) - - assert self.seq_lens is not None - assert len(self.seq_lens) == num_seqs - assert self.seq_lens_tensor is not None - assert self.seq_lens_tensor.shape == (num_seqs, ) - assert self.max_query_len == 1 - assert self.max_prefill_seq_len == 0 - - assert self.query_start_loc is not None - assert self.query_start_loc.shape == (num_queries + 1, ) - assert self.seq_start_loc is not None - assert self.seq_start_loc.shape == (num_seqs + 1, ) - - assert self.context_lens_tensor is not None - assert self.context_lens_tensor.shape == (num_queries, ) - - assert self.block_tables is not None - assert self.block_tables.shape[0] == num_seqs - - # Update query lengths. Note that we update only queries and not seqs, - # since tensors may be padded due to captured cuda graph batch size - for i in range(num_queries): - self.seq_lens[i] += 1 - self.max_decode_seq_len = max(self.seq_lens) - - ops.advance_step_flashattn(num_seqs=num_seqs, - num_queries=num_queries, - block_size=block_size, - input_tokens=model_input.input_tokens, - sampled_token_ids=sampled_token_ids, - input_positions=model_input.input_positions, - seq_lens=self.seq_lens_tensor, - slot_mapping=self.slot_mapping, - block_tables=self.block_tables) - - -class FlashAttentionMetadataBuilder( - AttentionMetadataBuilder[FlashAttentionMetadata]): - - def __init__(self, input_builder: "ModelInputForGPUBuilder"): - self.input_builder = input_builder - self.runner = input_builder.runner - self.sliding_window = input_builder.sliding_window - self.block_size = input_builder.block_size - - def prepare(self): - self.slot_mapping: List[int] = [] - self.prefill_seq_lens: List[int] = [] - self.context_lens: List[int] = [] - self.block_tables: List[List[int]] = [] - self.curr_seq_lens: List[int] = [] - self.multimodal_placeholder_maps: Dict[ - str, - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap) - self.num_prefills = 0 - self.num_prefill_tokens = 0 - self.num_decode_tokens = 0 - self.has_prefix_cache_hit = False - - def _add_seq_group( - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", - chunked_prefill_enabled: bool, prefix_cache_hit: bool): - """Add a sequence group to the metadata. Specifically update/append - 1. context length. - 2. block table. - 3. slot mapping. - """ - is_prompt = inter_data.is_prompt - block_tables = inter_data.block_tables - - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, - curr_sliding_window_block) in zip( - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens], - inter_data.orig_seq_lens, inter_data.seq_lens, - inter_data.query_lens, inter_data.context_lens, - inter_data.curr_sliding_window_blocks): - self.context_lens.append(context_len) - - if is_prompt: - mm_maps = inter_data.multi_modal_placeholder_maps - if mm_maps: - for modality, placeholders in mm_maps.items(): - self.multimodal_placeholder_maps[modality].extend( - placeholders) - - self.num_prefills += 1 - self.num_prefill_tokens += token_len - self.prefill_seq_lens.append(seq_len) - else: - self.num_decode_tokens += query_len - self.curr_seq_lens.append(curr_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: - # NOTE(woosuk): For flash-attn, the block table should - # include the entries for the incoming prefill tokens. - block_table = block_tables[seq_id] - elif ((chunked_prefill_enabled or not is_prompt) - and block_tables is not None): - if curr_sliding_window_block == 0: - block_table = block_tables[seq_id] - else: - block_table = block_tables[seq_id][ - -curr_sliding_window_block:] - 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) - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id, - seq_len, context_len, start_idx, - self.block_size, inter_data.block_tables) - - def _get_graph_runner_block_tables( - self, num_seqs: int, - block_tables: List[List[int]]) -> torch.Tensor: - # The shape of graph_block_tables is - # [max batch size, max context len // block size]. - max_batch_size, max_blocks = self.runner.graph_block_tables.shape - assert max_batch_size >= num_seqs - - graph_block_tables = self.runner.graph_block_tables[:num_seqs] - for i, block_table in enumerate(block_tables): - if block_table: - num_blocks = len(block_table) - if num_blocks <= max_blocks: - graph_block_tables[i, :num_blocks] = block_table - else: - # It may be possible to have more blocks allocated due - # to lookahead slots of multi-step, however, they are - # not used anyway, so can be safely ignored. - graph_block_tables[ - i, :max_blocks] = block_table[:max_blocks] - - return torch.from_numpy(graph_block_tables).to( - device=self.runner.device, non_blocking=True) - - def build(self, seq_lens: List[int], query_lens: List[int], - cuda_graph_pad_size: int, batch_size: int): - """Build attention metadata with on-device tensors. - - Args: - seq_lens: The maybe padded sequence lengths of the input sequences. - query_lens: The query lengths of the input sequences. - cuda_graph_pad_size: The padding size for cuda graph. - -1 if cuda graph is not used. - batch_size: The maybe padded batch size. - """ - prefix_cache_hit = any([ - inter_data.prefix_cache_hit - for inter_data in self.input_builder.inter_data_list - ]) - for inter_data in self.input_builder.inter_data_list: - self._add_seq_group(inter_data, - self.input_builder.chunked_prefill_enabled, - prefix_cache_hit) - - device = self.runner.device - use_captured_graph = cuda_graph_pad_size != -1 - - max_query_len = max(query_lens) - decode_query_lens = query_lens[self.num_prefills:] - if len(decode_query_lens) > 0: - max_decode_query_len = max(decode_query_lens) - else: - max_decode_query_len = 1 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0) - max_decode_seq_len = max(self.curr_seq_lens, default=0) - num_decode_tokens = self.num_decode_tokens - query_start_loc = list(accumulate(query_lens, initial=0)) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - - num_seqs = len(seq_lens) - 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 - self.num_prefill_tokens - block_tables = self._get_graph_runner_block_tables( - num_seqs, self.block_tables) - else: - block_tables = make_tensor_with_pad( - self.block_tables, - pad=0, - dtype=torch.int, - device=device, - ) - assert max_query_len > 0, ("query_lens: {}".format(query_lens)) - - assert device is not None - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int, - device, self.runner.pin_memory) - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device, - self.runner.pin_memory) - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long, - device, self.runner.pin_memory) - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32, - device, - self.runner.pin_memory) - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32, - device, self.runner.pin_memory) - placeholder_index_maps = { - modality: placeholder_map.index_map() - for modality, placeholder_map in - self.multimodal_placeholder_maps.items() - } - - return FlashAttentionMetadata( - num_prefills=self.num_prefills, - slot_mapping=slot_mapping_tensor, - num_prefill_tokens=self.num_prefill_tokens, - num_decode_tokens=num_decode_tokens, - seq_lens=seq_lens, - multi_modal_placeholder_index_maps=placeholder_index_maps, - enable_kv_scales_calculation=True, - seq_lens_tensor=seq_lens_tensor, - max_query_len=max_query_len, - max_decode_query_len=max_decode_query_len, - max_prefill_seq_len=max_prefill_seq_len, - max_decode_seq_len=max_decode_seq_len, - query_start_loc=query_start_loc_tensor, - seq_start_loc=seq_start_loc_tensor, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=use_captured_graph, - ) - - -class FlashAttentionImpl(AttentionImpl): - """ - If the input tensors contain prompt tokens, the layout is as follows: - |<--------------- num_prefill_tokens ----------------->| - |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->| - - Otherwise, the layout is as follows: - |<----------------- num_decode_tokens ------------------>| - |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->| - - Generation tokens can contain padding when cuda-graph is used. - Currently, prompt tokens don't contain any padding. - - The prompts might have different lengths, while the generation tokens - always have length 1. - - If chunked prefill is enabled, prefill tokens and decode tokens can be - batched together in a flattened 1D query. - - |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->| - |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->| - - Currently, cuda graph is disabled for chunked prefill, meaning there's no - padding between prefill and decode tokens. - """ - - def __init__( - self, - num_heads: int, - head_size: int, - scale: float, - num_kv_heads: int, - alibi_slopes: Optional[List[float]], - sliding_window: Optional[int], - kv_cache_dtype: str, - blocksparse_params: Optional[Dict[str, Any]] = None, - logits_soft_cap: Optional[float] = None, - attn_type: str = AttentionType.DECODER, - ) -> None: - if blocksparse_params is not None: - raise ValueError( - "FlashAttention does not support block-sparse attention.") - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads - if alibi_slopes is not None: - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) - self.alibi_slopes = alibi_slopes - self.sliding_window = ((sliding_window - 1, - 0) if sliding_window is not None else (-1, -1)) - self.kv_cache_dtype = kv_cache_dtype - if is_quantized_kv_cache(self.kv_cache_dtype): - raise NotImplementedError( - "FlashAttention with FP8 KV cache not yet supported") - if logits_soft_cap is None: - # In flash-attn, setting logits_soft_cap as 0 means no soft cap. - logits_soft_cap = 0 - self.logits_soft_cap = logits_soft_cap - - assert self.num_heads % self.num_kv_heads == 0 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads - - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes() - if head_size not in support_head_sizes: - raise ValueError( - f"Head size {head_size} is not supported by FlashAttention. " - f"Supported head sizes are: {support_head_sizes}.") - self.attn_type = attn_type - self.vllm_flash_attn_version = get_flash_attn_version() - - def forward( - self, - layer: AttentionLayer, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: FlashAttentionMetadata, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Forward pass with FlashAttention. - - Args: - query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] - output: shape = [num_tokens, num_heads, head_size] - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size] - NOTE: kv_cache will be an empty tensor with shape [0] - for profiling run. - attn_metadata: Metadata for attention. - NOTE: It in-place updates the output tensor. - """ - # NOTE(woosuk): FlashAttention does not support FP8 KV cache. - assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0, ( - "key/v_scale is not supported in FlashAttention.") - - assert output is not None, "Output tensor must be provided." - - attn_type = self.attn_type - if (attn_type == AttentionType.ENCODER - and (not attn_metadata.is_all_encoder_attn_metadata_set)): - raise AttributeError("Encoder attention requires setting " - "encoder metadata attributes.") - elif (attn_type == AttentionType.ENCODER_DECODER - and (not attn_metadata.is_all_cross_attn_metadata_set)): - raise AttributeError("Encoder/decoder cross-attention " - "requires setting cross-attention " - "metadata attributes.") - - kv_cache_dtype: str = self.kv_cache_dtype - softmax_scale: float = self.scale - window_size = self.sliding_window - alibi_slopes: Optional[torch.Tensor] = self.alibi_slopes - logits_soft_cap: Optional[float] = self.logits_soft_cap - - if kv_cache.numel() > 0: - key_cache = kv_cache[0] - value_cache = kv_cache[1] - # We skip updating the KV cache under two conditions: - # a. When the Attention Type is ENCODER. In this phase, we compute - # only the encoder attention without updating the cache. - # b. When both Key and Value are None. This occurs during - # cross-attention computation in the decoding phase, where the - # KV cache is already populated with the cross-attention - # tensor. Thus, we skip cache updates during this time. - if (attn_type != AttentionType.ENCODER) and (key is not None) and ( - value is not None): - if attn_type == AttentionType.ENCODER_DECODER: - # Update cross-attention KV cache (prefill-only) - updated_slot_mapping = attn_metadata.cross_slot_mapping - else: - # Update self-attention KV cache (prefill/decode) - updated_slot_mapping = attn_metadata.slot_mapping - - # Reshape the input keys and values and store them in the cache. - # If kv_cache is not provided, the new key and value tensors are - # not cached. This happens during the initial memory - # profiling run. - torch.ops._C_cache_ops.reshape_and_cache_flash( - key, - value, - kv_cache[0], - kv_cache[1], - updated_slot_mapping.flatten(), # type: ignore[union-attr] - kv_cache_dtype, - layer._k_scale, - layer._v_scale, - ) - - (num_prefill_query_tokens, num_prefill_kv_tokens, - num_decode_query_tokens) = \ - get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type) - decode_query = query[num_prefill_query_tokens:] - decode_output = output[num_prefill_query_tokens:] - # QKV for prefill. - query = query[:num_prefill_query_tokens] - prefill_output = output[:num_prefill_query_tokens] - assert query.shape[0] == num_prefill_query_tokens - assert decode_query.shape[0] == num_decode_query_tokens - - if prefill_meta := attn_metadata.prefill_metadata: - # Prompt run. - if (kv_cache.numel() == 0 or prefill_meta.block_tables is None - or prefill_meta.block_tables.numel() == 0): - # normal attention - # When block_tables are not filled, it means q and k are the - # prompt, and they have the same length. - q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \ - _get_query_key_seq_metadata(prefill_meta, True, attn_type) - - key = key[:num_prefill_kv_tokens] - value = value[:num_prefill_kv_tokens] - - flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=q_seq_start_loc, - cu_seqlens_k=k_seq_start_loc, - max_seqlen_q=q_seq_len, - max_seqlen_k=k_seq_len, - softmax_scale=softmax_scale, - causal=_get_causal_option(attn_type), - window_size=window_size, - alibi_slopes=alibi_slopes, - softcap=logits_soft_cap, - out=prefill_output, - fa_version=self.vllm_flash_attn_version, - ) - else: - # prefix-enabled attention - assert attn_type == AttentionType.DECODER, ( - "Only decoder-only models support prefix caching") - assert prefill_meta.seq_lens is not None - max_seq_len = max(prefill_meta.seq_lens) - flash_attn_varlen_func( # noqa - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=prefill_meta.query_start_loc, - max_seqlen_q=prefill_meta.max_query_len, - seqused_k=prefill_meta.seq_lens_tensor, - max_seqlen_k=max_seq_len, - softmax_scale=softmax_scale, - causal=True, - window_size=window_size, - alibi_slopes=alibi_slopes, - block_table=prefill_meta.block_tables, - softcap=logits_soft_cap, - out=prefill_output, - fa_version=self.vllm_flash_attn_version, - ) - - if decode_meta := attn_metadata.decode_metadata: - # Decoding run. - # Use flash_attn_varlen_func kernel for speculative decoding - # because different queries might have different lengths. - - assert decode_meta.max_decode_query_len is not None - # use only for actual varlen decoding - if decode_meta.max_decode_query_len > 1: - assert attn_type == AttentionType.DECODER, ( - "Only decoder-only models support max_decode_query_len > 1" - ) - flash_attn_varlen_func( - q=decode_query, - k=key_cache, - v=value_cache, - cu_seqlens_q=decode_meta.query_start_loc, - max_seqlen_q=decode_meta.max_decode_query_len, - seqused_k=decode_meta.seq_lens_tensor, - max_seqlen_k=decode_meta.max_decode_seq_len, - softmax_scale=softmax_scale, - causal=True, - window_size=window_size, - alibi_slopes=alibi_slopes, - softcap=logits_soft_cap, - block_table=decode_meta.block_tables, - out=decode_output, - fa_version=self.vllm_flash_attn_version, - ) - else: - # Use flash_attn_with_kvcache for normal decoding. - ( - seq_lens_arg, - _, - block_tables_arg, - ) = get_seq_len_block_table_args(decode_meta, False, attn_type) - flash_attn_with_kvcache( - q=decode_query.unsqueeze(1), - k_cache=key_cache, - v_cache=value_cache, - block_table=block_tables_arg, - cache_seqlens=seq_lens_arg, - softmax_scale=softmax_scale, - causal=True, - window_size=window_size, - alibi_slopes=alibi_slopes, - softcap=logits_soft_cap, - out=decode_output.unsqueeze(1), - fa_version=self.vllm_flash_attn_version, - ) - return output - - -def _get_query_key_seq_metadata( - attn_metadata, - is_prompt: bool, - attn_type: str, -) -> tuple: - """ - Returns sequence metadata for key and query based on the specified - attention type and whether input is a prompt. - - This function computes the starting locations and maximum sequence lengths - for key and query sequences for different attention types. - - Args: - attn_metadata: The attention metadata object - is_prompt (bool): A flag indicating if the input is a prompt - attn_type (AttentionType): The type of attention being used. - - Returns: - tuple: A tuple containing four integers: - - Starting location for the query sequence. - - Maximum sequence length for the query sequence. - - Starting location for the key sequence. - - Maximum sequence length for the key sequence. - - Raises: - AttributeError: If an invalid attention type is provided. - """ - if attn_type == AttentionType.DECODER: - # Decoder self-attention - # Choose max_seq_len based on whether we are in prompt_run - if is_prompt: - max_seq_len = attn_metadata.max_prefill_seq_len - else: - max_seq_len = attn_metadata.max_decode_seq_len - return (attn_metadata.seq_start_loc, max_seq_len, - attn_metadata.seq_start_loc, max_seq_len) - - elif attn_type == AttentionType.ENCODER_DECODER: - # This is cross attention between the where the key - # is the precomputed encoder attention and query - # is the input sequence. - # Choose query max length based on whether it is prompt - # or not. - if is_prompt: - max_seq_len = attn_metadata.max_prefill_seq_len - else: - max_seq_len = attn_metadata.max_decode_seq_len - return (attn_metadata.seq_start_loc, max_seq_len, - attn_metadata.encoder_seq_start_loc, - attn_metadata.max_encoder_seq_len) - elif attn_type == AttentionType.ENCODER: - # For encoder attention both the query and the key are same i.e the - # encoder sequence. - return (attn_metadata.encoder_seq_start_loc, - attn_metadata.max_encoder_seq_len, - attn_metadata.encoder_seq_start_loc, - attn_metadata.max_encoder_seq_len) - elif attn_type == AttentionType.ENCODER_ONLY: - assert is_prompt, "Should not have decode for encoder only model." - return (attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len, - attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len) - else: - raise AttributeError(f"Invalid attention type {str(attn_type)}") - - -def _get_causal_option(attn_type: str) -> bool: - """ - Determine whether the given attention type is suitable for causal - attention mechanisms. - - Args: - attn_type (AttentionType): The type of attention being evaluated - - Returns: - bool: Returns `True` if the attention type is suitable for causal - attention (i.e., not encoder, encoder-only, or encoder-decoder), - otherwise returns `False`. - """ - return not (attn_type == AttentionType.ENCODER - or attn_type == AttentionType.ENCODER_ONLY - or attn_type == AttentionType.ENCODER_DECODER) - - -class TextEncoderFlashAttentionImpl: - """Simplified Flash Attention implementation for text encoders in diffusion models. - - This implementation removes KV caching and other generation-specific features, - focusing only on the encoder attention pattern needed for text encoders in - diffusion models like DiT. - """ - - def __init__( - self, - num_heads: int, - head_size: int, - scale: Optional[float] = None, - ) -> None: - """Initialize the text encoder flash attention implementation. - - Args: - num_heads: Number of attention heads - head_size: Size of each attention head - scale: Attention scale factor. If None, defaults to 1/sqrt(head_size) - """ - self.num_heads = num_heads - self.head_size = head_size - self.num_kv_heads = num_heads # For encoder, typically num_kv_heads = num_heads - - # Set default scale if not provided - if scale is None: - scale = 1.0 / math.sqrt(head_size) - self.scale = float(scale) - - # Set defaults for optional parameters - self.alibi_slopes = None - self.sliding_window = (-1, -1) # No sliding window by default - self.logits_soft_cap = 0 # No soft cap by default - - # Always use encoder-only attention type - self.attn_type = AttentionType.ENCODER_ONLY - self.vllm_flash_attn_version = get_flash_attn_version() - - # Verify head size is supported - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes() - if head_size not in support_head_sizes: - raise ValueError( - f"Head size {head_size} is not supported by FlashAttention. " - f"Supported head sizes are: {support_head_sizes}.") - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - seq_lens: Optional[List[int]] = None, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Forward pass with FlashAttention for text encoder. - - Args: - query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] - seq_lens: List of sequence lengths for batched input - output: shape = [num_tokens, num_heads, head_size] - Pre-allocated output tensor (optional) - - Returns: - torch.Tensor: Output tensor with shape [num_tokens, num_heads, head_size] - """ - device = query.device - batch_size = len(seq_lens) if seq_lens is not None else 1 - - if output is None: - output = torch.empty_like(query) - - if seq_lens is None: - # Single sequence case - use standard flash attention - return flash_attn_varlen_func( - q=query, - k=key, - v=value, - softmax_scale=self.scale, - causal=False, # Non-causal for encoder - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - softcap=self.logits_soft_cap, - out=output, - fa_version=self.vllm_flash_attn_version, - ) - else: - # Multiple sequences case - use varlen flash attention - max_seq_len = max(seq_lens) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - seq_start_loc_tensor = torch.tensor(seq_start_loc, - dtype=torch.int32, - device=device) - - return flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=seq_start_loc_tensor, - cu_seqlens_k=seq_start_loc_tensor, - max_seqlen_q=max_seq_len, - max_seqlen_k=max_seq_len, - softmax_scale=self.scale, - causal=False, # Non-causal for encoder - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - softcap=self.logits_soft_cap, - out=output, - fa_version=self.vllm_flash_attn_version, - ) diff --git a/fastvideo/v1/attention/backends/utils.py b/fastvideo/v1/attention/backends/utils.py deleted file mode 100644 index 4374b5422..000000000 --- a/fastvideo/v1/attention/backends/utils.py +++ /dev/null @@ -1,619 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Attention backend utils""" -from collections import defaultdict -from contextlib import contextmanager -from itertools import accumulate -from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, TypeVar, Union - -import numpy as np -import torch - -from vllm import envs -from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder, - AttentionState) -from vllm.attention.backends.abstract import AttentionType -from vllm.logger import init_logger -from vllm.multimodal import MultiModalPlaceholderMap -from vllm.platforms import current_platform -from vllm.utils import async_tensor_h2d, make_tensor_with_pad - -logger = init_logger(__name__) - -if TYPE_CHECKING: - from vllm.worker.model_runner_base import ModelRunnerBase - -# Error string(s) for encoder/decoder -# unsupported attention scenarios -STR_NOT_IMPL_ENC_DEC_ROCM_HIP = ("ROCm/HIP is not currently supported " - "with encoder/decoder models.") - -PAD_SLOT_ID = -1 - -# Switch to numpy implementation of compute_slot_mapping -# if we have at least this many elements. Could be tuned further. -_COMPUTE_SLOT_MAPPING_NUMPY_NUMEL = 256 - -if TYPE_CHECKING: - from vllm.worker.model_runner import ModelInputForGPUBuilder - - -def is_block_tables_empty(block_tables: Union[None, Dict]): - """ - Check if block_tables is None or a dictionary with all None values. - """ - if block_tables is None: - return True - return (isinstance(block_tables, dict) - and all(value is None for value in block_tables.values())) - - -def compute_slot_mapping_start_idx(is_prompt: bool, query_len: int, - context_len: 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) - 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]]): - """ - Compute slot mapping. - """ - if is_profile_run: - # During memory profiling, the block tables are not - # initialized yet. In this case, we just use a dummy - # slot mapping. - # In embeddings, the block tables are {seq_id: None}. - slot_mapping.extend([PAD_SLOT_ID] * seq_len) - return - - # Mask the [0, start_idx) tokens of the prompt with - # PAD_SLOT_ID, where start_idx is max(0, seq_len - - # 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]. - 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) - - -TAttentionMetadata = TypeVar("TAttentionMetadata", bound='AttentionMetadata') - - -class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]): - - _metadata_cls: Type[TAttentionMetadata] - - def __init__(self, input_builder: "ModelInputForGPUBuilder"): - self.input_builder = input_builder - self.runner = input_builder.runner - - self.sliding_window = input_builder.sliding_window - self.block_size = input_builder.block_size - - def prepare(self): - self.slot_mapping: List[int] = [] - self.prefill_seq_lens: List[int] = [] - self.context_lens: List[int] = [] - self.block_tables: List[List[int]] = [] - self.curr_seq_lens: List[int] = [] - self.multimodal_placeholder_maps: Dict[ - str, - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap) - self.num_prefills = 0 - self.num_prefill_tokens = 0 - self.num_decode_tokens = 0 - - def _add_seq_group( - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", - chunked_prefill_enabled: bool): - is_prompt = inter_data.is_prompt - block_tables = inter_data.block_tables - - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, - curr_sliding_window_block) in zip( - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens], - inter_data.orig_seq_lens, inter_data.seq_lens, - inter_data.query_lens, inter_data.context_lens, - inter_data.curr_sliding_window_blocks): - self.context_lens.append(context_len) - if is_prompt: - mm_maps = inter_data.multi_modal_placeholder_maps - if mm_maps: - for modality, placeholders in mm_maps.items(): - self.multimodal_placeholder_maps[modality].extend( - placeholders) - - 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.curr_seq_lens.append(curr_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 inter_data.prefix_cache_hit: - block_table = block_tables[seq_id] - elif ((chunked_prefill_enabled or not is_prompt) - and block_tables is not None): - if curr_sliding_window_block == 0: - block_table = block_tables[seq_id] - else: - block_table = block_tables[seq_id][ - -curr_sliding_window_block:] - 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) - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id, - seq_len, context_len, start_idx, - self.block_size, inter_data.block_tables) - - def build(self, seq_lens: List[int], query_lens: List[int], - cuda_graph_pad_size: int, batch_size: int): - """Build attention metadata with on-device tensors. - - Args: - seq_lens: The maybe padded sequence lengths of the input sequences. - query_lens: The query lengths of the input sequences. - cuda_graph_pad_size: The padding size for cuda graph. - -1 if cuda graph is not used. - batch_size: The maybe padded batch size. - """ - for inter_data in self.input_builder.inter_data_list: - self._add_seq_group(inter_data, - self.input_builder.chunked_prefill_enabled) - - device = self.runner.device - use_captured_graph = cuda_graph_pad_size != -1 - - max_query_len = max(query_lens) - max_prefill_seq_len = max(self.prefill_seq_lens, default=0) - max_decode_seq_len = max(self.curr_seq_lens, default=0) - num_decode_tokens = self.num_decode_tokens - query_start_loc = list(accumulate(query_lens, initial=0)) - seq_start_loc = list(accumulate(seq_lens, initial=0)) - - 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 - - # The shape of graph_block_tables is - # [max batch size, max context len // block size]. - input_block_tables = self.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.from_numpy(input_block_tables).to( - device, non_blocking=True) - else: - block_tables = make_tensor_with_pad( - self.block_tables, - pad=0, - dtype=torch.int, - device=device, - ) - assert max_query_len > 0, "query_lens: {}".format(query_lens) - - assert device is not None - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int, - device, self.runner.pin_memory) - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device, - self.runner.pin_memory) - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long, - device, self.runner.pin_memory) - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32, - device, - self.runner.pin_memory) - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32, - device, self.runner.pin_memory) - placeholder_index_maps = { - modality: placeholder_map.index_map() - for modality, placeholder_map in - self.multimodal_placeholder_maps.items() - } - - return self._metadata_cls( # type: ignore - num_prefills=self.num_prefills, - slot_mapping=slot_mapping_tensor, - multi_modal_placeholder_index_maps=placeholder_index_maps, - enable_kv_scales_calculation=True, - num_prefill_tokens=self.num_prefill_tokens, - num_decode_tokens=num_decode_tokens, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - max_query_len=max_query_len, - max_prefill_seq_len=max_prefill_seq_len, - max_decode_seq_len=max_decode_seq_len, - query_start_loc=query_start_loc_tensor, - seq_start_loc=seq_start_loc_tensor, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=use_captured_graph, - ) - - -class CommonAttentionState(AttentionState): - - def __init__(self, runner: "ModelRunnerBase"): - self.runner = runner - self._is_graph_capturing = False - - @contextmanager - def graph_capture(self, max_batch_size: int): - - self._is_graph_capturing = True - - self._graph_slot_mapping = torch.full((max_batch_size, ), - PAD_SLOT_ID, - dtype=torch.long, - device=self.runner.device) - self._graph_seq_lens = torch.ones(max_batch_size, - dtype=torch.int32, - device=self.runner.device) - self._graph_block_tables = torch.from_numpy( - self.runner.graph_block_tables).to(device=self.runner.device) - - yield - - self._is_graph_capturing = False - del self._graph_slot_mapping - del self._graph_seq_lens - del self._graph_block_tables - - def graph_clone(self, batch_size: int) -> "CommonAttentionState": - assert self._is_graph_capturing - return self.__class__(self.runner) - - def graph_capture_get_metadata_for_batch( - self, batch_size: int, is_encoder_decoder_model: bool = False): - assert self._is_graph_capturing - attn_metadata = self.runner.attn_backend.make_metadata( - num_prefills=0, - num_prefill_tokens=0, - num_decode_tokens=batch_size, - slot_mapping=self._graph_slot_mapping[:batch_size], - multi_modal_placeholder_index_maps=None, - enable_kv_scales_calculation=True, - seq_lens=None, - seq_lens_tensor=self._graph_seq_lens[:batch_size], - max_query_len=1, - max_decode_query_len=1, - max_prefill_seq_len=0, - max_decode_seq_len=self.runner.max_seq_len_to_capture, - query_start_loc=None, - seq_start_loc=None, - context_lens_tensor=None, - block_tables=self._graph_block_tables[:batch_size], - use_cuda_graph=True, - ) - if is_encoder_decoder_model: - # The encoder decoder model works only with XFormers and - # Flash Attention backend. Assert the same. - assert self.runner.attn_backend.get_name() in\ - ["XFORMERS", "FLASH_ATTN"], \ - f"Expected attn_backend name to be either 'XFORMERS' or " \ - f"'FLASH_ATTN', but "\ - f"got '{self.runner.attn_backend.get_name()}'" - self._update_captured_metadata_for_enc_dec_model( - batch_size=batch_size, attn_metadata=attn_metadata) - - return attn_metadata - - def get_graph_input_buffers( - self, - attn_metadata, - is_encoder_decoder_model: bool = False) -> Dict[str, Any]: - input_buffers = { - "slot_mapping": attn_metadata.slot_mapping, - "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor, - "block_tables": attn_metadata.decode_metadata.block_tables, - } - if is_encoder_decoder_model: - # The encoder decoder model works only with XFormers and - # Flash Attention backend. Assert the same. - assert self.runner.attn_backend.get_name() in\ - ["XFORMERS", "FLASH_ATTN"], \ - f"Expected attn_backend name to be either 'XFORMERS' or "\ - f"'FLASH_ATTN', but "\ - f"got '{self.runner.attn_backend.get_name()}'" - self._add_additonal_input_buffers_for_enc_dec_model( - attn_metadata=attn_metadata, input_buffers=input_buffers) - return input_buffers - - def prepare_graph_input_buffers( - self, - input_buffers, - attn_metadata, - is_encoder_decoder_model: bool = False) -> None: - input_buffers["seq_lens_tensor"].copy_( - attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True) - input_buffers["block_tables"].copy_( - attn_metadata.decode_metadata.block_tables, non_blocking=True) - if is_encoder_decoder_model: - # The encoder decoder model works only with XFormers and - # Flash Attention backend. Assert the same. - assert self.runner.attn_backend.get_name() in\ - ["XFORMERS", "FLASH_ATTN"], \ - f"Expected attn_backend name to be either 'XFORMERS' or "\ - f"'FLASH_ATTN', but "\ - f"got '{self.runner.attn_backend.get_name()}'" - self._prepare_input_buffers_for_enc_dec_model( - attn_metadata, input_buffers) - - def begin_forward(self, model_input) -> None: - return - - def _update_captured_metadata_for_enc_dec_model(self, batch_size: int, - attn_metadata): - """ - Updates the attention metadata parameters for CUDA graph capture in an - encoder-decoder model. - - This method modifies attention-related tensors and metadata required - for CUDA graph capture in encoder-decoder models. Specifically, it - updates the cross-attention and encoder sequence tensors in the - AttentionMetadata object. - """ - # During decode phase the cross_slot_mapping will be empty. Hence set - # an empty tensor for CUDA Graph capture. - attn_metadata.cross_slot_mapping = torch.tensor( - [], dtype=torch.int).cuda() - attn_metadata.cross_block_tables = torch.full( - (batch_size, self.runner.get_max_block_per_batch()), - 1, - dtype=torch.int).cuda() - attn_metadata.encoder_seq_lens = torch.full((batch_size, ), - 1, - dtype=torch.int).cuda() - attn_metadata.encoder_seq_lens_tensor = torch.full( - (batch_size, ), 1, dtype=torch.int).cuda() - attn_metadata.max_encoder_seq_len = self.runner.max_seq_len_to_capture - attn_metadata.num_encoder_tokens = 0 - - def _add_additonal_input_buffers_for_enc_dec_model( - self, attn_metadata, input_buffers: Dict[str, Any]): - """ - Saves additional input buffers specific to the encoder-decoder model - from the attention metadata. - - This method extracts and stores encoder-decoder related input buffers - from the `attn_metadata` into the `input_buffers` dictionary. The - buffers include encoder sequence lengths, cross-slot mappings, and - cross-block tables, which are essential for the encoder-decoder model - during CUDA graph replay. - """ - input_buffers["encoder_seq_lens_tensor"] = ( - attn_metadata.decode_metadata.encoder_seq_lens_tensor) - input_buffers["cross_slot_mapping"] = ( - attn_metadata.decode_metadata.cross_slot_mapping) - input_buffers["cross_block_tables"] = ( - attn_metadata.decode_metadata.cross_block_tables) - - def _prepare_input_buffers_for_enc_dec_model(self, attn_metadata, - input_buffers: Dict[str, - Any]): - """ - Populates input buffers with data from the encoder-decoder model's - attention metadata. - - This method fills the input buffers with encoder-decoder specific - tensors. It copies data from the `attn_metadata` and keyword arguments - (`kwargs`) into corresponding buffers in the `input_buffers` dictionary. - The copied data includes attention-related metadata as well as input - IDs and positional information for the encoder. - """ - input_buffers["encoder_seq_lens_tensor"].copy_( - attn_metadata.decode_metadata.encoder_seq_lens_tensor, - non_blocking=True) - input_buffers["cross_slot_mapping"].copy_( - attn_metadata.decode_metadata.cross_slot_mapping, - non_blocking=True) - input_buffers["cross_block_tables"].copy_( - attn_metadata.decode_metadata.cross_block_tables, - non_blocking=True) - - -def is_all_encoder_attn_metadata_set(attn_metadata): - ''' - All attention metadata required for encoder attention is set. - ''' - return ((attn_metadata.encoder_seq_lens is not None) - and (attn_metadata.encoder_seq_lens_tensor is not None) - and (attn_metadata.max_encoder_seq_len is not None)) - - -def is_all_cross_attn_metadata_set(attn_metadata): - ''' - All attention metadata required for enc/dec cross-attention is set. - - Superset of encoder attention required metadata. - ''' - return (attn_metadata.is_all_encoder_attn_metadata_set - and (attn_metadata.cross_slot_mapping is not None) - and (attn_metadata.cross_block_tables is not None)) - - -def get_seq_len_block_table_args( - attn_metadata, - is_prompt: bool, - attn_type: str, -) -> tuple: - ''' - The particular choice of sequence-length- and block-table-related - attributes which should be extracted from attn_metadata is dependent - on the type of attention operation. - - Decoder attn -> select entirely decoder self-attention-related fields - Encoder/decoder cross-attn -> select encoder sequence lengths & - cross-attn block-tables fields - Encoder attn -> select encoder sequence lengths fields & no block tables - - Arguments: - - * attn_metadata: Attention metadata structure associated with attention op - * is_prompt: True if prefill, False otherwise - * attn_type: encoder attention, decoder self-attention, - encoder/decoder cross-attention - - Returns: - - * Appropriate sequence-lengths tensor - * Appropriate max sequence-length scalar - * Appropriate block tables (or None) - ''' - - if attn_type == AttentionType.DECODER: - # Decoder self-attention - # Choose max_seq_len based on whether we are in prompt_run - if is_prompt: - max_seq_len = attn_metadata.max_prefill_seq_len - else: - max_seq_len = attn_metadata.max_decode_seq_len - return (attn_metadata.seq_lens_tensor, max_seq_len, - attn_metadata.block_tables) - elif attn_type == AttentionType.ENCODER_DECODER: - # Enc/dec cross-attention KVs match encoder sequence length; - # cross-attention utilizes special "cross" block tables - return (attn_metadata.encoder_seq_lens_tensor, - attn_metadata.max_encoder_seq_len, - attn_metadata.cross_block_tables) - elif attn_type == AttentionType.ENCODER: - # No block tables associated with encoder attention - return (attn_metadata.encoder_seq_lens_tensor, - attn_metadata.max_encoder_seq_len, None) - else: - raise AttributeError(f"Invalid attention type {str(attn_type)}") - - -def get_num_prefill_decode_query_kv_tokens( - attn_metadata, - attn_type: str, -) -> Tuple[int, int, int]: - """ - Calculate the number of prefill and decode tokens for query, key/value - based on the attention metadata and the specified attention type. - - Args: - attn_metadata (FlashAttentionMetadata): Attention Metadata object. - attn_type (AttentionType): The type of attention being used. - Returns: - Tuple[int, int, int]: A tuple containing three integers: - - The number of prefill query tokens. - - The number of prefill key/value tokens. - - The number of decode query tokens. - - Raises: - AssertionError: If the number of encoder tokens in `attn_metadata` - is `None` when required for the calculations. - """ - num_prefill_query_tokens = 0 - num_decode_query_tokens = 0 - num_prefill_kv_tokens = 0 - if attn_type == AttentionType.ENCODER: - # Encoder attention is only invoked during prefill phase. - # The same input servers a both query and key. - assert attn_metadata.num_encoder_tokens is not None - num_prefill_query_tokens = attn_metadata.num_encoder_tokens - num_prefill_kv_tokens = attn_metadata.num_encoder_tokens - num_decode_query_tokens = 0 - elif attn_type == AttentionType.ENCODER_DECODER: - assert attn_metadata.num_encoder_tokens is not None - num_prefill_query_tokens = attn_metadata.num_prefill_tokens - # The key is the encoder/cross-attention. - num_prefill_kv_tokens = attn_metadata.num_encoder_tokens - num_decode_query_tokens = attn_metadata.num_decode_tokens - else: # attn_type == AttentionType.DECODER or - # attn_type == AttentionType.ENCODER_ONLY - num_prefill_query_tokens = attn_metadata.num_prefill_tokens - num_prefill_kv_tokens = attn_metadata.num_prefill_tokens - num_decode_query_tokens = attn_metadata.num_decode_tokens - - return (num_prefill_query_tokens, num_prefill_kv_tokens, - num_decode_query_tokens) - - -def get_flash_attn_version(): - try: - from vllm.vllm_flash_attn.flash_attn_interface import ( - fa_version_unsupported_reason, is_fa_version_supported) - - # if hopper default to FA3, otherwise stick to FA2 for now - # TODO(lucas): profile FA3 on ampere to see if it makes sense to - # use FA3 as default for both - if current_platform.get_device_capability()[0] == 9: - fa_version = 3 if is_fa_version_supported(3) else 2 - else: - fa_version = 2 - - if envs.VLLM_FLASH_ATTN_VERSION is not None: - assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3] - fa_version = envs.VLLM_FLASH_ATTN_VERSION - if (current_platform.get_device_capability()[0] == 10 - and envs.VLLM_FLASH_ATTN_VERSION == 3): - logger.warning("Cannot use FA version 3 on Blackwell platform", - "defaulting to FA version 2.") - fa_version = 2 - - if not is_fa_version_supported(fa_version): - logger.error("Cannot use FA version %d is not supported due to %s", - fa_version, fa_version_unsupported_reason(fa_version)) - - assert is_fa_version_supported(fa_version) - return fa_version - except (ImportError, AssertionError): - return None diff --git a/fastvideo/v1/attention/layer.py b/fastvideo/v1/attention/layer.py deleted file mode 100644 index 2fa49c7b9..000000000 --- a/fastvideo/v1/attention/layer.py +++ /dev/null @@ -1,135 +0,0 @@ -import torch -import torch.nn as nn -from typing import Optional, Dict, Any, Type - -from fastvideo.v1.attention.flash_attn import LocalAttention, DistributedAttention - - -class AttentionBackend: - """Base class for attention backends that provides implementation classes.""" - - @staticmethod - def get_impl_cls() -> Type: - """Returns the implementation class for the attention backend.""" - raise NotImplementedError("Subclasses must implement get_impl_cls") - - -class LocalAttentionBackend(AttentionBackend): - """Backend for local attention implementation.""" - - @staticmethod - def get_impl_cls() -> Type: - """Returns the LocalAttention implementation class.""" - return LocalAttention - - -class DistributedAttentionBackend(AttentionBackend): - """Backend for distributed attention implementation.""" - - @staticmethod - def get_impl_cls() -> Type: - """Returns the DistributedAttention implementation class.""" - return DistributedAttention - - -class AttentionLayer(nn.Module): - """Base attention layer that uses an implementation class from a backend.""" - - def __init__( - self, - attn_backend: AttentionBackend, - dropout_rate: float = 0.0, - causal: bool = False, - softmax_scale: Optional[float] = None, - **extra_impl_args - ): - super().__init__() - impl_cls = attn_backend.get_impl_cls() - self.impl = impl_cls( - dropout_rate=dropout_rate, - causal=causal, - softmax_scale=softmax_scale, - **extra_impl_args - ) - - def forward(self, *args, **kwargs): - """Forward pass delegated to the implementation.""" - return self.impl(*args, **kwargs) - - -class LocalAttentionLayer(AttentionLayer): - """Wrapper for local attention implementation.""" - - def __init__( - self, - dropout_rate: float = 0.0, - causal: bool = False, - softmax_scale: Optional[float] = None, - **extra_impl_args - ): - super().__init__( - LocalAttentionBackend(), - dropout_rate=dropout_rate, - causal=causal, - softmax_scale=softmax_scale, - **extra_impl_args - ) - - def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """Apply local attention. - - Args: - q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] - k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] - v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] - - Returns: - torch.Tensor: Output tensor after attention - """ - return self.impl(q, k, v) - - -class DistributedAttentionLayer(AttentionLayer): - """Wrapper for distributed attention implementation.""" - - def __init__( - self, - dropout_rate: float = 0.0, - causal: bool = False, - softmax_scale: Optional[float] = None, - **extra_impl_args - ): - super().__init__( - DistributedAttentionBackend(), - dropout_rate=dropout_rate, - causal=causal, - softmax_scale=softmax_scale, - **extra_impl_args - ) - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - replicated_q: Optional[torch.Tensor] = None, - replicated_k: Optional[torch.Tensor] = None, - replicated_v: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Forward pass for distributed attention. - - Args: - q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] - k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] - v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] - replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens - replicated_k (Optional[torch.Tensor]): Replicated key tensor - replicated_v (Optional[torch.Tensor]): Replicated value tensor - - Returns: - Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - - o (torch.Tensor): Output tensor after attention for the main sequence - - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided - """ - return self.impl(q, k, v, replicated_q, replicated_k, replicated_v) -