# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer ROCm GPUs.""" import itertools from dataclasses import dataclass from functools import cache from typing import List, Optional, Tuple, Type import torch import vllm.envs as envs from vllm import _custom_ops as ops from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionLayer, AttentionMetadata, AttentionType) from vllm.attention.backends.utils import (CommonAttentionState, CommonMetadataBuilder) from vllm.attention.ops.paged_attn import (PagedAttention, PagedAttentionMetadata) from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape) from vllm.platforms import current_platform logger = init_logger(__name__) _PARTITION_SIZE_ROCM = 256 @cache def is_rocm_aiter_paged_attn_enabled() -> bool: return envs.VLLM_ROCM_USE_AITER_PAGED_ATTN \ and envs.VLLM_ROCM_USE_AITER \ @cache def _get_paged_attn_module() -> PagedAttention: """ Initializes the appropriate PagedAttention module from `attention/ops`, which is used as helper function by `ROCmFlashAttentionImpl` and `ROCmFlashAttentionBackend`. The choice of attention module depends on whether AITER paged attention is enabled: - If enabled, `ROCmFlashAttentionImpl` uses `AITERPagedAttention`. - Otherwise, it defaults to using the original `PagedAttention`. """ if is_rocm_aiter_paged_attn_enabled(): # Import AITERPagedAttention only when the flag is enabled from vllm.attention.ops.rocm_aiter_paged_attn import ( AITERPagedAttention) return AITERPagedAttention() return PagedAttention() class ROCmFlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_name() -> str: return "ROCM_FLASH" @staticmethod def get_impl_cls() -> Type["ROCmFlashAttentionImpl"]: return ROCmFlashAttentionImpl @staticmethod def get_metadata_cls() -> Type["AttentionMetadata"]: return ROCmFlashAttentionMetadata @staticmethod def get_builder_cls() -> Type["ROCmFlashAttentionMetadataBuilder"]: return ROCmFlashAttentionMetadataBuilder @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, ...]: paged_attn = _get_paged_attn_module() return paged_attn.get_kv_cache_shape(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: paged_attn = _get_paged_attn_module() paged_attn.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst) @staticmethod def copy_blocks( kv_caches: List[torch.Tensor], src_to_dists: torch.Tensor, ) -> None: paged_attn = _get_paged_attn_module() paged_attn.copy_blocks(kv_caches, src_to_dists) @dataclass class ROCmFlashAttentionMetadata(AttentionMetadata, PagedAttentionMetadata): """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] # 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 # 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 # 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 query length in the batch. None for decoding. max_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 # (batch_size,) A tensor of context lengths (tokens that are computed # so far). context_lens_tensor: Optional[torch.Tensor] = None # Max number of query tokens among request in the batch. max_decode_query_len: Optional[int] = None _cached_prefill_metadata: Optional["ROCmFlashAttentionMetadata"] = None _cached_decode_metadata: Optional["ROCmFlashAttentionMetadata"] = 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 # 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 prefill_metadata(self) -> Optional["ROCmFlashAttentionMetadata"]: 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 assert self.seq_lens_tensor is not None assert self.block_tables is not None self._cached_prefill_metadata = ROCmFlashAttentionMetadata( num_prefills=self.num_prefills, num_prefill_tokens=self.num_prefill_tokens, num_decode_tokens=0, slot_mapping=self.slot_mapping[:self.num_prefill_tokens], multi_modal_placeholder_index_maps=self. multi_modal_placeholder_index_maps, enable_kv_scales_calculation=self.enable_kv_scales_calculation, seq_lens=self.seq_lens[:self.num_prefills], seq_lens_tensor=self.seq_lens_tensor[:self.num_prefills], max_query_len=self.max_query_len, max_prefill_seq_len=self.max_prefill_seq_len, max_decode_seq_len=0, query_start_loc=None if self.query_start_loc is None else self.query_start_loc[:self.num_prefills + 1], 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=self.block_tables[:self.num_prefills], 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, 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["ROCmFlashAttentionMetadata"]: if self.num_decode_tokens == 0: return None if self._cached_decode_metadata is not None: return self._cached_decode_metadata assert self.block_tables is not None assert self.seq_lens_tensor is not None self._cached_decode_metadata = ROCmFlashAttentionMetadata( num_prefills=0, num_prefill_tokens=0, num_decode_tokens=self.num_decode_tokens, slot_mapping=self.slot_mapping[self.num_prefill_tokens:], multi_modal_placeholder_index_maps=None, enable_kv_scales_calculation=True, seq_lens=None, seq_lens_tensor=self.seq_lens_tensor[self.num_prefills:], max_query_len=None, max_prefill_seq_len=0, max_decode_seq_len=self.max_decode_seq_len, query_start_loc=None, seq_start_loc=None, context_lens_tensor=None, block_tables=self.block_tables[self.num_prefills:], 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, max_encoder_seq_len=self.max_encoder_seq_len, cross_slot_mapping=self.cross_slot_mapping, cross_block_tables=self.cross_block_tables) # Batch may be composed of prefill|decodes, adjust query start indices # to refer to the start of decodes when the two are split apart. # E.g. in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6]. if self._cached_decode_metadata.query_start_loc is not None: qs = self._cached_decode_metadata.query_start_loc self._cached_decode_metadata.query_start_loc = qs - qs[0] return self._cached_decode_metadata class ROCmFlashAttentionMetadataBuilder( CommonMetadataBuilder[ROCmFlashAttentionMetadata]): _metadata_cls = ROCmFlashAttentionMetadata def _make_alibi_bias(alibi_slopes: torch.Tensor, dtype: torch.dtype, seq_lens: Optional[List[int]], make_attn_mask: bool = True) -> List[torch.Tensor]: attn_biases = [] if seq_lens: for seq_len in seq_lens: bias = torch.arange(seq_len, dtype=dtype) # NOTE(zhuohan): HF uses # `bias = bias[None, :].repeat(seq_len, 1)` # here. We find that both biases give the same results, but # the bias below more accurately follows the original ALiBi # paper. bias = bias[None, :] - bias[:, None] num_heads = alibi_slopes.shape[0] bias = bias[None, :].repeat( (num_heads, 1, 1)).to(alibi_slopes.device) bias.mul_(alibi_slopes[:, None, None]) if make_attn_mask: inf_mask = torch.empty( (1, seq_len, seq_len), dtype=bias.dtype).fill_(-torch.inf).triu_(diagonal=1).to( alibi_slopes.device) attn_biases.append((bias + inf_mask).to(dtype)) else: attn_biases.append(bias.to(dtype)) return attn_biases def _get_seq_len_block_table_args( attn_metadata: ROCmFlashAttentionMetadata, attn_type: str, ) -> tuple: ''' The particular choice of sequence-length 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 Encoder attn -> select encoder sequence lengths fields Encoder-only attn -> select prefill sequence lengths with bidirectional attention Arguments: * attn_metadata: Attention metadata structure associated with attention op * attn_type: encoder attention, decoder self-attention, encoder/decoder cross-attention, encoder-only Returns: * Appropriate sequence-lengths tensors for query and key * Appropriate max sequence-length scalar * Causal masking flag ''' if attn_type == AttentionType.ENCODER: assert attn_metadata.encoder_seq_lens is not None assert attn_metadata.encoder_seq_lens_tensor is not None query_seq_start_loc = torch.tensor( list(itertools.accumulate([0] + attn_metadata.encoder_seq_lens)), device=attn_metadata.encoder_seq_lens_tensor.device, dtype=attn_metadata.encoder_seq_lens_tensor.dtype) causal_mask = False # No block tables associated with encoder attention return (query_seq_start_loc, attn_metadata.max_encoder_seq_len, query_seq_start_loc, attn_metadata.max_encoder_seq_len, attn_metadata.encoder_seq_lens, causal_mask) elif attn_type == AttentionType.ENCODER_ONLY: # For encoder-only models, we use the prefill sequence lengths assert attn_metadata.seq_lens is not None assert attn_metadata.seq_lens_tensor is not None query_seq_start_loc = torch.tensor( list(itertools.accumulate([0] + attn_metadata.seq_lens)), device=attn_metadata.seq_lens_tensor.device, dtype=attn_metadata.seq_lens_tensor.dtype) max_seq_len = attn_metadata.max_prefill_seq_len # Encoder-only models typically use bidirectional attention causal_mask = False return (query_seq_start_loc, max_seq_len, query_seq_start_loc, max_seq_len, attn_metadata.seq_lens, causal_mask) elif attn_type == AttentionType.DECODER: # Decoder self-attention # Choose max_seq_len based on whether we are in prompt_run assert attn_metadata.seq_lens is not None assert attn_metadata.seq_lens_tensor is not None query_seq_start_loc = torch.tensor( list(itertools.accumulate([0] + attn_metadata.seq_lens)), device=attn_metadata.seq_lens_tensor.device, dtype=attn_metadata.seq_lens_tensor.dtype) max_seq_len = attn_metadata.max_prefill_seq_len causal_mask = True return (query_seq_start_loc, max_seq_len, query_seq_start_loc, max_seq_len, attn_metadata.seq_lens, causal_mask) elif attn_type == AttentionType.ENCODER_DECODER: assert attn_metadata.seq_lens is not None assert attn_metadata.encoder_seq_lens_tensor is not None query_start_loc = torch.tensor( list(itertools.accumulate([0] + attn_metadata.seq_lens)), device=attn_metadata.encoder_seq_lens_tensor.device, dtype=attn_metadata.encoder_seq_lens_tensor.dtype) assert attn_metadata.encoder_seq_lens is not None assert attn_metadata.seq_lens_tensor is not None key_seq_start_loc = torch.tensor( list(itertools.accumulate([0] + attn_metadata.encoder_seq_lens)), device=attn_metadata.seq_lens_tensor.device, dtype=attn_metadata.seq_lens_tensor.dtype) causal_mask = False # Enc/dec cross-attention KVs match encoder sequence length; # cross-attention utilizes special "cross" block tables return (query_start_loc, attn_metadata.max_prefill_seq_len, key_seq_start_loc, attn_metadata.max_encoder_seq_len, attn_metadata.seq_lens, causal_mask) else: raise AttributeError(f"Invalid attention type {str(attn_type)}") class ROCmFlashAttentionImpl(AttentionImpl): """ If the input tensors contain prompt tokens, the layout is as follows: |<--------------- num_prompt_tokens -------------->| |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->| Otherwise, the layout is as follows: |<------------------ num_generation_tokens (M) ----------------->| |<--generation_0-->|..........|<--generation_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 ----------->| |<-prompt_0->|...|<-prompt_N-1->|<-generation_0->|...|<-generation_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, logits_soft_cap: Optional[float] = None, attn_type: str = AttentionType.DECODER, kv_sharing_target_layer_name: Optional[str] = None, use_irope: bool = False, ) -> None: if kv_sharing_target_layer_name is not None: raise NotImplementedError("KV sharing is not supported in V0 " "ROCM_FLASH backend.") if use_irope: logger.warning_once( "Using irope in ROCm Flash Attention is not supported yet, it " "will fail back to global attention for long context.") if use_irope: logger.warning( "Using irope in V0 is not supported yet, it will fall back " "to global attention for long context.") if logits_soft_cap is None: # In flash-attn, setting logits_soft_cap as 0 means no soft cap. self.logits_soft_cap = 0.0 else: self.logits_soft_cap = logits_soft_cap self.attn_type = attn_type 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, sliding_window) if sliding_window is not None else (-1, -1)) self.kv_cache_dtype = kv_cache_dtype self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.paged_attn_module = _get_paged_attn_module() supported_head_sizes = self.paged_attn_module.get_supported_head_sizes( ) if head_size not in supported_head_sizes: raise ValueError( f"Head size {head_size} is not supported by PagedAttention. " f"Supported head sizes are: {supported_head_sizes}.") self.use_naive_attn = False # NOTE: Allow for switching between Triton and CK. Defaulting to triton. self.use_triton_flash_attn = envs.VLLM_USE_TRITON_FLASH_ATTN if self.use_triton_flash_attn: if logits_soft_cap is not None: raise ValueError( "ROCm Triton FlashAttention does not support attention" " logits soft capping." " please try using the ROCm CK " "FA backend instead by setting the env var " "`VLLM_USE_TRITON_FLASH_ATTN=0`") from vllm.attention.ops.triton_flash_attention import ( # noqa: F401 triton_attention) self.triton_attn_func = triton_attention logger.debug("Using Triton FA in ROCmBackend") if self.sliding_window != (-1, -1): logger.warning("ROCm Triton FA does not currently support " "sliding window attention. If using half " "precision, please try using the ROCm CK " "FA backend instead by setting the env var " "`VLLM_USE_TRITON_FLASH_ATTN=0`") else: # if not using triton, navi3x/navi21/navi10 do not use flash-attn # either if not current_platform.has_device_capability(90): self.use_naive_attn = True else: try: from flash_attn import flash_attn_varlen_func # noqa: F401 self.fa_attn_func = flash_attn_varlen_func logger.debug("Using CK FA in ROCmBackend") except ModuleNotFoundError: self.use_naive_attn = True if self.use_naive_attn: if logits_soft_cap is not None: raise ValueError( "ROCm Naive FlashAttention does not support " "attention logits soft capping.") self.sdpa_attn_func = _sdpa_attention logger.debug("Using naive (SDPA) attention in ROCmBackend") self.aiter_kv_scales_initialized = False self.force_fp8_attention = ( get_current_vllm_config() is not None and get_current_vllm_config().model_config.override_attention_dtype == "fp8") def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor: """torch.repeat_interleave(x, dim=1, repeats=n_rep)""" tokens, n_kv_heads, head_dim = x.shape return (x[:, :, None, :].expand(tokens, n_kv_heads, n_rep, head_dim).reshape(tokens, n_kv_heads * n_rep, head_dim)) def fused_output_quant_supported(self, dtype: torch.dtype, static: bool, group_shape: GroupShape): if self.use_triton_flash_attn: return dtype == current_platform.fp8_dtype( ) and static and group_shape == GroupShape.PER_TENSOR # Only supported in the Triton backend return False def forward( self, layer: AttentionLayer, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, kv_cache: torch.Tensor, attn_metadata: ROCmFlashAttentionMetadata, output: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward pass with FlashAttention and PagedAttention. For decoder-only models: query, key and value must be non-None. For encoder/decoder models: * ROCmFlashAttentionImpl.forward() may be invoked for both self- and cross-attention layers. * For self-attention: query, key and value must be non-None. * For cross-attention: * Query must be non-None * During prefill, key and value must be non-None; key and value get cached for use during decode. * During decode, key and value may be None, since: (1) key and value tensors were cached during prefill, and (2) cross-attention key and value tensors do not grow during decode A note on how the attn_type (attention type enum) argument impacts attention forward() behavior: * DECODER: normal decoder-only behavior; use decoder self-attention block table * ENCODER: no KV caching; pass encoder sequence attributes (encoder_seq_lens/encoder_seq_lens_tensor/ max_encoder_seq_len) to kernel, in lieu of decoder sequence attributes (seq_lens/seq_lens_tensor/max_seq_len) * ENCODER_DECODER: cross-attention behavior; use cross-attention block table for caching KVs derived from encoder hidden states; since KV sequence lengths will match encoder sequence lengths, pass encoder sequence attributes to kernel (encoder_seq_lens/encoder_seq_lens_tensor/ max_encoder_seq_len) * ENCODER_ONLY: bidirectional attention with no KV caching; use prefill sequence attributes 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] 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. attn_type: Select attention type, between encoder attention, decoder self-attention, or encoder/decoder cross- attention. Defaults to decoder self-attention, which is the vLLM default generally Returns: shape = [num_tokens, num_heads * head_size] """ assert output is not None, "Output tensor must be provided." if output_scale is not None and not self.use_triton_flash_attn: raise NotImplementedError( "fused output quantization only supported for Triton" " implementation in ROCMFlashAttentionImpl for now") query = query.view(-1, self.num_heads, self.head_size) if key is not None: assert value is not None key = key.view(-1, self.num_kv_heads, self.head_size) value = value.view(-1, self.num_kv_heads, self.head_size) else: assert value is None paged_attn = self.paged_attn_module # Reshaping kv tensors is required for AITER paged attention kernel # because it works on a different tensor shape, # when the size of one element is one byte (int8/fp8 dtypes). # This reshaping is only required on the first forward call # and the kv cache must not be empty. if (is_rocm_aiter_paged_attn_enabled() and kv_cache.dtype.itemsize == 1 and not self.aiter_kv_scales_initialized and kv_cache.shape != torch.Size([0])): num_blocks = kv_cache.shape[1] block_size = kv_cache.shape[2] // (self.num_kv_heads * self.head_size) k_scale = torch.empty((self.num_kv_heads, num_blocks * block_size), dtype=torch.float32, device=kv_cache.device) v_scale = torch.empty((self.num_kv_heads, num_blocks * block_size), dtype=torch.float32, device=kv_cache.device) self.aiter_kv_scales_initialized = True k_scale.fill_(layer._k_scale.item()) v_scale.fill_(layer._v_scale.item()) layer._k_scale = k_scale layer._v_scale = v_scale # Only update KV cache for decoder self-attention # and encoder-decoder cross-attention if self.attn_type not in [ AttentionType.ENCODER, AttentionType.ENCODER_ONLY ] and kv_cache.numel() > 0: key_cache, value_cache = paged_attn.split_kv_cache( kv_cache, self.num_kv_heads, self.head_size) if key is not None and value is not None: # 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. paged_attn.write_to_paged_cache( key, value, key_cache, value_cache, attn_metadata.slot_mapping if self.attn_type != AttentionType.ENCODER_DECODER else attn_metadata.cross_slot_mapping, self.kv_cache_dtype, layer._k_scale, layer._v_scale, ) if self.attn_type != AttentionType.ENCODER: num_prefill_tokens = attn_metadata.num_prefill_tokens elif self.attn_type == AttentionType.ENCODER_ONLY: # For encoder-only models, all tokens are processed in one go num_prefill_tokens = query.shape[0] else: assert attn_metadata.num_encoder_tokens is not None num_prefill_tokens = attn_metadata.num_encoder_tokens # Query for decode. KV is not needed because it is already cached. decode_query = query[num_prefill_tokens:] # QKV for prefill. query = query[:num_prefill_tokens] # For encoder-only and encoder models, # we process all tokens at once # For decoder and encoder-decoder, # we may need to limit key/value to prefill tokens if key is not None and value is not None \ and self.attn_type not in [AttentionType.ENCODER_DECODER, AttentionType.ENCODER_ONLY]: key = key[:num_prefill_tokens] value = value[:num_prefill_tokens] if prefill_meta := attn_metadata.prefill_metadata: # Prompt run. # normal attention and DECODER if self.attn_type == AttentionType.DECODER and ( kv_cache.numel() == 0 or prefill_meta.block_tables is None or prefill_meta.block_tables.numel() == 0): (query_seq_start_loc, query_max_seq_len, key_seq_start_loc, key_max_seq_len, seq_lens, causal_mask) = (prefill_meta.seq_start_loc, prefill_meta.max_prefill_seq_len, prefill_meta.seq_start_loc, prefill_meta.max_prefill_seq_len, attn_metadata.seq_lens, True) # prefix-enabled attention and ENCODER/ENCODER_DECODER else: (query_seq_start_loc, query_max_seq_len, key_seq_start_loc, key_max_seq_len, seq_lens, causal_mask) = _get_seq_len_block_table_args( prefill_meta, self.attn_type) # Prompt run. if kv_cache.numel() == 0 or prefill_meta.block_tables.numel() == 0: # triton attention # When block_tables are not filled, it means q and k are the # prompt, and they have the same length. attn_masks = None if self.use_triton_flash_attn: if self.alibi_slopes is not None: attn_masks = _make_alibi_bias( self.alibi_slopes, query.dtype, seq_lens, make_attn_mask=causal_mask) # type: ignore use_fp8_scales = (layer._q_scale and layer._k_scale and layer._v_scale and layer._prob_scale and (self.kv_cache_dtype == "fp8" or self.force_fp8_attention)) full_scales = ( layer._q_scale.item(), layer._k_scale.item(), layer._v_scale.item(), layer._prob_scale.item()) if use_fp8_scales else None self.triton_attn_func( query, key, value, output[:num_prefill_tokens], query_seq_start_loc, key_seq_start_loc, query_max_seq_len, key_max_seq_len, causal_mask, self.scale, attn_masks[0][None] if attn_masks is not None else None, full_scales, output_scale, ) elif self.use_naive_attn: if self.num_kv_heads != self.num_heads: # Interleave for MQA workaround. key = self.repeat_kv(key, self.num_queries_per_kv) value = self.repeat_kv(value, self.num_queries_per_kv) if self.alibi_slopes is not None: attn_masks = _make_alibi_bias( self.alibi_slopes, query.dtype, attn_metadata.seq_lens, make_attn_mask=causal_mask) # type: ignore query = query.movedim(0, query.dim() - 2) key = key.movedim(0, key.dim() - 2) value = value.movedim(0, value.dim() - 2) # sdpa math backend attention self.sdpa_attn_func( query, key, value, output[:num_prefill_tokens], query_seq_start_loc, num_prefill_tokens, self.num_heads, self.head_size, self.scale, attn_masks, ) else: # upstream FA does not support an output arg, copy output[:num_prefill_tokens] = self.fa_attn_func( q=query, k=key, v=value, cu_seqlens_q=query_seq_start_loc, cu_seqlens_k=key_seq_start_loc, max_seqlen_q=prefill_meta.max_prefill_seq_len, max_seqlen_k=key_max_seq_len, softmax_scale=self.scale, causal=causal_mask, window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, softcap=self.logits_soft_cap, ) else: # prefix-enabled attention - # not applicable for encoder-only models if self.attn_type != AttentionType.ENCODER_ONLY: output[:num_prefill_tokens] = paged_attn.forward_prefix( query, key, value, self.kv_cache_dtype, key_cache, value_cache, prefill_meta.block_tables, prefill_meta.query_start_loc, prefill_meta.seq_lens_tensor, prefill_meta.max_query_len, self.alibi_slopes, self.sliding_window[0], layer._k_scale, layer._v_scale, ) # Skip decode phase for encoder-only models if (decode_meta := attn_metadata.decode_metadata) and ( self.attn_type != AttentionType.ENCODER_ONLY): # Decoding run. # Whether to use rocm custom paged attention or not num_seqs, num_heads, head_size = decode_query.shape block_size = value_cache.shape[3] gqa_ratio = num_heads // self.num_kv_heads from vllm.platforms.rocm import use_rocm_custom_paged_attention use_custom = use_rocm_custom_paged_attention( decode_query.dtype, head_size, block_size, gqa_ratio, decode_meta.max_decode_seq_len, self.sliding_window, self.kv_cache_dtype, self.alibi_slopes) if use_custom: max_seq_len = (decode_meta.max_decode_seq_len if self.attn_type != AttentionType.ENCODER_DECODER else decode_meta.max_encoder_seq_len) assert max_seq_len is not None max_num_partitions = ( (max_seq_len + _PARTITION_SIZE_ROCM - 1) // _PARTITION_SIZE_ROCM) assert _PARTITION_SIZE_ROCM % block_size == 0 tmp_output = torch.empty( size=(num_seqs, num_heads, max_num_partitions, head_size), dtype=query.dtype, device=output.device, ) exp_sums = torch.empty( size=(num_seqs, num_heads, max_num_partitions), dtype=torch.float32, device=output.device, ) max_logits = torch.empty_like(exp_sums) query_start_loc = None ops.paged_attention_rocm( output[num_prefill_tokens:], exp_sums, max_logits, tmp_output, decode_query, key_cache, value_cache, self.num_kv_heads, self.scale, decode_meta.block_tables if self.attn_type != AttentionType.ENCODER_DECODER else decode_meta.cross_block_tables, decode_meta.seq_lens_tensor if self.attn_type != AttentionType.ENCODER_DECODER else decode_meta.encoder_seq_lens_tensor, query_start_loc, block_size, max_seq_len, self.alibi_slopes, self.kv_cache_dtype, layer._k_scale, layer._v_scale, output_scale, ) else: # PagedAttention does not support fused quant, manually quantize if output_scale is None: out_pa = output[num_prefill_tokens:] else: out_pa = torch.empty_like(output[num_prefill_tokens:], dtype=query.dtype) out_pa[:] = paged_attn.forward_decode( decode_query, key_cache, value_cache, decode_meta.block_tables if self.attn_type != AttentionType.ENCODER_DECODER else decode_meta.cross_block_tables, decode_meta.seq_lens_tensor if self.attn_type != AttentionType.ENCODER_DECODER else decode_meta.encoder_seq_lens_tensor, decode_meta.max_decode_seq_len if self.attn_type != AttentionType.ENCODER_DECODER else decode_meta.max_encoder_seq_len, self.kv_cache_dtype, self.num_kv_heads, self.scale, self.alibi_slopes, layer._k_scale, layer._v_scale, ) # Manually perform quantization if output_scale is not None: out_uq = out_pa.view(-1, self.num_heads * self.head_size) out_q = output.view(-1, self.num_heads * self.head_size) ops.scaled_fp8_quant(out_uq, output_scale, output=out_q[num_prefill_tokens:]) # Reshape the output tensor. return output.view(-1, self.num_heads * self.head_size) def _sdpa_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, seq_lens: torch.Tensor, num_tokens: int, num_heads: int, head_size: int, scale: float, attn_masks: Optional[List[torch.Tensor]] = None, ) -> torch.Tensor: start = 0 assert output.shape == (num_tokens, num_heads, head_size) assert output.dtype == query.dtype assert output.device == query.device for i, seq_len in enumerate(seq_lens): end = start + seq_len with torch.nn.attention.sdpa_kernel( torch.nn.attention.SDPBackend.MATH): sub_out = torch.nn.functional.scaled_dot_product_attention( query[:, start:end, :], key[:, start:end, :], value[:, start:end, :], dropout_p=0.0, is_causal=attn_masks is None, attn_mask=attn_masks[i] if attn_masks else None, scale=scale).movedim(query.dim() - 2, 0) output[start:end, :, :] = sub_out start = end return output