| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention layer with PagedAttention and Triton prefix prefill.""" |
| 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 |
| 12 | from vllm.config.cache import CacheDType |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 15 | QuantKey, |
| 16 | kFp8StaticTensorSym, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache |
| 20 | from vllm.v1.attention.backend import ( |
| 21 | AttentionBackend, |
| 22 | AttentionCGSupport, |
| 23 | AttentionImpl, |
| 24 | AttentionLayer, |
| 25 | AttentionMetadataBuilder, |
| 26 | AttentionType, |
| 27 | CommonAttentionMetadata, |
| 28 | MultipleOf, |
| 29 | ) |
| 30 | from vllm.v1.attention.ops.chunked_prefill_paged_decode import ( |
| 31 | chunked_prefill_paged_decode, |
| 32 | has_native_kv_cache_layout, |
| 33 | ) |
| 34 | from vllm.v1.attention.ops.paged_attn import PagedAttention |
| 35 | from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( |
| 36 | triton_reshape_and_cache_flash, |
| 37 | ) |
| 38 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout |
| 39 | |
| 40 | logger = init_logger(__name__) |
| 41 | |
| 42 | |
| 43 | @dataclass |
| 44 | class RocmAttentionMetadata: |
| 45 | # NOTE(sang): Definition of context_len, query_len, and seq_len. |
| 46 | # |---------- N-1 iteration --------| |
| 47 | # |---------------- N iteration ---------------------| |
| 48 | # |- tokenA -|......................|-- newTokens ---| |
| 49 | # |---------- context_len ----------| |
| 50 | # |-------------------- seq_len ---------------------| |
| 51 | # |-- query_len ---| |
| 52 | |
| 53 | num_actual_tokens: int # Number of tokens excluding padding. |
| 54 | max_query_len: int |
| 55 | query_start_loc: torch.Tensor |
| 56 | max_seq_len: int |
| 57 | seq_lens: torch.Tensor |
| 58 | block_table: torch.Tensor |
| 59 | slot_mapping: torch.Tensor |
| 60 | |
| 61 | # For cascade attention. |
| 62 | use_cascade: bool |
| 63 | common_prefix_len: int |
| 64 | cu_prefix_query_lens: torch.Tensor | None |
| 65 | prefix_kv_lens: torch.Tensor | None |
| 66 | suffix_kv_lens: torch.Tensor | None |
| 67 | |
| 68 | # Optional aot scheduling |
| 69 | scheduler_metadata: torch.Tensor | None = None |
| 70 | prefix_scheduler_metadata: torch.Tensor | None = None |
| 71 | |
| 72 | # DFlash drafting sets this to False via CommonAttentionMetadata. |
| 73 | causal: bool = True |
| 74 | |
| 75 | |
| 76 | class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadata]): |
| 77 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 78 | |
| 79 | def __init__( |
| 80 | self, |
| 81 | kv_cache_spec: AttentionSpec, |
| 82 | layer_names: list[str], |
| 83 | vllm_config: VllmConfig, |
| 84 | device: torch.device, |
| 85 | ): |
| 86 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 87 | |
| 88 | self.block_size = kv_cache_spec.block_size |
| 89 | |
| 90 | model_config = vllm_config.model_config |
| 91 | self.num_heads_q = model_config.get_num_attention_heads( |
| 92 | vllm_config.parallel_config |
| 93 | ) |
| 94 | self.num_heads_kv = model_config.get_num_kv_heads(vllm_config.parallel_config) |
| 95 | self.headdim = model_config.get_head_size() |
| 96 | |
| 97 | def build_for_cudagraph_capture( |
| 98 | self, common_attn_metadata: CommonAttentionMetadata |
| 99 | ) -> RocmAttentionMetadata: |
| 100 | attn_metadata = self.build(0, common_attn_metadata) |
| 101 | # When doing full graph capture, setting seq_lens to |
| 102 | # max_model_len will cause graph capture to be extremely |
| 103 | # slow, so here we set it to 1. |
| 104 | attn_metadata.seq_lens.fill_(1) |
| 105 | |
| 106 | # Zero device query start locations to avoid invalid memory access in |
| 107 | # the prefix prefill kernel during graph capture (#25985). |
| 108 | common_attn_metadata.query_start_loc.zero_() |
| 109 | |
| 110 | return attn_metadata |
| 111 | |
| 112 | def build( |
| 113 | self, |
| 114 | common_prefix_len: int, |
| 115 | common_attn_metadata: CommonAttentionMetadata, |
| 116 | fast_build: bool = False, |
| 117 | ) -> RocmAttentionMetadata: |
| 118 | num_actual_tokens = common_attn_metadata.num_actual_tokens |
| 119 | max_query_len = common_attn_metadata.max_query_len |
| 120 | |
| 121 | max_seq_len = common_attn_metadata.max_seq_len |
| 122 | query_start_loc = common_attn_metadata.query_start_loc |
| 123 | seq_lens = common_attn_metadata.seq_lens |
| 124 | block_table_tensor = common_attn_metadata.block_table_tensor |
| 125 | slot_mapping = common_attn_metadata.slot_mapping |
| 126 | |
| 127 | use_cascade = common_prefix_len > 0 |
| 128 | |
| 129 | if use_cascade: |
| 130 | cu_prefix_query_lens = torch.tensor( |
| 131 | [0, num_actual_tokens], dtype=torch.int32, device=self.device |
| 132 | ) |
| 133 | prefix_kv_lens = torch.tensor( |
| 134 | [common_prefix_len], dtype=torch.int32, device=self.device |
| 135 | ) |
| 136 | suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len |
| 137 | suffix_kv_lens = suffix_kv_lens.to(self.device) |
| 138 | else: |
| 139 | cu_prefix_query_lens = None |
| 140 | prefix_kv_lens = None |
| 141 | suffix_kv_lens = None |
| 142 | prefix_scheduler_metadata = None |
| 143 | |
| 144 | attn_metadata = RocmAttentionMetadata( |
| 145 | num_actual_tokens=num_actual_tokens, |
| 146 | max_query_len=max_query_len, |
| 147 | query_start_loc=query_start_loc, |
| 148 | max_seq_len=max_seq_len, |
| 149 | seq_lens=seq_lens, |
| 150 | block_table=block_table_tensor, |
| 151 | slot_mapping=slot_mapping, |
| 152 | use_cascade=use_cascade, |
| 153 | common_prefix_len=common_prefix_len, |
| 154 | cu_prefix_query_lens=cu_prefix_query_lens, |
| 155 | prefix_kv_lens=prefix_kv_lens, |
| 156 | suffix_kv_lens=suffix_kv_lens, |
| 157 | prefix_scheduler_metadata=prefix_scheduler_metadata, |
| 158 | causal=common_attn_metadata.causal, |
| 159 | ) |
| 160 | return attn_metadata |
| 161 | |
| 162 | |
| 163 | class RocmAttentionBackend(AttentionBackend): |
| 164 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 165 | torch.float16, |
| 166 | torch.bfloat16, |
| 167 | torch.float32, |
| 168 | ] |
| 169 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 170 | "auto", |
| 171 | "float16", |
| 172 | "bfloat16", |
| 173 | "fp8", |
| 174 | "fp8_e4m3", |
| 175 | "fp8_e5m2", |
| 176 | ] |
| 177 | |
| 178 | @staticmethod |
| 179 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 180 | # ROCM paged attention native C++ kernel only supports block sizes 16 and 32 |
| 181 | # due to shared memory (LDS) constraints on AMD GPUs. |
| 182 | # See csrc/rocm/attention.cu CALL_CUSTOM_LAUNCHER_BLK macro. |
| 183 | # However, vLLM allows support for any multiple of 16 via the Triton path. |
| 184 | # As addressed in PR: https://github.com/vllm-project/vllm/pull/31380, |
| 185 | # non-standard models (like qwen3-next with block_size 544, or qwen3_5 |
| 186 | # with 784 and 1056) are dynamically routed to our optimized Triton kernel |
| 187 | # in `do_kv_cache_update`. |
| 188 | return [MultipleOf(16)] |
| 189 | |
| 190 | @classmethod |
| 191 | def get_supported_head_sizes(cls) -> list[int]: |
| 192 | return [32, 64, 80, 96, 128, 160, 192, 224, 256] |
| 193 | |
| 194 | @classmethod |
| 195 | def supports_mm_prefix(cls) -> bool: |
| 196 | # Not implemented |
| 197 | return False |
| 198 | |
| 199 | @classmethod |
| 200 | def supports_sink(cls) -> bool: |
| 201 | # ROCM custom attention kernel does not support sinks. |
| 202 | # Callink this backend with sinks will cause it to fall back to the Triton |
| 203 | # kernel, which is less efficient than the proper triton backends. |
| 204 | return False |
| 205 | |
| 206 | @classmethod |
| 207 | def supports_non_causal(cls) -> bool: |
| 208 | return True |
| 209 | |
| 210 | @classmethod |
| 211 | def supports_kv_connector(cls) -> bool: |
| 212 | # ROCM_ATTN uses (2, num_blocks, ...) KV cache layout which is |
| 213 | # incompatible with KV connectors that require blocks-first layout. |
| 214 | return False |
| 215 | |
| 216 | forward_includes_kv_cache_update: bool = False |
| 217 | |
| 218 | @staticmethod |
| 219 | def get_name() -> str: |
| 220 | return "ROCM_ATTN" |
| 221 | |
| 222 | @classmethod |
| 223 | def supports_sliding_window(cls) -> bool: |
| 224 | return True |
| 225 | |
| 226 | @staticmethod |
| 227 | def get_impl_cls() -> type["RocmAttentionImpl"]: |
| 228 | return RocmAttentionImpl |
| 229 | |
| 230 | @classmethod |
| 231 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 232 | """ENCODER_DECODER is not supported because |
| 233 | chunked_prefill_paged_decode's prefill kernel (context_attention_fwd) |
| 234 | assumes self-attention semantics: it treats passed K/V as new tokens |
| 235 | to mix with cached K/V. For cross-attention layers the encoder K/V |
| 236 | are already fully cached, so mixing them again produces incorrect |
| 237 | results when max_query_len > 1 (e.g. beam search). |
| 238 | """ |
| 239 | return attn_type in ( |
| 240 | AttentionType.DECODER, |
| 241 | AttentionType.ENCODER, |
| 242 | AttentionType.ENCODER_ONLY, |
| 243 | ) |
| 244 | |
| 245 | @classmethod |
| 246 | def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: |
| 247 | """K and V as two head groups so the native HIP kernels address each side |
| 248 | as one contiguous region (x-packed interior applied in split_kv_cache).""" |
| 249 | if spec.state_content_bytes is not None: |
| 250 | return spec |
| 251 | assert spec.head_size == spec.head_size_v, ( |
| 252 | "Separate K/V head groups require symmetric K/V head sizes." |
| 253 | ) |
| 254 | return replace( |
| 255 | spec, |
| 256 | num_head_slots=2, |
| 257 | state_content_bytes=spec.num_kv_heads |
| 258 | * spec.head_size |
| 259 | * get_dtype_size(spec.dtype), |
| 260 | ) |
| 261 | |
| 262 | @classmethod |
| 263 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 264 | # The native HIP kernels hardcode group-contiguous block addressing, so |
| 265 | # they need the K and V groups to span all blocks (H outermost within the layer) |
| 266 | # The Triton fallbacks are stride-aware, so block-interior LBHNC is also |
| 267 | # supported, just without the native kernels. |
| 268 | return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC) |
| 269 | |
| 270 | @staticmethod |
| 271 | def use_cascade_attention(*args, **kwargs) -> bool: |
| 272 | return False |
| 273 | |
| 274 | @staticmethod |
| 275 | def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]: |
| 276 | return RocmAttentionMetadataBuilder |
| 277 | |
| 278 | |
| 279 | class RocmAttentionImpl(AttentionImpl): |
| 280 | def fused_output_quant_supported(self, quant_key: QuantKey): |
| 281 | return quant_key == kFp8StaticTensorSym |
| 282 | |
| 283 | def __init__( |
| 284 | self, |
| 285 | num_heads: int, |
| 286 | head_size: int, |
| 287 | scale: float, |
| 288 | num_kv_heads: int, |
| 289 | alibi_slopes: list[float] | None, |
| 290 | sliding_window: int | None, |
| 291 | kv_cache_dtype: str, |
| 292 | logits_soft_cap: float | None = None, |
| 293 | attn_type: AttentionType = AttentionType.DECODER, |
| 294 | kv_sharing_target_layer_name: int | None = None, |
| 295 | sinks: torch.Tensor | None = None, |
| 296 | ) -> None: |
| 297 | self.attn_type = attn_type |
| 298 | self.num_heads = num_heads |
| 299 | self.head_size = head_size |
| 300 | self.scale = float(scale) |
| 301 | self.num_kv_heads = num_kv_heads |
| 302 | if alibi_slopes is not None: |
| 303 | alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) |
| 304 | self.alibi_slopes = alibi_slopes |
| 305 | if sliding_window is None: |
| 306 | self.sliding_window = (-1, -1) |
| 307 | elif attn_type in (AttentionType.ENCODER, AttentionType.ENCODER_ONLY): |
| 308 | self.sliding_window = (sliding_window - 1, sliding_window - 1) |
| 309 | else: |
| 310 | self.sliding_window = (sliding_window - 1, 0) |
| 311 | self.kv_cache_dtype = kv_cache_dtype |
| 312 | if logits_soft_cap is None: |
| 313 | # In flash-attn, setting logits_soft_cap as 0 means no soft cap. |
| 314 | logits_soft_cap = 0 |
| 315 | self.logits_soft_cap = logits_soft_cap |
| 316 | self.kv_sharing_target_layer_name = kv_sharing_target_layer_name |
| 317 | |
| 318 | self.num_queries_per_kv = self.num_heads // self.num_kv_heads |
| 319 | |
| 320 | self.fp8_dtype = current_platform.fp8_dtype() |
| 321 | |
| 322 | self.sinks = sinks |
| 323 | if sinks is not None: |
| 324 | assert sinks.shape[0] == num_heads, ( |
| 325 | "Sinks must have the same number of heads as the number of " |
| 326 | f"heads in the layer. Sinks shape: {sinks.shape}, " |
| 327 | f"num_heads: {num_heads}." |
| 328 | ) |
| 329 | |
| 330 | def _forward_encoder_attention( |
| 331 | self, |
| 332 | query: torch.Tensor, |
| 333 | key: torch.Tensor, |
| 334 | value: torch.Tensor, |
| 335 | output: torch.Tensor, |
| 336 | attn_metadata: RocmAttentionMetadata, |
| 337 | layer: torch.nn.Module, |
| 338 | ) -> torch.Tensor: |
| 339 | """Forward pass for encoder attention without KV cache. |
| 340 | |
| 341 | Args: |
| 342 | query: shape = [num_encoder_tokens, num_heads, head_size] |
| 343 | key: shape = [num_encoder_tokens, num_kv_heads, head_size] |
| 344 | value: shape = [num_encoder_tokens, num_kv_heads, head_size] |
| 345 | output: shape = [num_encoder_tokens, num_heads, head_size] |
| 346 | attn_metadata: Encoder attention metadata |
| 347 | layer: The attention layer |
| 348 | |
| 349 | """ |
| 350 | # For encoder attention, process FP8 quantization if needed |
| 351 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 352 | raise NotImplementedError( |
| 353 | "quantization is not supported for encoder attention" |
| 354 | ) |
| 355 | |
| 356 | # Use encoder-specific metadata for sequence information |
| 357 | query_start_loc = attn_metadata.query_start_loc |
| 358 | seq_lens = attn_metadata.seq_lens |
| 359 | max_query_len = attn_metadata.max_query_len |
| 360 | |
| 361 | # Call flash attention directly on Q, K, V tensors |
| 362 | from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd |
| 363 | |
| 364 | context_attention_fwd( |
| 365 | q=query, |
| 366 | k=key, |
| 367 | v=value, |
| 368 | o=output, |
| 369 | b_start_loc=query_start_loc, |
| 370 | b_seq_len=seq_lens, |
| 371 | max_input_len=max_query_len, |
| 372 | is_causal=False, |
| 373 | softmax_scale=self.scale, |
| 374 | sliding_window_q=self.sliding_window[0], |
| 375 | sliding_window_k=self.sliding_window[1], |
| 376 | sinks=self.sinks, |
| 377 | ) |
| 378 | return output |
| 379 | |
| 380 | def forward( |
| 381 | self, |
| 382 | layer: torch.nn.Module, |
| 383 | query: torch.Tensor, |
| 384 | key: torch.Tensor, |
| 385 | value: torch.Tensor, |
| 386 | kv_cache: torch.Tensor, |
| 387 | attn_metadata: RocmAttentionMetadata, |
| 388 | output: torch.Tensor, |
| 389 | output_scale: torch.Tensor | None = None, |
| 390 | output_block_scale: torch.Tensor | None = None, |
| 391 | ) -> torch.Tensor: |
| 392 | """Forward pass with FlashAttention. |
| 393 | |
| 394 | Args: |
| 395 | layer: The attention layer, providing the q/k/v quantization scales. |
| 396 | query: shape = [num_tokens, num_heads, head_size] |
| 397 | key: shape = [num_tokens, num_kv_heads, head_size] |
| 398 | value: shape = [num_tokens, num_kv_heads, head_size] |
| 399 | kv_cache: logical [num_blocks, 2, block_size, num_kv_heads * |
| 400 | head_size] under LHBNC (physically K/V-group-first) |
| 401 | attn_metadata: Metadata for attention. |
| 402 | output: Tensor that the attention result is written into. |
| 403 | output_scale: Scale for fused output quantization. |
| 404 | output_block_scale: Block scale for fused output quantization; |
| 405 | not supported by this backend. |
| 406 | |
| 407 | Returns: |
| 408 | shape = [num_tokens, num_heads * head_size] |
| 409 | |
| 410 | """ |
| 411 | if output_block_scale is not None: |
| 412 | raise NotImplementedError( |
| 413 | "fused block_scale output quantization is not yet supported" |
| 414 | " for RocmAttentionImpl" |
| 415 | ) |
| 416 | |
| 417 | if attn_metadata is None: |
| 418 | # Profiling run. |
| 419 | return output.fill_(0) |
| 420 | |
| 421 | assert attn_metadata.use_cascade is False |
| 422 | |
| 423 | # IMPORTANT! |
| 424 | # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in |
| 425 | # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead |
| 426 | # in this method. For example, `view` and `slice` (or `[:n]`) operations |
| 427 | # are surprisingly slow even in the case they do not invoke any GPU ops. |
| 428 | # Minimize the PyTorch ops in this method as much as possible. |
| 429 | # Whenever making a change in this method, please benchmark the |
| 430 | # performance to make sure it does not introduce any overhead. |
| 431 | |
| 432 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 433 | |
| 434 | if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 435 | return self._forward_encoder_attention( |
| 436 | query[:num_actual_tokens], |
| 437 | key[:num_actual_tokens], |
| 438 | value[:num_actual_tokens], |
| 439 | output[:num_actual_tokens], |
| 440 | attn_metadata, |
| 441 | layer, |
| 442 | ) |
| 443 | |
| 444 | # The bound view is logical [B, 2, N, H*hs]; split_kv_cache expects |
| 445 | # the K/V groups first. |
| 446 | key_cache, value_cache = PagedAttention.split_kv_cache( |
| 447 | kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size |
| 448 | ) |
| 449 | |
| 450 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 451 | key_cache = key_cache.view(self.fp8_dtype) |
| 452 | value_cache = value_cache.view(self.fp8_dtype) |
| 453 | # q_scale only applies to an fp8 query; this path keeps the query |
| 454 | # in full precision, so a non-1.0 q_scale is not applicable here. |
| 455 | if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0: |
| 456 | raise NotImplementedError( |
| 457 | "A non 1.0 q_scale with an fp8 query is not currently " |
| 458 | "supported by RocmAttentionImpl." |
| 459 | ) |
| 460 | |
| 461 | cu_seqlens_q = attn_metadata.query_start_loc |
| 462 | seqused_k = attn_metadata.seq_lens |
| 463 | max_seqlen_q = attn_metadata.max_query_len |
| 464 | max_seqlen_k = attn_metadata.max_seq_len |
| 465 | block_table = attn_metadata.block_table |
| 466 | |
| 467 | # Compute attention and update output up to `num_actual_tokens`. |
| 468 | chunked_prefill_paged_decode( |
| 469 | query=query[:num_actual_tokens], |
| 470 | key=key[:num_actual_tokens] if key is not None else None, |
| 471 | value=value[:num_actual_tokens] if value is not None else None, |
| 472 | output=output[:num_actual_tokens], |
| 473 | kv_cache_dtype=self.kv_cache_dtype, |
| 474 | key_cache=key_cache, |
| 475 | value_cache=value_cache, |
| 476 | block_table=block_table, |
| 477 | query_start_loc=cu_seqlens_q, |
| 478 | seq_lens=seqused_k, |
| 479 | max_seq_len=max_seqlen_k, |
| 480 | max_query_len=max_seqlen_q, |
| 481 | k_scale=layer._k_scale, |
| 482 | v_scale=layer._v_scale, |
| 483 | alibi_slopes=self.alibi_slopes, |
| 484 | sliding_window=self.sliding_window[0], |
| 485 | sm_scale=self.scale, |
| 486 | output_scale=output_scale, |
| 487 | sinks=self.sinks, |
| 488 | causal=attn_metadata.causal, |
| 489 | ) |
| 490 | |
| 491 | return output |
| 492 | |
| 493 | def do_kv_cache_update( |
| 494 | self, |
| 495 | layer: AttentionLayer, |
| 496 | key: torch.Tensor, |
| 497 | value: torch.Tensor, |
| 498 | kv_cache: torch.Tensor, |
| 499 | slot_mapping: torch.Tensor, |
| 500 | ): |
| 501 | if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 502 | return |
| 503 | key_cache, value_cache = PagedAttention.split_kv_cache( |
| 504 | kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size |
| 505 | ) |
| 506 | |
| 507 | # Reshape the input keys and values and store them in the cache. |
| 508 | # Get the actual block_size from value_cache |
| 509 | # value_cache shape: [num_blocks, num_heads, head_size, block_size] |
| 510 | block_size = value_cache.shape[3] |
| 511 | has_native_layout = has_native_kv_cache_layout(key_cache, value_cache) |
| 512 | |
| 513 | if block_size in (16, 32) and has_native_layout: |
| 514 | # Normal 16, 32 with contiguous blocks: use vLLM native HIP C++ logic. |
| 515 | PagedAttention.write_to_paged_cache( |
| 516 | key, |
| 517 | value, |
| 518 | key_cache, |
| 519 | value_cache, |
| 520 | slot_mapping, |
| 521 | self.kv_cache_dtype, |
| 522 | layer._k_scale, |
| 523 | layer._v_scale, |
| 524 | ) |
| 525 | else: |
| 526 | # Non-standard blocks and hybrid attention/Mamba layouts need the |
| 527 | # stride-aware Triton writer. The native reshape_and_cache kernel |
| 528 | # assumes contiguous block storage and writes to the wrong hybrid |
| 529 | # cache blocks. |
| 530 | triton_reshape_and_cache_flash( |
| 531 | key, |
| 532 | value, |
| 533 | key_cache, |
| 534 | value_cache, |
| 535 | slot_mapping, |
| 536 | self.kv_cache_dtype, |
| 537 | layer._k_scale, |
| 538 | layer._v_scale, |
| 539 | ) |
| 540 | |
| 541 | def fused_rope_kvcache_supported(self): |
| 542 | return rocm_aiter_ops.is_enabled() |
| 543 | |
| 544 | def do_rope_and_kv_cache_update( |
| 545 | self, |
| 546 | layer: AttentionLayer, |
| 547 | query: torch.Tensor, |
| 548 | key: torch.Tensor, |
| 549 | value: torch.Tensor, |
| 550 | positions: torch.Tensor, |
| 551 | cos_sin_cache: torch.Tensor, |
| 552 | is_neox: bool, |
| 553 | kv_cache: torch.Tensor, |
| 554 | layer_slot_mapping: torch.Tensor, |
| 555 | ): |
| 556 | if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 557 | return |
| 558 | key_cache, value_cache = PagedAttention.split_kv_cache( |
| 559 | kv_cache.transpose(0, 1), |
| 560 | layer.num_kv_heads, # type: ignore[attr-defined] |
| 561 | layer.head_size, # type: ignore[attr-defined] |
| 562 | ) |
| 563 | flash_layout = False |
| 564 | |
| 565 | is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) |
| 566 | if is_fp8_kv_cache: |
| 567 | key_cache = key_cache.view(self.fp8_dtype) |
| 568 | value_cache = value_cache.view(self.fp8_dtype) |
| 569 | |
| 570 | rocm_aiter_ops.triton_rope_and_cache( |
| 571 | query, |
| 572 | key, |
| 573 | value, |
| 574 | positions, |
| 575 | cos_sin_cache, |
| 576 | is_neox, |
| 577 | key_cache, |
| 578 | value_cache, |
| 579 | layer_slot_mapping, |
| 580 | layer._k_scale, |
| 581 | layer._v_scale, |
| 582 | flash_layout, |
| 583 | is_fp8_kv_cache, |
| 584 | ) |