| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention layer with FlashAttention.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.logger import init_logger |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.utils.torch_utils import ( |
| 10 | canonicalize_singleton_dim_strides, |
| 11 | is_quantized_kv_cache, |
| 12 | ) |
| 13 | from vllm.v1.attention.backend import AttentionType |
| 14 | from vllm.v1.attention.backends.fa_utils import ( |
| 15 | get_flash_attn_version, |
| 16 | is_flash_attn_varlen_func_available, |
| 17 | ) |
| 18 | from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( |
| 19 | triton_reshape_and_cache_flash_diffkv, |
| 20 | ) |
| 21 | |
| 22 | if is_flash_attn_varlen_func_available(): |
| 23 | from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func |
| 24 | |
| 25 | from .flash_attn import ( |
| 26 | FlashAttentionBackend, |
| 27 | FlashAttentionImpl, |
| 28 | FlashAttentionMetadata, |
| 29 | cascade_attention, |
| 30 | ) |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | |
| 35 | class FlashAttentionDiffKVBackend(FlashAttentionBackend): |
| 36 | # Default to 128 for this backend |
| 37 | head_size_v: int = 128 |
| 38 | |
| 39 | @classmethod |
| 40 | def set_head_size_v(cls, head_size_v: int) -> None: |
| 41 | cls.head_size_v = head_size_v |
| 42 | |
| 43 | @classmethod |
| 44 | def is_supported_on_current_device( |
| 45 | cls, |
| 46 | head_size: int, |
| 47 | head_size_v: int, |
| 48 | has_sinks: bool, |
| 49 | ) -> bool: |
| 50 | """Check whether FA3/4 with this DiffKV config is usable here. |
| 51 | |
| 52 | DiffKV (hdim_qk != hdim_v) requires FA3 or FA4 |
| 53 | """ |
| 54 | if not is_flash_attn_varlen_func_available(): |
| 55 | return False |
| 56 | try: |
| 57 | version = get_flash_attn_version( |
| 58 | requires_alibi=False, |
| 59 | head_size=head_size, |
| 60 | head_size_v=head_size_v, |
| 61 | has_sinks=has_sinks, |
| 62 | ) |
| 63 | except Exception: |
| 64 | return False |
| 65 | return version in (3, 4) |
| 66 | |
| 67 | @staticmethod |
| 68 | def get_name() -> str: |
| 69 | return "FLASH_ATTN_DIFFKV" |
| 70 | |
| 71 | @staticmethod |
| 72 | def get_impl_cls() -> type["FlashAttentionImpl"]: |
| 73 | return FlashAttentionDiffKVImpl |
| 74 | |
| 75 | |
| 76 | class FlashAttentionDiffKVImpl(FlashAttentionImpl): |
| 77 | vllm_flash_attn_version: int | None |
| 78 | |
| 79 | def __init__(self, *args, **kwargs) -> None: |
| 80 | super().__init__(*args, **kwargs) |
| 81 | # Re-derive the FA version with diff-kv context so that |
| 82 | # get_flash_attn_version can apply the FA3 -> FA4 upgrade rule |
| 83 | # for sinks + hdim != hdim_v. |
| 84 | self.vllm_flash_attn_version = get_flash_attn_version( |
| 85 | requires_alibi=self.alibi_slopes is not None, |
| 86 | head_size=self.head_size, |
| 87 | head_size_v=FlashAttentionDiffKVBackend.head_size_v, |
| 88 | has_sinks=self.sinks is not None, |
| 89 | ) |
| 90 | self.fa4_hd256 = False |
| 91 | |
| 92 | def do_kv_cache_update( |
| 93 | self, |
| 94 | layer: torch.nn.Module, |
| 95 | key: torch.Tensor, |
| 96 | value: torch.Tensor, |
| 97 | kv_cache: torch.Tensor, |
| 98 | slot_mapping: torch.Tensor, |
| 99 | ) -> None: |
| 100 | if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 101 | # For encoder attention, |
| 102 | # we use direct Q, K, V tensors without caching |
| 103 | return |
| 104 | |
| 105 | # DiffKV packs K and V into a single tensor along the last dim: |
| 106 | # kv_cache shape: [num_blocks, block_size, num_kv_heads, |
| 107 | # head_size_k + head_size_v] |
| 108 | # (B, H, N, C) -> (B, N, H, C) for kernel compatibility. |
| 109 | triton_reshape_and_cache_flash_diffkv( |
| 110 | key, |
| 111 | value, |
| 112 | kv_cache.transpose(1, 2), |
| 113 | slot_mapping, |
| 114 | self.kv_cache_dtype, |
| 115 | layer._k_scale, |
| 116 | layer._v_scale, |
| 117 | ) |
| 118 | |
| 119 | def forward( |
| 120 | self, |
| 121 | layer: torch.nn.Module, |
| 122 | query: torch.Tensor, |
| 123 | key: torch.Tensor, |
| 124 | value: torch.Tensor, |
| 125 | kv_cache: torch.Tensor, |
| 126 | attn_metadata: FlashAttentionMetadata, |
| 127 | output: torch.Tensor, |
| 128 | output_scale: torch.Tensor | None = None, |
| 129 | output_block_scale: torch.Tensor | None = None, |
| 130 | ) -> torch.Tensor: |
| 131 | """Forward pass with FlashAttention. |
| 132 | |
| 133 | Args: |
| 134 | layer: The attention layer, providing the q/k/v quantization scales. |
| 135 | query: shape = [num_tokens, num_heads, head_size] |
| 136 | key: shape = [num_tokens, num_kv_heads, head_size] |
| 137 | value: shape = [num_tokens, num_kv_heads, head_size_v] |
| 138 | kv_cache: shape = |
| 139 | [num_blocks, block_size, num_kv_heads, head_size + head_size_v] |
| 140 | attn_metadata: Metadata for attention. |
| 141 | output: Tensor that the attention result is written into. |
| 142 | output_scale: Scale for fused output quantization; not supported |
| 143 | by this backend. |
| 144 | output_block_scale: Block scale for fused output quantization; |
| 145 | not supported by this backend. |
| 146 | |
| 147 | Returns: |
| 148 | shape = [num_tokens, num_heads * head_size_v] |
| 149 | NOTE: FP8 quantization, flash-attn expect the size of |
| 150 | {q,k,v}_descale to be (num_sequences, num_kv_heads). |
| 151 | We use torch's .expand() to avoid duplicating values |
| 152 | |
| 153 | """ |
| 154 | assert self.vllm_flash_attn_version is not None, ( |
| 155 | "FlashAttention version not detected." |
| 156 | ) |
| 157 | |
| 158 | if output_scale is not None or output_block_scale is not None: |
| 159 | raise NotImplementedError( |
| 160 | "fused output quantization is not yet supported for FlashAttentionImpl" |
| 161 | ) |
| 162 | |
| 163 | if attn_metadata is None: |
| 164 | # Profiling run. |
| 165 | return output.fill_(0) |
| 166 | |
| 167 | attn_type = self.attn_type |
| 168 | |
| 169 | # IMPORTANT! |
| 170 | # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in |
| 171 | # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead |
| 172 | # in this method. For example, `view` and `slice` (or `[:n]`) operations |
| 173 | # are surprisingly slow even in the case they do not invoke any GPU ops. |
| 174 | # Minimize the PyTorch ops in this method as much as possible. |
| 175 | # Whenever making a change in this method, please benchmark the |
| 176 | # performance to make sure it does not introduce any overhead. |
| 177 | |
| 178 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 179 | |
| 180 | # Handle encoder attention differently - no KV cache needed |
| 181 | if attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 182 | # For encoder attention, |
| 183 | # we use direct Q, K, V tensors without caching |
| 184 | return self._forward_encoder_attention( |
| 185 | query[:num_actual_tokens], |
| 186 | key[:num_actual_tokens], |
| 187 | value[:num_actual_tokens], |
| 188 | output[:num_actual_tokens], |
| 189 | attn_metadata, |
| 190 | layer, |
| 191 | ) |
| 192 | |
| 193 | # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) |
| 194 | key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) |
| 195 | # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). |
| 196 | # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. |
| 197 | # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides. |
| 198 | fixed_k = canonicalize_singleton_dim_strides(key_cache) |
| 199 | fixed_v = canonicalize_singleton_dim_strides(value_cache) |
| 200 | if fixed_k is not key_cache or fixed_v is not value_cache: |
| 201 | logger.debug( |
| 202 | "Canonicalized degenerate KV cache strides (FlashAttentionDiffKV): " |
| 203 | "shape=%s, key strides before=%s after=%s, " |
| 204 | "value strides before=%s after=%s", |
| 205 | key_cache.shape, |
| 206 | key_cache.stride(), |
| 207 | fixed_k.stride(), |
| 208 | value_cache.stride(), |
| 209 | fixed_v.stride(), |
| 210 | ) |
| 211 | key_cache, value_cache = fixed_k, fixed_v |
| 212 | |
| 213 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 214 | # queries are quantized in the attention layer |
| 215 | key_cache = key_cache.view(current_platform.fp8_dtype()) |
| 216 | value_cache = value_cache.view(current_platform.fp8_dtype()) |
| 217 | |
| 218 | if not attn_metadata.use_cascade: |
| 219 | cu_seqlens_q = attn_metadata.query_start_loc |
| 220 | seqused_k = attn_metadata.seq_lens |
| 221 | max_seqlen_q = attn_metadata.max_query_len |
| 222 | max_seqlen_k = attn_metadata.max_seq_len |
| 223 | block_table = attn_metadata.block_table |
| 224 | scheduler_metadata = attn_metadata.scheduler_metadata |
| 225 | |
| 226 | descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) |
| 227 | |
| 228 | if self.dcp_world_size > 1: |
| 229 | self._forward_with_dcp( |
| 230 | query[:num_actual_tokens], |
| 231 | key[:num_actual_tokens], |
| 232 | value[:num_actual_tokens], |
| 233 | key_cache, |
| 234 | value_cache, |
| 235 | output[:num_actual_tokens], |
| 236 | attn_metadata, |
| 237 | q_descale=layer._q_scale.expand(descale_shape), |
| 238 | k_descale=layer._k_scale.expand(descale_shape), |
| 239 | v_descale=layer._v_scale.expand(descale_shape), |
| 240 | ) |
| 241 | return output |
| 242 | else: |
| 243 | sliding_window_size = ( |
| 244 | list(self.sliding_window) |
| 245 | if self.sliding_window is not None |
| 246 | else None |
| 247 | ) |
| 248 | flash_attn_varlen_func( |
| 249 | q=query[:num_actual_tokens], |
| 250 | k=key_cache, |
| 251 | v=value_cache, |
| 252 | out=output[:num_actual_tokens], |
| 253 | cu_seqlens_q=cu_seqlens_q, |
| 254 | max_seqlen_q=max_seqlen_q, |
| 255 | seqused_k=seqused_k, |
| 256 | max_seqlen_k=max_seqlen_k, |
| 257 | softmax_scale=self.scale, |
| 258 | causal=attn_metadata.causal, |
| 259 | alibi_slopes=self.alibi_slopes, |
| 260 | window_size=sliding_window_size, |
| 261 | block_table=block_table, |
| 262 | softcap=self.logits_soft_cap, |
| 263 | scheduler_metadata=scheduler_metadata, |
| 264 | fa_version=self.vllm_flash_attn_version, |
| 265 | q_descale=layer._q_scale.expand(descale_shape), |
| 266 | k_descale=layer._k_scale.expand(descale_shape), |
| 267 | v_descale=layer._v_scale.expand(descale_shape), |
| 268 | num_splits=attn_metadata.max_num_splits, |
| 269 | s_aux=self.sinks, |
| 270 | ) |
| 271 | return output |
| 272 | |
| 273 | # Cascade attention (rare case). |
| 274 | cascade_attention( |
| 275 | output[:num_actual_tokens], |
| 276 | query[:num_actual_tokens], |
| 277 | key_cache, |
| 278 | value_cache, |
| 279 | cu_query_lens=attn_metadata.query_start_loc, |
| 280 | max_query_len=attn_metadata.max_query_len, |
| 281 | cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens, |
| 282 | prefix_kv_lens=attn_metadata.prefix_kv_lens, |
| 283 | suffix_kv_lens=attn_metadata.suffix_kv_lens, |
| 284 | max_kv_len=attn_metadata.max_seq_len, |
| 285 | softmax_scale=self.scale, |
| 286 | alibi_slopes=self.alibi_slopes, |
| 287 | sliding_window=self.sliding_window, |
| 288 | logits_soft_cap=self.logits_soft_cap, |
| 289 | block_table=attn_metadata.block_table, |
| 290 | common_prefix_len=attn_metadata.common_prefix_len, |
| 291 | max_num_splits=attn_metadata.max_num_splits, |
| 292 | fa_version=self.vllm_flash_attn_version, |
| 293 | prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata, |
| 294 | suffix_scheduler_metadata=attn_metadata.scheduler_metadata, |
| 295 | q_descale=layer._q_scale, |
| 296 | k_descale=layer._k_scale, |
| 297 | v_descale=layer._v_scale, |
| 298 | s_aux=self.sinks, |
| 299 | ) |
| 300 | return output |