| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention layer with AiterFlashAttention.""" |
| 4 | |
| 5 | from dataclasses import dataclass, replace |
| 6 | from typing import ClassVar |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm._aiter_ops import rocm_aiter_ops |
| 11 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 12 | from vllm.config.cache import CacheDType |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.attention import Attention |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.platforms.interface import DeviceCapability |
| 17 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 18 | from vllm.utils.math_utils import cdiv |
| 19 | from vllm.utils.platform_utils import num_compute_units |
| 20 | from vllm.utils.torch_utils import ( |
| 21 | async_tensor_h2d, |
| 22 | get_dtype_size, |
| 23 | is_quantized_kv_cache, |
| 24 | ) |
| 25 | from vllm.v1.attention.backend import ( |
| 26 | AttentionBackend, |
| 27 | AttentionCGSupport, |
| 28 | AttentionImpl, |
| 29 | AttentionLayer, |
| 30 | AttentionMetadataBuilder, |
| 31 | AttentionType, |
| 32 | CommonAttentionMetadata, |
| 33 | MultipleOf, |
| 34 | ) |
| 35 | from vllm.v1.attention.backends.utils import ( |
| 36 | split_decodes_prefills_and_extends, |
| 37 | ) |
| 38 | from vllm.v1.attention.ops.merge_attn_states import merge_attn_states |
| 39 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout |
| 40 | |
| 41 | _PARTITION_SIZE_ROCM = 256 |
| 42 | _CP_TOKENS_PER_ITER_ROCM = 32 * 1024 |
| 43 | if current_platform.is_rocm(): |
| 44 | from vllm.triton_utils import tl, triton |
| 45 | |
| 46 | def block_size(x, head_dim): |
| 47 | return min(65536 // x.element_size(), triton.next_power_of_2(head_dim)) |
| 48 | |
| 49 | def num_programs(total_tokens): |
| 50 | return min(total_tokens, num_compute_units()) |
| 51 | |
| 52 | @triton.jit |
| 53 | def cp_mha_gather_cache_kernel( |
| 54 | key_cache_ptr, # [num_blocks, page_size, num_head, head_size] |
| 55 | value_cache_ptr, # [num_blocks, page_size, num_head, head_size] |
| 56 | key_ptr, # [num_tokens, num_heads, head_size] |
| 57 | value_ptr, # [num_tokens, num_heads, head_size] |
| 58 | block_table_ptr, # [num_batches, max_block_num] |
| 59 | cu_seqlens_kv_ptr, # [num_batches + 1] |
| 60 | token_to_batch_ptr, # [max_cum_tokens] |
| 61 | seq_start_ptr, # [num_batches] |
| 62 | k_scale_ptr, # [1] / [num_blocks, num_kv_heads, page_size] |
| 63 | v_scale_ptr, |
| 64 | num_heads, |
| 65 | head_size, |
| 66 | x, |
| 67 | max_block_num, |
| 68 | k_cache_stride0, |
| 69 | k_cache_stride1, |
| 70 | k_cache_stride2, |
| 71 | v_cache_stride0, |
| 72 | v_cache_stride1, |
| 73 | v_cache_stride2, |
| 74 | DEQUANT: tl.constexpr, |
| 75 | PAGE_SIZE: tl.constexpr, |
| 76 | CACHE_FORMAT: tl.constexpr, |
| 77 | BLOCK_SIZE: tl.constexpr, |
| 78 | ): |
| 79 | token_id = tl.program_id(0) |
| 80 | head_id = tl.program_id(1) |
| 81 | col_offsets = tl.arange(0, BLOCK_SIZE) |
| 82 | |
| 83 | key_ptr_offset = ( |
| 84 | key_ptr + token_id * head_size * num_heads + head_id * head_size |
| 85 | ) |
| 86 | value_ptr_offset = ( |
| 87 | value_ptr + token_id * head_size * num_heads + head_id * head_size |
| 88 | ) |
| 89 | batch_idx = tl.load(token_to_batch_ptr + token_id) |
| 90 | batch_start = tl.load(seq_start_ptr + batch_idx) |
| 91 | token_start = tl.load(cu_seqlens_kv_ptr + batch_idx) |
| 92 | batch_offset = token_id - token_start + batch_start |
| 93 | block_offset = batch_offset // PAGE_SIZE |
| 94 | block_id = tl.load( |
| 95 | block_table_ptr + max_block_num * batch_idx + block_offset |
| 96 | ).to(tl.int64) |
| 97 | slot_id = batch_offset % PAGE_SIZE |
| 98 | |
| 99 | if CACHE_FORMAT == "NHD": |
| 100 | # for kv cache layout as |
| 101 | # K: [num_blocks, page_size, num_head, head_dim] |
| 102 | # V: [num_blocks, page_size, num_head, head_dim] |
| 103 | key_cache_ptr_offset = ( |
| 104 | key_cache_ptr |
| 105 | + block_id * k_cache_stride0 |
| 106 | + slot_id * k_cache_stride1 |
| 107 | + head_id * k_cache_stride2 |
| 108 | ) |
| 109 | value_cache_ptr_offset = ( |
| 110 | value_cache_ptr |
| 111 | + block_id * v_cache_stride0 |
| 112 | + slot_id * v_cache_stride1 |
| 113 | + head_id * v_cache_stride2 |
| 114 | ) |
| 115 | k_reg = tl.load(key_cache_ptr_offset + col_offsets) |
| 116 | v_reg = tl.load(value_cache_ptr_offset + col_offsets) |
| 117 | if DEQUANT: |
| 118 | k_scale = tl.load(k_scale_ptr) |
| 119 | v_scale = tl.load(v_scale_ptr) |
| 120 | k_reg = (k_reg.to(tl.float32) * k_scale).to( |
| 121 | key_ptr_offset.dtype.element_ty |
| 122 | ) |
| 123 | v_reg = (v_reg.to(tl.float32) * v_scale).to( |
| 124 | value_ptr_offset.dtype.element_ty |
| 125 | ) |
| 126 | tl.store(key_ptr_offset + col_offsets, k_reg) |
| 127 | tl.store(value_ptr_offset + col_offsets, v_reg) |
| 128 | |
| 129 | elif CACHE_FORMAT == "SHUFFLE": |
| 130 | # for kv cache layout as |
| 131 | # K: [num_blocks, num_head, head_dim // x, page_size, x] |
| 132 | # V: [num_blocks, num_head, page_size // x, head_dim, x] |
| 133 | # Take the block stride from the tensor, as the NHD branch does: |
| 134 | # the two sides only sit a dense page apart when they are separate |
| 135 | # planes, and models that pin a block-compact layout interleave |
| 136 | # them instead. |
| 137 | key_cache_ptr_offset = ( |
| 138 | key_cache_ptr |
| 139 | + block_id * k_cache_stride0 |
| 140 | + head_id * head_size * PAGE_SIZE |
| 141 | + slot_id * x |
| 142 | ) |
| 143 | value_cache_ptr_offset = ( |
| 144 | value_cache_ptr |
| 145 | + block_id * v_cache_stride0 |
| 146 | + head_id * head_size * PAGE_SIZE |
| 147 | + (slot_id // x) * head_size * x |
| 148 | + slot_id % x |
| 149 | ) |
| 150 | k_reg_offset = col_offsets // x * PAGE_SIZE * x + col_offsets % x |
| 151 | v_reg_offset = col_offsets * x |
| 152 | k_reg = tl.load(key_cache_ptr_offset + k_reg_offset) |
| 153 | v_reg = tl.load(value_cache_ptr_offset + v_reg_offset) |
| 154 | if DEQUANT: |
| 155 | k_scale = 1.0 |
| 156 | v_scale = 1.0 |
| 157 | k_reg = k_reg.to(tl.float32) * k_scale |
| 158 | v_reg = v_reg.to(tl.float32) * v_scale |
| 159 | tl.store(key_ptr_offset + col_offsets, k_reg) |
| 160 | tl.store(value_ptr_offset + col_offsets, v_reg) |
| 161 | |
| 162 | def cp_mha_gather_cache( |
| 163 | key_cache: torch.Tensor, |
| 164 | value_cache: torch.Tensor, |
| 165 | key: torch.Tensor, |
| 166 | value: torch.Tensor, |
| 167 | block_tables: torch.Tensor, |
| 168 | k_scales: torch.Tensor, |
| 169 | v_scales: torch.Tensor, |
| 170 | cu_seqlens_kv: torch.Tensor, |
| 171 | token_to_batch: torch.Tensor, |
| 172 | seq_starts: torch.Tensor, |
| 173 | dequant: bool, |
| 174 | kv_cache_layout: str, |
| 175 | total_tokens: int, |
| 176 | ): |
| 177 | assert kv_cache_layout in ["NHD", "SHUFFLE"], ( |
| 178 | "kv_cache_layout only supports NHD, SHUFFLE" |
| 179 | ) |
| 180 | head_dim = key.shape[2] |
| 181 | x = 16 // key_cache.element_size() |
| 182 | # assert dequant is True, "Currently, we only support "\ |
| 183 | # "gather cache with dequant" |
| 184 | # For k cache layout: [num_blocks, page_size, num_heads, head_dim] |
| 185 | assert head_dim == key_cache.shape[3], ( |
| 186 | "We assume your kv cache layout is [num_blocks, " |
| 187 | "page_size, num_heads, head_dim], but got otherwise" |
| 188 | ) |
| 189 | page_size = key_cache.shape[1] |
| 190 | num_heads = key_cache.shape[2] |
| 191 | |
| 192 | # Pass actual tensor strides so the kernel works correctly |
| 193 | # even when the cache is non-contiguous (e.g, for hybrid model) |
| 194 | k_strides = key_cache.stride() |
| 195 | v_strides = value_cache.stride() |
| 196 | |
| 197 | grid = lambda meta: (total_tokens, num_heads) |
| 198 | cp_mha_gather_cache_kernel[grid]( |
| 199 | key_cache, |
| 200 | value_cache, |
| 201 | key, |
| 202 | value, |
| 203 | block_tables, |
| 204 | cu_seqlens_kv, |
| 205 | token_to_batch, |
| 206 | seq_starts, |
| 207 | k_scales, |
| 208 | v_scales, |
| 209 | num_heads, |
| 210 | head_dim, |
| 211 | x, |
| 212 | block_tables.size(1), |
| 213 | k_strides[0], |
| 214 | k_strides[1], |
| 215 | k_strides[2], |
| 216 | v_strides[0], |
| 217 | v_strides[1], |
| 218 | v_strides[2], |
| 219 | DEQUANT=dequant, |
| 220 | PAGE_SIZE=page_size, |
| 221 | CACHE_FORMAT=kv_cache_layout, |
| 222 | BLOCK_SIZE=head_dim, |
| 223 | ) |
| 224 | |
| 225 | @triton.jit |
| 226 | def reshape_and_cache_shuffle_kernel( |
| 227 | key_ptr, # [num_tokens, num_kv_heads, head_size] |
| 228 | value_ptr, # [num_tokens, num_kv_heads, head_size] |
| 229 | key_cache_ptr, # [num_blocks, num_kv_heads, head_size // x, block_size, x] |
| 230 | value_cache_ptr, # [num_blocks, num_kv_heads, block_size // x, head_size, x] |
| 231 | slot_mapping_ptr, # [num_tokens] |
| 232 | k_scale_ptr, # [num_blocks, num_kv_heads, block_size] |
| 233 | v_scale_ptr, # [num_blocks, num_kv_heads, block_size] |
| 234 | x, |
| 235 | k_stride0, |
| 236 | v_stride0, |
| 237 | k_cache_block_stride, |
| 238 | v_cache_block_stride, |
| 239 | block_size, |
| 240 | head_size, |
| 241 | num_kv_heads, |
| 242 | BLOCK_SIZE: tl.constexpr, |
| 243 | QUANT: tl.constexpr, |
| 244 | ): |
| 245 | tid = tl.program_id(0) |
| 246 | head_id = tl.program_id(1) |
| 247 | offset = tl.arange(0, BLOCK_SIZE) |
| 248 | src_offset_k = tid * k_stride0 + head_id * head_size |
| 249 | src_offset_v = tid * v_stride0 + head_id * head_size |
| 250 | slot_id = tl.load(slot_mapping_ptr + tid) |
| 251 | if slot_id < 0: |
| 252 | return |
| 253 | block_id = slot_id // block_size |
| 254 | block_offset = slot_id % block_size |
| 255 | k_dst_offset = ( |
| 256 | block_id * k_cache_block_stride + head_id * head_size * block_size |
| 257 | ) |
| 258 | dst_k_shuffle_offset = ( |
| 259 | k_dst_offset + offset // x * block_size * x + block_offset * x + offset % x |
| 260 | ) |
| 261 | v_dst_offset = ( |
| 262 | block_id * v_cache_block_stride + head_id * head_size * block_size |
| 263 | ) |
| 264 | dst_v_shuffle_offset = ( |
| 265 | v_dst_offset |
| 266 | + block_offset // x * head_size * x |
| 267 | + offset * x |
| 268 | + block_offset % x |
| 269 | ) |
| 270 | k_val = tl.load(key_ptr + src_offset_k + offset) |
| 271 | v_val = tl.load(value_ptr + src_offset_v + offset) |
| 272 | if QUANT: |
| 273 | k_scale = 1.0 |
| 274 | v_scale = 1.0 |
| 275 | k_dtype = key_cache_ptr.type.element_ty |
| 276 | v_dtype = value_cache_ptr.type.element_ty |
| 277 | k_val = (k_val.to(tl.float32) / k_scale).to(k_dtype) |
| 278 | v_val = (v_val.to(tl.float32) / v_scale).to(v_dtype) |
| 279 | tl.store(key_cache_ptr + dst_k_shuffle_offset, k_val) |
| 280 | tl.store(value_cache_ptr + dst_v_shuffle_offset, v_val) |
| 281 | |
| 282 | def reshape_and_cache_shuffle_triton( |
| 283 | key: torch.Tensor, |
| 284 | value: torch.Tensor, |
| 285 | key_cache: torch.Tensor, |
| 286 | value_cache: torch.Tensor, |
| 287 | slot_mapping: torch.Tensor, |
| 288 | kv_cache_dtype: str, |
| 289 | k_scales: torch.Tensor, |
| 290 | v_scales: torch.Tensor, |
| 291 | ): |
| 292 | num_tokens = slot_mapping.shape[0] |
| 293 | _, num_kv_heads, head_size = key.shape |
| 294 | num_blocks, block_size, _, _ = key_cache.shape |
| 295 | x = 16 // key_cache.element_size() |
| 296 | QUANT = False |
| 297 | if is_quantized_kv_cache(kv_cache_dtype): |
| 298 | QUANT = True |
| 299 | grid = ( |
| 300 | num_tokens, |
| 301 | num_kv_heads, |
| 302 | ) |
| 303 | reshape_and_cache_shuffle_kernel[grid]( |
| 304 | key, |
| 305 | value, |
| 306 | key_cache, |
| 307 | value_cache, |
| 308 | slot_mapping, |
| 309 | k_scales, |
| 310 | v_scales, |
| 311 | x, |
| 312 | key.stride(0), |
| 313 | value.stride(0), |
| 314 | key_cache.stride(0), |
| 315 | value_cache.stride(0), |
| 316 | block_size, |
| 317 | head_size, |
| 318 | num_kv_heads, |
| 319 | BLOCK_SIZE=head_size, |
| 320 | QUANT=QUANT, |
| 321 | ) |
| 322 | |
| 323 | |
| 324 | logger = init_logger(__name__) |
| 325 | |
| 326 | |
| 327 | @dataclass |
| 328 | class AiterFlashAttentionDecodeMetadata: |
| 329 | max_query_len: int |
| 330 | |
| 331 | |
| 332 | @dataclass |
| 333 | class AiterFlashAttentionPrefillMetadata: |
| 334 | max_query_len: int |
| 335 | max_seq_len: int |
| 336 | query_start_loc: torch.Tensor |
| 337 | |
| 338 | |
| 339 | @dataclass |
| 340 | class AiterChunkSlidingWindowMetadata: |
| 341 | swa_cu_seqlens: torch.Tensor |
| 342 | swa_seq_starts: torch.Tensor |
| 343 | swa_token_to_batch: torch.Tensor |
| 344 | swa_max_seqlens: int |
| 345 | swa_total_tokens: int |
| 346 | swa_workspace: torch.Tensor |
| 347 | |
| 348 | |
| 349 | @dataclass |
| 350 | class AiterChunkContextMetadata: |
| 351 | workspace: torch.Tensor |
| 352 | cu_seq_lens_chunk: torch.Tensor |
| 353 | chunk_starts: torch.Tensor |
| 354 | token_to_batch: torch.Tensor |
| 355 | max_seq_lens: list[int] |
| 356 | num_chunks: int |
| 357 | total_token_per_batch: list[int] |
| 358 | swa_metadata: AiterChunkSlidingWindowMetadata | None |
| 359 | |
| 360 | |
| 361 | @dataclass |
| 362 | class AiterFlashAttentionChunkPrefillMetadata: |
| 363 | max_query_len: int |
| 364 | max_seq_len: int |
| 365 | query_start_loc: torch.Tensor |
| 366 | chunk_context_metadata: AiterChunkContextMetadata |
| 367 | |
| 368 | |
| 369 | @dataclass |
| 370 | class AiterKVSharingMetadata: |
| 371 | workspace: torch.Tensor |
| 372 | query_start_loc: torch.Tensor |
| 373 | token_to_batch: torch.Tensor |
| 374 | seq_starts: torch.Tensor |
| 375 | |
| 376 | |
| 377 | @dataclass |
| 378 | class AiterFlashAttentionMetadata: |
| 379 | # NOTE(sang): Definition of context_len, query_len, and seq_len. |
| 380 | # |---------- N-1 iteration --------| |
| 381 | # |---------------- N iteration ---------------------| |
| 382 | # |- tokenA -|......................|-- newTokens ---| |
| 383 | # |---------- context_len ----------| |
| 384 | # |-------------------- seq_len ---------------------| |
| 385 | # |-- query_len ---| |
| 386 | |
| 387 | num_actual_tokens: int # Number of tokens excluding padding. |
| 388 | query_start_loc: torch.Tensor |
| 389 | max_seq_len: int |
| 390 | seq_lens: torch.Tensor |
| 391 | slot_mapping: torch.Tensor |
| 392 | block_table: torch.Tensor |
| 393 | causal: bool |
| 394 | |
| 395 | # prefill and decode split |
| 396 | num_decodes: int |
| 397 | num_decode_tokens: int |
| 398 | num_prefills: int |
| 399 | num_extends: int |
| 400 | num_extend_tokens: int |
| 401 | |
| 402 | decode_metadata: AiterFlashAttentionDecodeMetadata | None |
| 403 | prefill_metadata: AiterFlashAttentionPrefillMetadata | None |
| 404 | extend_metadata: AiterFlashAttentionChunkPrefillMetadata | None |
| 405 | |
| 406 | # For cascade attention. |
| 407 | use_cascade: bool |
| 408 | |
| 409 | # Only for fp8 shuffle layout kv cache, we allocate kv_scale for each layer |
| 410 | # since we might integrate per token quant for kv cache in the future. |
| 411 | k_scale: dict[str, torch.Tensor] | None |
| 412 | v_scale: dict[str, torch.Tensor] | None |
| 413 | kv_sharing_metadata: AiterKVSharingMetadata | None = None |
| 414 | |
| 415 | |
| 416 | class AiterFlashAttentionMetadataBuilder( |
| 417 | AttentionMetadataBuilder[AiterFlashAttentionMetadata] |
| 418 | ): |
| 419 | _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH |
| 420 | |
| 421 | def __init__( |
| 422 | self, |
| 423 | kv_cache_spec: AttentionSpec, |
| 424 | layer_names: list[str], |
| 425 | vllm_config: VllmConfig, |
| 426 | device: torch.device, |
| 427 | ): |
| 428 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 429 | |
| 430 | self.model_config = vllm_config.model_config |
| 431 | self.parallel_config = vllm_config.parallel_config |
| 432 | self.cache_config = vllm_config.cache_config |
| 433 | |
| 434 | self.num_heads_q = self.model_config.get_num_attention_heads( |
| 435 | self.parallel_config |
| 436 | ) |
| 437 | self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config) |
| 438 | self.headdim = self.model_config.get_head_size() |
| 439 | self.block_size = kv_cache_spec.block_size |
| 440 | # Sliding window size to be used with the AOT scheduler will be |
| 441 | # populated on first build() call. |
| 442 | self.aot_sliding_window: tuple[int, int] | None = None |
| 443 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 444 | |
| 445 | sliding_window_configs: set[tuple[int, int] | None] = set() |
| 446 | kv_sharing_shape = None |
| 447 | layers = get_layers_from_vllm_config(self.vllm_config, Attention) |
| 448 | for name, layer in layers.items(): |
| 449 | if name not in layer_names: |
| 450 | continue |
| 451 | assert isinstance(layer.impl, AiterFlashAttentionImpl), ( |
| 452 | "Aiter Flash Attention Metadata Builder can only be used " |
| 453 | "with Aiter Flash Attention Impl." |
| 454 | ) |
| 455 | sliding_window_configs.add(layer.impl.sliding_window) |
| 456 | if layer.kv_sharing_target_layer_name is not None: |
| 457 | kv_sharing_shape = (layer.impl.num_kv_heads, layer.impl.head_size) |
| 458 | |
| 459 | while len(sliding_window_configs) > 0: |
| 460 | sliding_window_config = sliding_window_configs.pop() |
| 461 | if sliding_window_config is not None and sliding_window_config[0] != -1: |
| 462 | assert self.aot_sliding_window is None, ( |
| 463 | "Aiter Flash ATTENTION can only support one valid sliding window!" |
| 464 | ) |
| 465 | self.aot_sliding_window = sliding_window_config |
| 466 | |
| 467 | self.extend_workspace = torch.empty( |
| 468 | [2, _CP_TOKENS_PER_ITER_ROCM, self.num_heads_kv, self.headdim], |
| 469 | dtype=self.model_config.dtype, |
| 470 | device=device, |
| 471 | ) |
| 472 | self.scale = torch.tensor([1.0], dtype=torch.float, device=self.device) |
| 473 | self.kv_sharing_workspace = ( |
| 474 | torch.empty( |
| 475 | ( |
| 476 | 2, |
| 477 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 478 | *kv_sharing_shape, |
| 479 | ), |
| 480 | dtype=self.model_config.dtype, |
| 481 | device=device, |
| 482 | ) |
| 483 | if kv_sharing_shape is not None |
| 484 | else None |
| 485 | ) |
| 486 | |
| 487 | def build_for_cudagraph_capture( |
| 488 | self, common_attn_metadata: CommonAttentionMetadata |
| 489 | ): |
| 490 | return self.build( |
| 491 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 492 | ) |
| 493 | |
| 494 | def build( |
| 495 | self, |
| 496 | common_prefix_len: int, |
| 497 | common_attn_metadata: CommonAttentionMetadata, |
| 498 | fast_build: bool = False, |
| 499 | ) -> "AiterFlashAttentionMetadata": |
| 500 | assert self.reorder_batch_threshold is not None |
| 501 | split_ret = split_decodes_prefills_and_extends( |
| 502 | common_attn_metadata, |
| 503 | decode_threshold=self.reorder_batch_threshold, |
| 504 | ) |
| 505 | # Allocate scales for fp8 shuffle kv cache with shuffle_kv_cache enabled |
| 506 | if ( |
| 507 | rocm_aiter_ops.is_shuffle_kv_cache_enabled() |
| 508 | and self.scale.numel() == 1 |
| 509 | and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype) |
| 510 | ): |
| 511 | # Size the scales from a layer this builder owns. The draft model |
| 512 | # runs its own builder over its own KV cache, so the first layer of |
| 513 | # the whole config can carry an unrelated block count. |
| 514 | kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[ |
| 515 | self.layer_names[0] |
| 516 | ].kv_cache.shape |
| 517 | num_blocks = kv_cache_shape[0] |
| 518 | self.scale = torch.ones( |
| 519 | [num_blocks, self.num_heads_kv, self.block_size], |
| 520 | dtype=torch.float32, |
| 521 | device=self.device, |
| 522 | ) |
| 523 | ( |
| 524 | num_decodes, |
| 525 | num_extends, |
| 526 | num_prefills, |
| 527 | num_decode_tokens, |
| 528 | num_extend_tokens, |
| 529 | _, |
| 530 | ) = split_ret |
| 531 | |
| 532 | query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu |
| 533 | |
| 534 | with gpu_sync_allowed(): |
| 535 | # Only copy seq_lens to CPU when prefill or extend is present to avoid a |
| 536 | # blocking device→host transfer. |
| 537 | seq_lens = ( |
| 538 | common_attn_metadata.seq_lens.cpu() |
| 539 | if num_prefills > 0 or num_extends > 0 |
| 540 | else None |
| 541 | ) |
| 542 | |
| 543 | query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] |
| 544 | |
| 545 | kv_sharing_metadata = None |
| 546 | if self.kv_sharing_workspace is not None and seq_lens is not None: |
| 547 | query_lens = query_lens_cpu[num_decodes:] |
| 548 | token_to_batch = torch.repeat_interleave( |
| 549 | torch.arange(len(query_lens), dtype=torch.int32, device="cpu"), |
| 550 | query_lens, |
| 551 | ) |
| 552 | query_start_loc = common_attn_metadata.query_start_loc[num_decodes:] |
| 553 | kv_sharing_metadata = AiterKVSharingMetadata( |
| 554 | workspace=self.kv_sharing_workspace, |
| 555 | query_start_loc=query_start_loc - query_start_loc[0], |
| 556 | token_to_batch=async_tensor_h2d(token_to_batch, self.device), |
| 557 | seq_starts=async_tensor_h2d( |
| 558 | seq_lens[num_decodes:] - query_lens, self.device |
| 559 | ), |
| 560 | ) |
| 561 | |
| 562 | decode_metadata = None |
| 563 | if num_decodes > 0: |
| 564 | decode_metadata = AiterFlashAttentionDecodeMetadata( |
| 565 | max_query_len=query_lens_cpu[:num_decodes].max().item(), |
| 566 | ) |
| 567 | |
| 568 | prefill_metadata = None |
| 569 | if num_prefills > 0: |
| 570 | assert seq_lens is not None |
| 571 | query_lens_for_prefill = query_lens_cpu[num_decodes + num_extends :] |
| 572 | query_start_loc_device = common_attn_metadata.query_start_loc[ |
| 573 | num_decodes + num_extends : |
| 574 | ] |
| 575 | prefill_metadata = AiterFlashAttentionPrefillMetadata( |
| 576 | max_query_len=query_lens_for_prefill.max().item(), |
| 577 | max_seq_len=seq_lens[num_decodes + num_extends :].max().item(), |
| 578 | query_start_loc=query_start_loc_device - query_start_loc_device[0], |
| 579 | ) |
| 580 | |
| 581 | extend_metadata = None |
| 582 | if num_extends > 0: |
| 583 | assert seq_lens is not None |
| 584 | num_extends_slice = slice(num_decodes, num_decodes + num_extends) |
| 585 | query_lens_for_extend = query_lens_cpu[num_extends_slice] |
| 586 | seq_lens_for_extend = seq_lens[num_extends_slice] |
| 587 | computed_kv_lens = seq_lens_for_extend - query_lens_for_extend |
| 588 | swa_metadata = None |
| 589 | if self.aot_sliding_window is not None: |
| 590 | swa_seqlen_for_extend = torch.minimum( |
| 591 | seq_lens_for_extend, |
| 592 | query_lens_for_extend + self.aot_sliding_window[0] + 1, |
| 593 | ) |
| 594 | cu_seq_lens = torch.zeros( |
| 595 | num_extends + 1, |
| 596 | dtype=torch.int32, |
| 597 | device=seq_lens_for_extend.device, |
| 598 | pin_memory=True, |
| 599 | ) |
| 600 | torch.cumsum( |
| 601 | swa_seqlen_for_extend, |
| 602 | dim=0, |
| 603 | dtype=cu_seq_lens.dtype, |
| 604 | out=cu_seq_lens[1:], |
| 605 | ) |
| 606 | token_to_seq = torch.arange( |
| 607 | 0, |
| 608 | num_extends, |
| 609 | dtype=torch.int32, |
| 610 | device=seq_lens_for_extend.device, |
| 611 | ) |
| 612 | token_to_seq = torch.repeat_interleave( |
| 613 | token_to_seq, swa_seqlen_for_extend |
| 614 | ) |
| 615 | fetched_shape = cu_seq_lens[-1].item() |
| 616 | # TODO(ganyi): Maybe reuse these 2 buffer from extend_workspace |
| 617 | swa_workspace = torch.empty( |
| 618 | (2, fetched_shape, self.num_heads_kv, self.headdim), |
| 619 | dtype=self.vllm_config.model_config.dtype, |
| 620 | device=self.device, |
| 621 | ) |
| 622 | |
| 623 | seq_starts = seq_lens_for_extend - swa_seqlen_for_extend |
| 624 | max_seqlen_k = swa_seqlen_for_extend.max().item() |
| 625 | total_tokens = cu_seq_lens[-1].item() |
| 626 | |
| 627 | swa_metadata = AiterChunkSlidingWindowMetadata( |
| 628 | swa_cu_seqlens=cu_seq_lens.to(self.device, non_blocking=True), |
| 629 | swa_seq_starts=async_tensor_h2d(seq_starts, self.device), |
| 630 | swa_token_to_batch=async_tensor_h2d(token_to_seq, self.device), |
| 631 | swa_max_seqlens=max_seqlen_k, |
| 632 | swa_total_tokens=total_tokens, |
| 633 | swa_workspace=swa_workspace, |
| 634 | ) |
| 635 | |
| 636 | # allocate the equal amount of workspace for |
| 637 | # each chunk prefill request |
| 638 | max_context_chunk = _CP_TOKENS_PER_ITER_ROCM // num_extends |
| 639 | num_chunks = cdiv(computed_kv_lens.max().item(), max_context_chunk) |
| 640 | |
| 641 | chunk_starts = ( |
| 642 | torch.arange(num_chunks, dtype=torch.int32) |
| 643 | .unsqueeze(1) |
| 644 | .expand(-1, num_extends) |
| 645 | * max_context_chunk |
| 646 | ) |
| 647 | chunk_ends = torch.min( |
| 648 | computed_kv_lens.unsqueeze(0), chunk_starts + max_context_chunk |
| 649 | ) |
| 650 | chunk_seq_lens = chunk_ends - chunk_starts |
| 651 | chunk_seq_lens.clamp_(min=0) # [num_chunks, num_extends] |
| 652 | cu_seq_lens_cpu = torch.zeros( |
| 653 | [num_chunks, num_extends + 1], dtype=torch.int32, pin_memory=True |
| 654 | ) |
| 655 | torch.cumsum( |
| 656 | chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32 |
| 657 | ) |
| 658 | # Avoid .max() on an empty tensor when there is no context |
| 659 | # (num_chunks == 0, e.g. Whisper encoder's first pass). |
| 660 | max_cum_tokens = ( |
| 661 | cu_seq_lens_cpu[:, -1].max().item() if num_chunks > 0 else 0 |
| 662 | ) |
| 663 | |
| 664 | range_idx = torch.arange(max_cum_tokens, dtype=torch.int32)[None, None, :] |
| 665 | idx_to_batch_tensor = range_idx == cu_seq_lens_cpu[:, 1:][:, :, None] |
| 666 | idx_to_batch_tensor = idx_to_batch_tensor.sum( |
| 667 | dim=1 |
| 668 | ) # [num_chunks, max_cum_tokens] |
| 669 | token_to_batch_tensor = torch.cumsum(idx_to_batch_tensor, dim=1) |
| 670 | |
| 671 | chunk_context_metadata = AiterChunkContextMetadata( |
| 672 | workspace=self.extend_workspace, |
| 673 | cu_seq_lens_chunk=cu_seq_lens_cpu.to(self.device, non_blocking=True), |
| 674 | chunk_starts=async_tensor_h2d(chunk_starts, self.device), |
| 675 | token_to_batch=async_tensor_h2d(token_to_batch_tensor, self.device), |
| 676 | max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(), |
| 677 | num_chunks=num_chunks, |
| 678 | total_token_per_batch=cu_seq_lens_cpu[:, -1].tolist(), |
| 679 | swa_metadata=swa_metadata, |
| 680 | ) |
| 681 | |
| 682 | query_start_loc_device = common_attn_metadata.query_start_loc[ |
| 683 | num_decodes : num_decodes + num_extends + 1 |
| 684 | ] |
| 685 | seq_lens_device = common_attn_metadata.seq_lens[num_extends_slice] |
| 686 | cu_seq_lens = torch.zeros( |
| 687 | num_extends + 1, dtype=torch.int32, device=seq_lens_device.device |
| 688 | ) |
| 689 | torch.cumsum( |
| 690 | seq_lens_device, dim=0, dtype=cu_seq_lens.dtype, out=cu_seq_lens[1:] |
| 691 | ) |
| 692 | extend_metadata = AiterFlashAttentionChunkPrefillMetadata( |
| 693 | max_query_len=query_lens_for_extend.max().item(), |
| 694 | max_seq_len=seq_lens[num_extends_slice].max().item(), |
| 695 | query_start_loc=query_start_loc_device - query_start_loc_device[0], |
| 696 | chunk_context_metadata=chunk_context_metadata, |
| 697 | ) |
| 698 | |
| 699 | use_cascade = common_prefix_len > 0 |
| 700 | |
| 701 | attn_metadata = AiterFlashAttentionMetadata( |
| 702 | num_actual_tokens=common_attn_metadata.num_actual_tokens, |
| 703 | query_start_loc=common_attn_metadata.query_start_loc, |
| 704 | max_seq_len=common_attn_metadata.max_seq_len, |
| 705 | seq_lens=common_attn_metadata.seq_lens, |
| 706 | block_table=common_attn_metadata.block_table_tensor, |
| 707 | causal=common_attn_metadata.causal, |
| 708 | slot_mapping=common_attn_metadata.slot_mapping, |
| 709 | num_decodes=num_decodes, |
| 710 | num_decode_tokens=num_decode_tokens, |
| 711 | num_prefills=num_prefills, |
| 712 | num_extends=num_extends, |
| 713 | num_extend_tokens=num_extend_tokens, |
| 714 | decode_metadata=decode_metadata, |
| 715 | prefill_metadata=prefill_metadata, |
| 716 | extend_metadata=extend_metadata, |
| 717 | use_cascade=use_cascade, |
| 718 | k_scale=self.scale, |
| 719 | v_scale=self.scale, |
| 720 | kv_sharing_metadata=kv_sharing_metadata, |
| 721 | ) |
| 722 | return attn_metadata |
| 723 | |
| 724 | def build_for_drafting( |
| 725 | self, |
| 726 | common_attn_metadata: CommonAttentionMetadata, |
| 727 | draft_index: int, |
| 728 | ) -> AiterFlashAttentionMetadata: |
| 729 | """Build attention metadata for draft model without CPU-GPU sync. |
| 730 | |
| 731 | During EAGLE drafting all requests are uniform decodes, so we can |
| 732 | skip split_decodes_prefills_and_extends() and avoid all .cpu() / |
| 733 | .item() calls that would otherwise break CUDA graph capture. |
| 734 | """ |
| 735 | num_reqs = common_attn_metadata.num_reqs |
| 736 | num_tokens = common_attn_metadata.num_actual_tokens |
| 737 | |
| 738 | decode_metadata = AiterFlashAttentionDecodeMetadata( |
| 739 | max_query_len=common_attn_metadata.max_query_len, |
| 740 | ) |
| 741 | |
| 742 | return AiterFlashAttentionMetadata( |
| 743 | num_actual_tokens=num_tokens, |
| 744 | query_start_loc=common_attn_metadata.query_start_loc, |
| 745 | max_seq_len=common_attn_metadata.max_seq_len, |
| 746 | seq_lens=common_attn_metadata.seq_lens, |
| 747 | block_table=common_attn_metadata.block_table_tensor, |
| 748 | causal=common_attn_metadata.causal, |
| 749 | slot_mapping=common_attn_metadata.slot_mapping, |
| 750 | num_decodes=num_reqs, |
| 751 | num_decode_tokens=num_tokens, |
| 752 | num_prefills=0, |
| 753 | num_extends=0, |
| 754 | num_extend_tokens=0, |
| 755 | decode_metadata=decode_metadata, |
| 756 | prefill_metadata=None, |
| 757 | extend_metadata=None, |
| 758 | use_cascade=False, |
| 759 | k_scale=self.scale, |
| 760 | v_scale=self.scale, |
| 761 | ) |
| 762 | |
| 763 | def use_cascade_attention(self, *args, **kwargs) -> bool: |
| 764 | return False |
| 765 | |
| 766 | |
| 767 | class AiterFlashAttentionBackend(AttentionBackend): |
| 768 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] |
| 769 | |
| 770 | @classmethod |
| 771 | def supports_sink(cls) -> bool: |
| 772 | return True |
| 773 | |
| 774 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 775 | "auto", |
| 776 | "float16", |
| 777 | "bfloat16", |
| 778 | "fp8", |
| 779 | "fp8_e4m3", |
| 780 | "fp8_e5m2", |
| 781 | ] |
| 782 | |
| 783 | @classmethod |
| 784 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 785 | """ENCODER_DECODER is not supported because the prefill path uses |
| 786 | flash_attn_varlen_func with cu_seqlens_k set to decoder |
| 787 | query_start_loc (not encoder seq lens) and causal=True, both of |
| 788 | which are incorrect for cross-attention layers. |
| 789 | """ |
| 790 | return attn_type in (AttentionType.DECODER,) |
| 791 | |
| 792 | @staticmethod |
| 793 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 794 | return [16, 32] |
| 795 | |
| 796 | @classmethod |
| 797 | def get_supported_head_sizes(cls) -> list[int]: |
| 798 | return [64, 128, 256] |
| 799 | |
| 800 | forward_includes_kv_cache_update: bool = False |
| 801 | |
| 802 | @staticmethod |
| 803 | def get_name() -> str: |
| 804 | return "FLASH_ATTN" |
| 805 | |
| 806 | @classmethod |
| 807 | def supports_sliding_window(cls) -> bool: |
| 808 | return True |
| 809 | |
| 810 | @staticmethod |
| 811 | def get_impl_cls() -> type["AiterFlashAttentionImpl"]: |
| 812 | return AiterFlashAttentionImpl |
| 813 | |
| 814 | @staticmethod |
| 815 | def get_builder_cls() -> type["AiterFlashAttentionMetadataBuilder"]: |
| 816 | return AiterFlashAttentionMetadataBuilder |
| 817 | |
| 818 | @classmethod |
| 819 | def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: |
| 820 | """Validate the block size the ROCm gather kernels require, and under |
| 821 | the shuffle layout publish K and V as two head slots. |
| 822 | |
| 823 | The shuffle read/write kernels rearrange a whole ``head_size x |
| 824 | block_size`` tile per side, so each side needs its own contiguous run |
| 825 | in the page. Packed into the content dim they interleave per token |
| 826 | instead, which no layout can undo. |
| 827 | """ |
| 828 | # block_size == 1 is the per-token page-size probe (see |
| 829 | # Platform.get_page_size_bytes); real blocks must be gatherable in |
| 830 | # 16-token units by the ROCm kernel. |
| 831 | if spec.block_size != 1 and spec.block_size % 16 != 0: |
| 832 | raise ValueError("Block size must be a multiple of 16.") |
| 833 | if not rocm_aiter_ops.is_shuffle_kv_cache_enabled(): |
| 834 | return spec |
| 835 | if spec.state_content_bytes is not None: |
| 836 | return spec |
| 837 | assert spec.head_size == spec.head_size_v, ( |
| 838 | "Separate K/V head slots require symmetric K/V head sizes." |
| 839 | ) |
| 840 | return replace( |
| 841 | spec, |
| 842 | num_head_slots=2, |
| 843 | state_content_bytes=spec.num_kv_heads |
| 844 | * spec.head_size |
| 845 | * get_dtype_size(spec.dtype), |
| 846 | ) |
| 847 | |
| 848 | @classmethod |
| 849 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 850 | if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): |
| 851 | # pa_fwd_asm strides between pages by a whole dense page, so the two |
| 852 | # head slots customize_spec publishes have to be separate planes |
| 853 | # spanning every block: H outermost. Under LBHNC the slots alternate |
| 854 | # block by block, doubling the stride the kernel assumes. LBHNC |
| 855 | # stays listed for models whose mixed HNC shapes need a |
| 856 | # block-compact layout; those read K/V by stride instead. |
| 857 | return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC) |
| 858 | # K and V come out of the content dim as transposed views rather than |
| 859 | # copies, so the head dim may sit on either side of the block dim, but |
| 860 | # the layer must stay outermost. |
| 861 | return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) |
| 862 | |
| 863 | @classmethod |
| 864 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 865 | from vllm.platforms.rocm import get_cdna_version |
| 866 | |
| 867 | # DeviceCapability is currently created using torch.cuda.get_device_capability() |
| 868 | # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is |
| 869 | # more reliable. |
| 870 | return get_cdna_version() > 2 |
| 871 | |
| 872 | @classmethod |
| 873 | def supports_non_causal(cls) -> bool: |
| 874 | return True |
| 875 | |
| 876 | |
| 877 | class AiterFlashAttentionImpl(AttentionImpl): |
| 878 | def __init__( |
| 879 | self, |
| 880 | num_heads: int, |
| 881 | head_size: int, |
| 882 | scale: float, |
| 883 | num_kv_heads: int, |
| 884 | alibi_slopes: list[float] | None, |
| 885 | sliding_window: int | None, |
| 886 | kv_cache_dtype: str, |
| 887 | logits_soft_cap: float | None = None, |
| 888 | attn_type: AttentionType = AttentionType.DECODER, |
| 889 | kv_sharing_target_layer_name: int | None = None, |
| 890 | sinks: torch.Tensor | None = None, |
| 891 | ) -> None: |
| 892 | self.num_heads = num_heads |
| 893 | self.head_size = head_size |
| 894 | self.scale = float(scale) |
| 895 | self.num_kv_heads = num_kv_heads |
| 896 | if alibi_slopes is not None: |
| 897 | alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) |
| 898 | self.alibi_slopes = alibi_slopes |
| 899 | if sliding_window is None: |
| 900 | self.sliding_window = (-1, -1) |
| 901 | else: |
| 902 | self.sliding_window = (sliding_window - 1, 0) |
| 903 | self.kv_cache_dtype = kv_cache_dtype |
| 904 | if logits_soft_cap is None: |
| 905 | # In flash-attn, setting logits_soft_cap as 0 means no soft cap. |
| 906 | logits_soft_cap = 0.0 |
| 907 | self.logits_soft_cap = logits_soft_cap |
| 908 | self.kv_sharing_target_layer_name = kv_sharing_target_layer_name |
| 909 | self.sinks = sinks |
| 910 | |
| 911 | assert self.num_heads % self.num_kv_heads == 0 |
| 912 | self.num_queries_per_kv = self.num_heads // self.num_kv_heads |
| 913 | |
| 914 | if attn_type != AttentionType.DECODER: |
| 915 | raise NotImplementedError( |
| 916 | "Only decoder self-attention is supported for " |
| 917 | "AiterFlashAttentionImpl. ENCODER_DECODER is not supported " |
| 918 | "because the prefill path uses cu_seqlens_k set to decoder " |
| 919 | "query_start_loc with causal=True, which is incorrect for " |
| 920 | "cross-attention." |
| 921 | ) |
| 922 | |
| 923 | def _get_kv_cache_descales( |
| 924 | self, |
| 925 | layer: AttentionLayer, |
| 926 | num_decodes: int, |
| 927 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 928 | # AITER treats non-null descales as an instruction to dequantize. |
| 929 | if not is_quantized_kv_cache(self.kv_cache_dtype): |
| 930 | return None, None |
| 931 | descale_shape = (num_decodes, self.num_kv_heads) |
| 932 | return ( |
| 933 | layer._k_scale.expand(descale_shape), |
| 934 | layer._v_scale.expand(descale_shape), |
| 935 | ) |
| 936 | |
| 937 | def extend_for_sliding_window( |
| 938 | self, |
| 939 | attn_metadata: AiterFlashAttentionMetadata, |
| 940 | query: torch.Tensor, |
| 941 | key_cache, |
| 942 | value_cache, |
| 943 | output: torch.Tensor, |
| 944 | cu_seqlens_q: torch.Tensor, |
| 945 | max_seqlen_q: int, |
| 946 | block_table: torch.Tensor, |
| 947 | k_scale: float, |
| 948 | v_scale: float, |
| 949 | ): |
| 950 | assert attn_metadata.extend_metadata is not None |
| 951 | assert attn_metadata.extend_metadata.chunk_context_metadata is not None |
| 952 | chunked_metadata = attn_metadata.extend_metadata.chunk_context_metadata |
| 953 | swa_metadata = chunked_metadata.swa_metadata |
| 954 | assert swa_metadata is not None |
| 955 | swa_cu_seqlens = swa_metadata.swa_cu_seqlens |
| 956 | swa_seq_starts = swa_metadata.swa_seq_starts |
| 957 | swa_token_to_batch = swa_metadata.swa_token_to_batch |
| 958 | swa_max_seqlens = swa_metadata.swa_max_seqlens |
| 959 | swa_total_tokens = swa_metadata.swa_total_tokens |
| 960 | key_fetched, value_fetched = ( |
| 961 | swa_metadata.swa_workspace[0], |
| 962 | swa_metadata.swa_workspace[1], |
| 963 | ) |
| 964 | cp_mha_gather_cache( |
| 965 | key_cache=key_cache, |
| 966 | value_cache=value_cache, |
| 967 | key=key_fetched, |
| 968 | value=value_fetched, |
| 969 | block_tables=block_table, |
| 970 | k_scales=k_scale, |
| 971 | v_scales=v_scale, |
| 972 | cu_seqlens_kv=swa_cu_seqlens, |
| 973 | token_to_batch=swa_token_to_batch, |
| 974 | seq_starts=swa_seq_starts, |
| 975 | dequant=is_quantized_kv_cache(self.kv_cache_dtype), |
| 976 | kv_cache_layout="NHD", |
| 977 | total_tokens=swa_total_tokens, |
| 978 | ) |
| 979 | |
| 980 | rocm_aiter_ops.flash_attn_varlen_func( |
| 981 | q=query, |
| 982 | k=key_fetched, |
| 983 | v=value_fetched, |
| 984 | cu_seqlens_q=cu_seqlens_q, |
| 985 | cu_seqlens_k=swa_cu_seqlens, |
| 986 | max_seqlen_q=max_seqlen_q, |
| 987 | max_seqlen_k=swa_max_seqlens, |
| 988 | min_seqlen_q=1, |
| 989 | dropout_p=0.0, |
| 990 | softmax_scale=self.scale, |
| 991 | causal=True, |
| 992 | window_size=self.sliding_window, |
| 993 | alibi_slopes=self.alibi_slopes, |
| 994 | return_lse=False, |
| 995 | out=output, |
| 996 | sink_ptr=self.sinks, |
| 997 | ) |
| 998 | |
| 999 | def extend_forward( |
| 1000 | self, |
| 1001 | attn_metadata: AiterFlashAttentionMetadata, |
| 1002 | query: torch.Tensor, |
| 1003 | key: torch.Tensor, |
| 1004 | value: torch.Tensor, |
| 1005 | key_cache: torch.Tensor, |
| 1006 | value_cache: torch.Tensor, |
| 1007 | output: torch.Tensor, |
| 1008 | cu_seqlens_q: torch.Tensor, |
| 1009 | max_seqlen_q: int, |
| 1010 | min_seqlen_q: int, |
| 1011 | max_seqlen_k: int, |
| 1012 | block_table: torch.Tensor, |
| 1013 | slot_mapping: torch.Tensor, |
| 1014 | k_scale: torch.Tensor, |
| 1015 | v_scale: torch.Tensor, |
| 1016 | ): |
| 1017 | if self.sliding_window[0] != -1: |
| 1018 | self.extend_for_sliding_window( |
| 1019 | attn_metadata, |
| 1020 | query, |
| 1021 | key_cache, |
| 1022 | value_cache, |
| 1023 | output, |
| 1024 | cu_seqlens_q, |
| 1025 | max_seqlen_q, |
| 1026 | block_table, |
| 1027 | k_scale, |
| 1028 | v_scale, |
| 1029 | ) |
| 1030 | return |
| 1031 | out, lse = rocm_aiter_ops.flash_attn_varlen_func( |
| 1032 | q=query, |
| 1033 | k=key, |
| 1034 | v=value, |
| 1035 | cu_seqlens_q=cu_seqlens_q, |
| 1036 | cu_seqlens_k=cu_seqlens_q, |
| 1037 | max_seqlen_q=max_seqlen_q, |
| 1038 | max_seqlen_k=max_seqlen_q, |
| 1039 | min_seqlen_q=min_seqlen_q, |
| 1040 | dropout_p=0.0, |
| 1041 | softmax_scale=self.scale, |
| 1042 | causal=True, |
| 1043 | window_size=self.sliding_window, |
| 1044 | alibi_slopes=self.alibi_slopes, |
| 1045 | return_lse=True, |
| 1046 | sink_ptr=self.sinks, |
| 1047 | ) |
| 1048 | assert attn_metadata.extend_metadata is not None |
| 1049 | chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata |
| 1050 | num_chunks = chunk_context_metadata.num_chunks |
| 1051 | workspace = chunk_context_metadata.workspace |
| 1052 | cu_seqlens_kv = chunk_context_metadata.cu_seq_lens_chunk |
| 1053 | max_seqlens = chunk_context_metadata.max_seq_lens |
| 1054 | chunk_starts = chunk_context_metadata.chunk_starts |
| 1055 | token_to_batch = chunk_context_metadata.token_to_batch |
| 1056 | total_token_per_batch = chunk_context_metadata.total_token_per_batch |
| 1057 | key_fetched, value_fetched = workspace[0], workspace[1] |
| 1058 | chunked_output = None |
| 1059 | chunked_lse = None |
| 1060 | for chunk_idx in range(num_chunks): |
| 1061 | cp_mha_gather_cache( |
| 1062 | key_cache=key_cache, |
| 1063 | value_cache=value_cache, |
| 1064 | key=key_fetched, |
| 1065 | value=value_fetched, |
| 1066 | block_tables=block_table, |
| 1067 | k_scales=k_scale, |
| 1068 | v_scales=v_scale, |
| 1069 | cu_seqlens_kv=cu_seqlens_kv[chunk_idx], |
| 1070 | token_to_batch=token_to_batch[chunk_idx], |
| 1071 | seq_starts=chunk_starts[chunk_idx], |
| 1072 | dequant=is_quantized_kv_cache(self.kv_cache_dtype), |
| 1073 | kv_cache_layout="SHUFFLE" |
| 1074 | if rocm_aiter_ops.is_shuffle_kv_cache_enabled() |
| 1075 | else "NHD", |
| 1076 | total_tokens=total_token_per_batch[chunk_idx], |
| 1077 | ) |
| 1078 | |
| 1079 | suf_out, suf_lse = rocm_aiter_ops.flash_attn_varlen_func( |
| 1080 | q=query, |
| 1081 | k=key_fetched, |
| 1082 | v=value_fetched, |
| 1083 | cu_seqlens_q=cu_seqlens_q, |
| 1084 | cu_seqlens_k=cu_seqlens_kv[chunk_idx], |
| 1085 | max_seqlen_q=max_seqlen_q, |
| 1086 | max_seqlen_k=max_seqlens[chunk_idx], |
| 1087 | min_seqlen_q=min_seqlen_q, |
| 1088 | dropout_p=0.0, |
| 1089 | softmax_scale=self.scale, |
| 1090 | causal=False, |
| 1091 | window_size=self.sliding_window, |
| 1092 | alibi_slopes=self.alibi_slopes, |
| 1093 | return_lse=True, |
| 1094 | sink_ptr=self.sinks, |
| 1095 | ) |
| 1096 | if chunked_output is None: |
| 1097 | chunked_output = suf_out |
| 1098 | chunked_lse = suf_lse |
| 1099 | else: |
| 1100 | tmp_output = torch.empty_like(out) |
| 1101 | tmp_lse = torch.empty_like(lse) |
| 1102 | merge_attn_states( |
| 1103 | output=tmp_output, |
| 1104 | output_lse=tmp_lse, |
| 1105 | prefix_output=chunked_output, |
| 1106 | prefix_lse=chunked_lse, |
| 1107 | suffix_output=suf_out, |
| 1108 | suffix_lse=suf_lse, |
| 1109 | ) |
| 1110 | chunked_output = tmp_output |
| 1111 | chunked_lse = tmp_lse |
| 1112 | |
| 1113 | merge_attn_states( |
| 1114 | output=output, |
| 1115 | prefix_output=chunked_output, |
| 1116 | prefix_lse=chunked_lse, |
| 1117 | suffix_output=out, |
| 1118 | suffix_lse=lse, |
| 1119 | ) |
| 1120 | |
| 1121 | def forward( |
| 1122 | self, |
| 1123 | layer: torch.nn.Module, |
| 1124 | query: torch.Tensor, |
| 1125 | key: torch.Tensor | None, |
| 1126 | value: torch.Tensor | None, |
| 1127 | kv_cache: torch.Tensor, |
| 1128 | attn_metadata: AiterFlashAttentionMetadata, |
| 1129 | output: torch.Tensor, |
| 1130 | output_scale: torch.Tensor | None = None, |
| 1131 | output_block_scale: torch.Tensor | None = None, |
| 1132 | ) -> torch.Tensor: |
| 1133 | """Forward pass with AiterFlashAttention. |
| 1134 | |
| 1135 | Args: |
| 1136 | layer: The attention layer, providing the q/k/v quantization scales. |
| 1137 | query: shape = [num_tokens, num_heads, head_size] |
| 1138 | key: shape = [num_tokens, num_kv_heads, head_size] |
| 1139 | value: shape = [num_tokens, num_kv_heads, head_size] |
| 1140 | kv_cache: shape = |
| 1141 | [num_blocks, 2, block_size, num_kv_heads, head_size] |
| 1142 | attn_metadata: Metadata for attention. |
| 1143 | output: Tensor that the attention result is written into. |
| 1144 | output_scale: Scale for fused output quantization; not supported |
| 1145 | by this backend. |
| 1146 | output_block_scale: Block scale for fused output quantization; |
| 1147 | not supported by this backend. |
| 1148 | |
| 1149 | Returns: |
| 1150 | shape = [num_tokens, num_heads * head_size] |
| 1151 | NOTE: FP8 quantization, flash-attn expect the size of |
| 1152 | {q,k,v}_descale to be (num_sequences, num_kv_heads). |
| 1153 | We use torch's .expand() to avoid duplicating values |
| 1154 | |
| 1155 | """ |
| 1156 | if output_scale is not None or output_block_scale is not None: |
| 1157 | raise NotImplementedError( |
| 1158 | "fused output quantization is not yet supported " |
| 1159 | "for AiterFlashAttentionImpl" |
| 1160 | ) |
| 1161 | |
| 1162 | if attn_metadata is None: |
| 1163 | # Profiling run. |
| 1164 | return output.fill_(0) |
| 1165 | |
| 1166 | # IMPORTANT! |
| 1167 | # NOTE(woosuk): With piece-wise CUDA graphs, this method is |
| 1168 | # executed in eager-mode PyTorch. Thus, we need to be careful |
| 1169 | # about any CPU overhead in this method. For example, `view` |
| 1170 | # and `slice` (or `[:n]`) operations are surprisingly slow even |
| 1171 | # in the case they do not invoke any GPU ops. |
| 1172 | # Minimize the PyTorch ops in this method as much as possible. |
| 1173 | # Whenever making a change in this method, please benchmark the |
| 1174 | # performance to make sure it does not introduce any overhead. |
| 1175 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 1176 | |
| 1177 | key_cache, value_cache = self._split_kv_cache(kv_cache) |
| 1178 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 1179 | key_cache = key_cache.view(current_platform.fp8_dtype()) |
| 1180 | value_cache = value_cache.view(current_platform.fp8_dtype()) |
| 1181 | |
| 1182 | # decode:extend:prefill |
| 1183 | query = query[:num_actual_tokens] |
| 1184 | if key is not None: |
| 1185 | key = key[:num_actual_tokens] |
| 1186 | if value is not None: |
| 1187 | value = value[:num_actual_tokens] |
| 1188 | |
| 1189 | output_actual_tokens = output[:num_actual_tokens] |
| 1190 | |
| 1191 | num_decodes = attn_metadata.num_decodes |
| 1192 | num_prefills = attn_metadata.num_prefills |
| 1193 | num_extends = attn_metadata.num_extends |
| 1194 | |
| 1195 | num_decode_tokens = attn_metadata.num_decode_tokens |
| 1196 | num_extend_tokens = attn_metadata.num_extend_tokens |
| 1197 | if self.kv_sharing_target_layer_name is not None and ( |
| 1198 | num_prefills > 0 or num_extends > 0 |
| 1199 | ): |
| 1200 | # Shared layers project only Q. Fetch the new K/V tokens from the |
| 1201 | # target cache for the direct prefill and extend suffix kernels. |
| 1202 | shared = attn_metadata.kv_sharing_metadata |
| 1203 | assert shared is not None |
| 1204 | key, value = shared.workspace[:, :num_actual_tokens].unbind(0) |
| 1205 | cp_mha_gather_cache( |
| 1206 | key_cache=key_cache, |
| 1207 | value_cache=value_cache, |
| 1208 | key=key[num_decode_tokens:], |
| 1209 | value=value[num_decode_tokens:], |
| 1210 | block_tables=attn_metadata.block_table[num_decodes:], |
| 1211 | k_scales=layer._k_scale, |
| 1212 | v_scales=layer._v_scale, |
| 1213 | cu_seqlens_kv=shared.query_start_loc, |
| 1214 | token_to_batch=shared.token_to_batch, |
| 1215 | seq_starts=shared.seq_starts, |
| 1216 | dequant=is_quantized_kv_cache(self.kv_cache_dtype), |
| 1217 | kv_cache_layout="SHUFFLE" |
| 1218 | if rocm_aiter_ops.is_shuffle_kv_cache_enabled() |
| 1219 | else "NHD", |
| 1220 | total_tokens=num_actual_tokens - num_decode_tokens, |
| 1221 | ) |
| 1222 | if not attn_metadata.use_cascade: |
| 1223 | # calculate for pure prefills |
| 1224 | if num_prefills > 0: |
| 1225 | assert attn_metadata.prefill_metadata is not None |
| 1226 | assert key is not None and value is not None |
| 1227 | |
| 1228 | prefill_query = query[num_decode_tokens + num_extend_tokens :] |
| 1229 | prefill_key = key[num_decode_tokens + num_extend_tokens :] |
| 1230 | prefill_value = value[num_decode_tokens + num_extend_tokens :] |
| 1231 | |
| 1232 | rocm_aiter_ops.flash_attn_varlen_func( |
| 1233 | q=prefill_query, |
| 1234 | k=prefill_key, |
| 1235 | v=prefill_value, |
| 1236 | cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc, |
| 1237 | cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc, |
| 1238 | max_seqlen_q=attn_metadata.prefill_metadata.max_query_len, |
| 1239 | max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len, |
| 1240 | min_seqlen_q=1, |
| 1241 | dropout_p=0.0, |
| 1242 | softmax_scale=self.scale, |
| 1243 | causal=attn_metadata.causal, |
| 1244 | window_size=self.sliding_window, |
| 1245 | alibi_slopes=self.alibi_slopes, |
| 1246 | out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], |
| 1247 | sink_ptr=self.sinks, |
| 1248 | ) |
| 1249 | |
| 1250 | # calculate for extends |
| 1251 | if num_extends > 0: |
| 1252 | assert attn_metadata.extend_metadata is not None |
| 1253 | assert key is not None and value is not None |
| 1254 | extend_tokens_slice = slice( |
| 1255 | num_decode_tokens, num_decode_tokens + num_extend_tokens |
| 1256 | ) |
| 1257 | extend_queries = query[extend_tokens_slice] |
| 1258 | extend_keys = key[extend_tokens_slice] |
| 1259 | extend_values = value[extend_tokens_slice] |
| 1260 | extend_outputs = output[extend_tokens_slice] |
| 1261 | k_scale = layer._k_scale |
| 1262 | v_scale = layer._v_scale |
| 1263 | if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): |
| 1264 | k_scale = attn_metadata.k_scale |
| 1265 | v_scale = attn_metadata.v_scale |
| 1266 | self.extend_forward( |
| 1267 | attn_metadata=attn_metadata, |
| 1268 | query=extend_queries, |
| 1269 | key=extend_keys, |
| 1270 | value=extend_values, |
| 1271 | key_cache=key_cache, |
| 1272 | value_cache=value_cache, |
| 1273 | output=extend_outputs, |
| 1274 | cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc, |
| 1275 | max_seqlen_q=attn_metadata.extend_metadata.max_query_len, |
| 1276 | min_seqlen_q=1, |
| 1277 | max_seqlen_k=attn_metadata.extend_metadata.max_seq_len, |
| 1278 | block_table=attn_metadata.block_table[ |
| 1279 | num_decodes : num_decodes + num_extends |
| 1280 | ], |
| 1281 | slot_mapping=attn_metadata.slot_mapping[ |
| 1282 | num_decodes : num_decodes + num_extends |
| 1283 | ], |
| 1284 | k_scale=k_scale, |
| 1285 | v_scale=v_scale, |
| 1286 | ) |
| 1287 | |
| 1288 | # calculate for decodes |
| 1289 | if num_decodes > 0: |
| 1290 | assert attn_metadata.decode_metadata is not None |
| 1291 | decode_max_query_len = attn_metadata.decode_metadata.max_query_len |
| 1292 | |
| 1293 | # Use unified_attention for the decodes the paged kernels can't |
| 1294 | # take: sliding window, sinks, or a multi-token batch |
| 1295 | # (pa_fwd_asm and paged_attention_v1 don't support sinks). |
| 1296 | if ( |
| 1297 | self.sliding_window[0] != -1 |
| 1298 | or decode_max_query_len > 1 |
| 1299 | or self.sinks is not None |
| 1300 | ): |
| 1301 | k_descale, v_descale = self._get_kv_cache_descales( |
| 1302 | layer, num_decodes |
| 1303 | ) |
| 1304 | assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( |
| 1305 | "Shuffle KV cache layout is not supported with sliding " |
| 1306 | "window, sinks, or speculative decoding (multi-token decode)." |
| 1307 | ) |
| 1308 | if not attn_metadata.causal: |
| 1309 | from aiter.ops.triton.attention.mha_v3 import ( |
| 1310 | flash_attn_with_kvcache, |
| 1311 | ) |
| 1312 | |
| 1313 | decode_query = query[:num_decode_tokens].reshape( |
| 1314 | num_decodes, |
| 1315 | decode_max_query_len, |
| 1316 | query.shape[1], |
| 1317 | query.shape[2], |
| 1318 | ) |
| 1319 | decode_out = flash_attn_with_kvcache( |
| 1320 | q=decode_query, |
| 1321 | k_cache=key_cache, |
| 1322 | v_cache=value_cache, |
| 1323 | cache_seqlens=attn_metadata.seq_lens[:num_decodes], |
| 1324 | softmax_scale=self.scale, |
| 1325 | causal=attn_metadata.causal, |
| 1326 | window_size=self.sliding_window, |
| 1327 | softcap=self.logits_soft_cap, |
| 1328 | q_descale=None, |
| 1329 | k_descale=k_descale, |
| 1330 | v_descale=v_descale, |
| 1331 | page_table=attn_metadata.block_table[:num_decodes], |
| 1332 | ) |
| 1333 | output[:num_decode_tokens].copy_( |
| 1334 | decode_out.reshape( |
| 1335 | num_decode_tokens, |
| 1336 | query.shape[1], |
| 1337 | query.shape[2], |
| 1338 | ) |
| 1339 | ) |
| 1340 | else: |
| 1341 | # Non-uniform query lengths can appear in real serving |
| 1342 | # traffic (e.g. mixed datasets). Fall back to varlen |
| 1343 | # unified_attention instead of asserting. |
| 1344 | from aiter.ops.triton.unified_attention import ( |
| 1345 | unified_attention, |
| 1346 | ) |
| 1347 | |
| 1348 | unified_attention( |
| 1349 | q=query[:num_decode_tokens], |
| 1350 | k=key_cache, |
| 1351 | v=value_cache, |
| 1352 | out=output[:num_decode_tokens], |
| 1353 | cu_seqlens_q=attn_metadata.query_start_loc[ |
| 1354 | : num_decodes + 1 |
| 1355 | ], |
| 1356 | max_seqlen_q=decode_max_query_len, |
| 1357 | seqused_k=attn_metadata.seq_lens[:num_decodes], |
| 1358 | max_seqlen_k=attn_metadata.max_seq_len, |
| 1359 | softmax_scale=self.scale, |
| 1360 | causal=True, |
| 1361 | alibi_slopes=self.alibi_slopes, |
| 1362 | window_size=self.sliding_window, |
| 1363 | block_table=attn_metadata.block_table[:num_decodes], |
| 1364 | softcap=self.logits_soft_cap, |
| 1365 | q_descale=None, |
| 1366 | k_descale=k_descale, |
| 1367 | v_descale=v_descale, |
| 1368 | sinks=self.sinks, |
| 1369 | ) |
| 1370 | return |
| 1371 | |
| 1372 | # The ll4mi kernel in paged_attention_v1 requires |
| 1373 | # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4). |
| 1374 | # For smaller head sizes or sliding window attention, |
| 1375 | # fall back to the unified_attention triton kernel which |
| 1376 | # handles both correctly. |
| 1377 | _MIN_HEAD_SIZE_FOR_LL4MI = 64 |
| 1378 | use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI |
| 1379 | |
| 1380 | if use_unified_attention: |
| 1381 | k_descale, v_descale = self._get_kv_cache_descales( |
| 1382 | layer, num_decodes |
| 1383 | ) |
| 1384 | assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( |
| 1385 | "unified_attention fallback with shuffle layout " |
| 1386 | "is not supported yet." |
| 1387 | ) |
| 1388 | from aiter.ops.triton.unified_attention import ( |
| 1389 | unified_attention, |
| 1390 | ) |
| 1391 | |
| 1392 | decode_cu_seqlens_q = attn_metadata.query_start_loc[ |
| 1393 | : num_decodes + 1 |
| 1394 | ] |
| 1395 | unified_attention( |
| 1396 | q=query[:num_decode_tokens], |
| 1397 | k=key_cache, |
| 1398 | v=value_cache, |
| 1399 | out=output[:num_decode_tokens], |
| 1400 | cu_seqlens_q=decode_cu_seqlens_q, |
| 1401 | max_seqlen_q=1, |
| 1402 | seqused_k=attn_metadata.seq_lens[:num_decodes], |
| 1403 | max_seqlen_k=attn_metadata.max_seq_len, |
| 1404 | softmax_scale=self.scale, |
| 1405 | causal=True, |
| 1406 | alibi_slopes=self.alibi_slopes, |
| 1407 | window_size=self.sliding_window, |
| 1408 | block_table=attn_metadata.block_table[:num_decodes], |
| 1409 | softcap=self.logits_soft_cap, |
| 1410 | q_descale=None, |
| 1411 | k_descale=k_descale, |
| 1412 | v_descale=v_descale, |
| 1413 | ) |
| 1414 | elif rocm_aiter_ops.is_shuffle_kv_cache_enabled(): |
| 1415 | _, num_heads, head_size = query.shape |
| 1416 | num_blocks, block_size, num_kv_heads, _ = key_cache.shape |
| 1417 | x = 16 // key_cache.element_size() |
| 1418 | new_key_cache = key_cache.reshape( |
| 1419 | num_blocks, num_kv_heads, head_size // x, block_size, x |
| 1420 | ) |
| 1421 | new_value_cache = value_cache.reshape( |
| 1422 | num_blocks, num_kv_heads, block_size // x, head_size, x |
| 1423 | ) |
| 1424 | |
| 1425 | # This kernel derives block addresses arithmetically, so it |
| 1426 | # only reads the right bytes when each side is a dense |
| 1427 | # plane of pages. |
| 1428 | assert new_key_cache.stride(0) == ( |
| 1429 | num_kv_heads * head_size * block_size |
| 1430 | ), ( |
| 1431 | "paged_attention_common needs the K/V sides as dense " |
| 1432 | "planes; the resolved KV cache layout interleaves them " |
| 1433 | "within a block." |
| 1434 | ) |
| 1435 | num_seqs = attn_metadata.seq_lens.shape[0] |
| 1436 | max_num_partitions = ( |
| 1437 | attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 |
| 1438 | ) // _PARTITION_SIZE_ROCM |
| 1439 | tmp_out = torch.empty( |
| 1440 | (num_seqs, num_heads, max_num_partitions, head_size), |
| 1441 | dtype=query.dtype, |
| 1442 | device=query.device, |
| 1443 | ) |
| 1444 | exp_sums = torch.empty( |
| 1445 | (num_seqs, num_heads, max_num_partitions), |
| 1446 | dtype=torch.float32, |
| 1447 | device=query.device, |
| 1448 | ) |
| 1449 | max_logits = torch.empty_like(exp_sums) |
| 1450 | k_qscale = ( |
| 1451 | layer._k_scale |
| 1452 | if attn_metadata.k_scale is None |
| 1453 | else attn_metadata.k_scale |
| 1454 | ) |
| 1455 | v_qscale = ( |
| 1456 | layer._v_scale |
| 1457 | if attn_metadata.v_scale is None |
| 1458 | else attn_metadata.v_scale |
| 1459 | ) |
| 1460 | rocm_aiter_ops.paged_attention_common( |
| 1461 | Q=query[:num_decode_tokens], |
| 1462 | K=new_key_cache, |
| 1463 | V=new_value_cache, |
| 1464 | tmp_out=tmp_out, |
| 1465 | max_logits=max_logits, |
| 1466 | exp_sums=exp_sums, |
| 1467 | max_seq_len=attn_metadata.max_seq_len, |
| 1468 | block_tables=attn_metadata.block_table[:num_decodes], |
| 1469 | context_lens=attn_metadata.seq_lens[:num_decodes], |
| 1470 | block_tables_stride0=attn_metadata.block_table[ |
| 1471 | :num_decodes |
| 1472 | ].stride(0), |
| 1473 | scale=self.scale, |
| 1474 | K_QScale_hip=k_qscale, |
| 1475 | V_QScale_hip=v_qscale, |
| 1476 | K_QScale_asm=k_qscale, |
| 1477 | V_QScale_asm=v_qscale, |
| 1478 | out_=output[:num_decode_tokens], |
| 1479 | kv_cache_dtype=self.kv_cache_dtype, |
| 1480 | ) |
| 1481 | else: |
| 1482 | _, num_heads, head_size = query.shape |
| 1483 | nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8 |
| 1484 | num_seqs = attn_metadata.seq_lens.shape[0] |
| 1485 | max_num_partitions = ( |
| 1486 | attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 |
| 1487 | ) // _PARTITION_SIZE_ROCM |
| 1488 | |
| 1489 | workspace_buffer = torch.empty( |
| 1490 | (num_seqs * num_heads * max_num_partitions * head_size) |
| 1491 | * nbytes_per_qo_elem |
| 1492 | + 2 * (num_seqs * num_heads * max_num_partitions) * 4, |
| 1493 | dtype=torch.uint8, |
| 1494 | device=output.device, |
| 1495 | ) |
| 1496 | |
| 1497 | # import so that aiter register the op to the namespace of |
| 1498 | # torch.ops.aiter |
| 1499 | import aiter # noqa: F401 |
| 1500 | |
| 1501 | torch.ops.aiter.paged_attention_v1( |
| 1502 | output[:num_decode_tokens], |
| 1503 | workspace_buffer, |
| 1504 | query[:num_decode_tokens], |
| 1505 | key_cache, |
| 1506 | value_cache, |
| 1507 | self.scale, |
| 1508 | attn_metadata.block_table[:num_decodes], |
| 1509 | attn_metadata.query_start_loc[:num_decodes], |
| 1510 | attn_metadata.seq_lens[:num_decodes], |
| 1511 | attn_metadata.max_seq_len, |
| 1512 | self.alibi_slopes, |
| 1513 | self.kv_cache_dtype, |
| 1514 | "NHD", |
| 1515 | self.logits_soft_cap, |
| 1516 | layer._k_scale, |
| 1517 | layer._v_scale, |
| 1518 | None, |
| 1519 | _PARTITION_SIZE_ROCM, |
| 1520 | 1, |
| 1521 | self.sliding_window[0] + 1, |
| 1522 | ) |
| 1523 | else: |
| 1524 | raise NotImplementedError( |
| 1525 | "Cascade attention is not implemented for ROCM AITER" |
| 1526 | ) |
| 1527 | |
| 1528 | return output |
| 1529 | |
| 1530 | def _split_kv_cache( |
| 1531 | self, kv_cache: torch.Tensor |
| 1532 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1533 | if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): |
| 1534 | # (B, 2, N, H*hs) -> one (B, N, H, hs) per side, which is what the |
| 1535 | # shuffle read/write kernels reinterpret in place. The heads are |
| 1536 | # folded into the content dim by customize_spec. |
| 1537 | key_cache, value_cache = kv_cache.unflatten( |
| 1538 | -1, (self.num_kv_heads, self.head_size) |
| 1539 | ).unbind(1) |
| 1540 | return key_cache, value_cache |
| 1541 | # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) |
| 1542 | return kv_cache.transpose(1, 2).split(self.head_size, dim=-1) |
| 1543 | |
| 1544 | def do_kv_cache_update( |
| 1545 | self, |
| 1546 | layer: AttentionLayer, |
| 1547 | key: torch.Tensor, |
| 1548 | value: torch.Tensor, |
| 1549 | kv_cache: torch.Tensor, |
| 1550 | slot_mapping: torch.Tensor, |
| 1551 | ): |
| 1552 | key_cache, value_cache = self._split_kv_cache(kv_cache) |
| 1553 | |
| 1554 | # key and value may be None in the case of cross attention. They are |
| 1555 | # calculated once based on the output from the encoder and then cached |
| 1556 | # in KV cache. |
| 1557 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 1558 | key_cache = key_cache.view(current_platform.fp8_dtype()) |
| 1559 | value_cache = value_cache.view(current_platform.fp8_dtype()) |
| 1560 | # Reshape the input keys and values and store them in the cache. |
| 1561 | # Skip this if sharing KV cache with an earlier attention layer. |
| 1562 | # NOTE(woosuk): Here, key and value are padded while slot_mapping |
| 1563 | # is not padded. However, we don't need to do |
| 1564 | # key[:num_actual_tokens] and value[:num_actual_tokens] because |
| 1565 | # the reshape_and_cache_flash op uses the slot_mapping's shape |
| 1566 | # to determine the number of actual tokens. |
| 1567 | if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): |
| 1568 | # We may calculate per token quant scale in |
| 1569 | # reshape_and_cache_shuffle_triton which might differ from |
| 1570 | # vllm's style when shuffle layout is used. |
| 1571 | k_scale = layer._k_scale |
| 1572 | v_scale = layer._v_scale |
| 1573 | assert k_scale is not None and v_scale is not None, ( |
| 1574 | "k_scale and v_scale are required for shuffled update" |
| 1575 | ) |
| 1576 | reshape_and_cache_shuffle_triton( |
| 1577 | key, |
| 1578 | value, |
| 1579 | key_cache, |
| 1580 | value_cache, |
| 1581 | slot_mapping, |
| 1582 | self.kv_cache_dtype, |
| 1583 | k_scale, |
| 1584 | v_scale, |
| 1585 | ) |
| 1586 | else: |
| 1587 | torch.ops._C_cache_ops.reshape_and_cache_flash( |
| 1588 | key, |
| 1589 | value, |
| 1590 | key_cache, |
| 1591 | value_cache, |
| 1592 | slot_mapping, |
| 1593 | self.kv_cache_dtype, |
| 1594 | layer._k_scale, |
| 1595 | layer._v_scale, |
| 1596 | ) |
| 1597 | |
| 1598 | def fused_rope_kvcache_supported(self): |
| 1599 | # Only support fusion when shuffle KV cache layout is not used; |
| 1600 | # shuffle layout uses a different cache update path. |
| 1601 | return ( |
| 1602 | rocm_aiter_ops.is_enabled() |
| 1603 | and not rocm_aiter_ops.is_shuffle_kv_cache_enabled() |
| 1604 | ) |
| 1605 | |
| 1606 | def fused_qk_norm_rope_kvcache_supported(self): |
| 1607 | # Only fuse when shuffle layout is off; the shuffle write path uses a |
| 1608 | # dedicated cache update, mirroring fused_rope_kvcache_supported. |
| 1609 | return ( |
| 1610 | rocm_aiter_ops.is_enabled() |
| 1611 | and not rocm_aiter_ops.is_shuffle_kv_cache_enabled() |
| 1612 | ) |
| 1613 | |
| 1614 | def do_qk_norm_rope_kvcache_update( |
| 1615 | self, |
| 1616 | layer: AttentionLayer, |
| 1617 | qkv: torch.Tensor, |
| 1618 | q_out: torch.Tensor, |
| 1619 | k_out: torch.Tensor, |
| 1620 | positions: torch.Tensor, |
| 1621 | q_weight: torch.Tensor, |
| 1622 | k_weight: torch.Tensor, |
| 1623 | rms_norm_eps: float, |
| 1624 | cos_sin_cache: torch.Tensor, |
| 1625 | is_neox: bool, |
| 1626 | kv_cache: torch.Tensor, |
| 1627 | layer_slot_mapping: torch.Tensor, |
| 1628 | ): |
| 1629 | key_cache, value_cache = self._split_kv_cache(kv_cache) |
| 1630 | rocm_aiter_ops.do_qk_norm_rope_kvcache_update( |
| 1631 | qkv=qkv, |
| 1632 | q_weight=q_weight, |
| 1633 | k_weight=k_weight, |
| 1634 | cos_sin_cache=cos_sin_cache, |
| 1635 | positions=positions, |
| 1636 | num_heads_q=self.num_heads, |
| 1637 | num_heads_k=self.num_kv_heads, |
| 1638 | head_dim=self.head_size, |
| 1639 | is_neox=is_neox, |
| 1640 | rms_norm_eps=rms_norm_eps, |
| 1641 | q_out=q_out, |
| 1642 | k_out=k_out, |
| 1643 | key_cache=key_cache, |
| 1644 | value_cache=value_cache, |
| 1645 | slot_mapping=layer_slot_mapping, |
| 1646 | k_scale=layer._k_scale_cpu, |
| 1647 | v_scale=layer._v_scale_cpu, |
| 1648 | kv_cache_dtype=self.kv_cache_dtype, |
| 1649 | use_shuffle_layout=rocm_aiter_ops.is_shuffle_kv_cache_enabled(), |
| 1650 | ) |
| 1651 | |
| 1652 | def do_rope_and_kv_cache_update( |
| 1653 | self, |
| 1654 | layer: AttentionLayer, |
| 1655 | query: torch.Tensor, |
| 1656 | key: torch.Tensor, |
| 1657 | value: torch.Tensor, |
| 1658 | positions: torch.Tensor, |
| 1659 | cos_sin_cache: torch.Tensor, |
| 1660 | is_neox: bool, |
| 1661 | kv_cache: torch.Tensor, |
| 1662 | layer_slot_mapping: torch.Tensor, |
| 1663 | ): |
| 1664 | # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) |
| 1665 | key_cache, value_cache = self._split_kv_cache(kv_cache) |
| 1666 | flash_layout = True |
| 1667 | |
| 1668 | is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) |
| 1669 | if is_fp8_kv_cache: |
| 1670 | key_cache = key_cache.view(current_platform.fp8_dtype()) |
| 1671 | value_cache = value_cache.view(current_platform.fp8_dtype()) |
| 1672 | |
| 1673 | rocm_aiter_ops.triton_rope_and_cache( |
| 1674 | query, |
| 1675 | key, |
| 1676 | value, |
| 1677 | positions, |
| 1678 | cos_sin_cache, |
| 1679 | is_neox, |
| 1680 | key_cache, |
| 1681 | value_cache, |
| 1682 | layer_slot_mapping, |
| 1683 | layer._k_scale, |
| 1684 | layer._v_scale, |
| 1685 | flash_layout, |
| 1686 | is_fp8_kv_cache, |
| 1687 | ) |