| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """High-Performance Triton-only Attention layer.""" |
| 4 | |
| 5 | from dataclasses import dataclass, replace |
| 6 | from typing import ClassVar |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm._aiter_ops import rocm_aiter_ops |
| 12 | from vllm.config import CUDAGraphMode, VllmConfig |
| 13 | from vllm.config.cache import CacheDType |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 16 | QuantKey, |
| 17 | kFp8StaticTensorSym, |
| 18 | ) |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.platforms.interface import DeviceCapability |
| 21 | from vllm.utils.math_utils import cdiv, next_power_of_2 |
| 22 | from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache |
| 23 | from vllm.v1.attention.backend import ( |
| 24 | AttentionBackend, |
| 25 | AttentionCGSupport, |
| 26 | AttentionImpl, |
| 27 | AttentionLayer, |
| 28 | AttentionMetadataBuilder, |
| 29 | AttentionType, |
| 30 | CommonAttentionMetadata, |
| 31 | MultipleOf, |
| 32 | ) |
| 33 | from vllm.v1.attention.backends.utils import ( |
| 34 | compute_mm_prefix_range_tensor, |
| 35 | get_num_attention_heads_from_layers, |
| 36 | ) |
| 37 | from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd |
| 38 | from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( |
| 39 | triton_reshape_and_cache_flash, |
| 40 | triton_reshape_and_cache_flash_per_token_head_quant, |
| 41 | ) |
| 42 | from vllm.v1.attention.ops.triton_unified_attention import unified_attention |
| 43 | from vllm.v1.kv_cache_interface import ( |
| 44 | AttentionSpec, |
| 45 | KVQuantMode, |
| 46 | get_kv_quant_mode, |
| 47 | ) |
| 48 | |
| 49 | logger = init_logger(__name__) |
| 50 | |
| 51 | |
| 52 | # constants |
| 53 | MIN_LAUNCH_GRID_SIZE_2D = 128 # Minimum launch grid size of 2D kernel |
| 54 | NUM_PAR_SOFTMAX_SEGMENTS = 16 # Number of parallel tiled softmax segments |
| 55 | |
| 56 | |
| 57 | @dataclass |
| 58 | class TritonAttentionMetadata: |
| 59 | # NOTE(sang): Definition of context_len, query_len, and seq_len. |
| 60 | # |---------- N-1 iteration --------| |
| 61 | # |---------------- N iteration ---------------------| |
| 62 | # |- tokenA -|......................|-- newTokens ---| |
| 63 | # |---------- context_len ----------| |
| 64 | # |-------------------- seq_len ---------------------| |
| 65 | # |-- query_len ---| |
| 66 | |
| 67 | num_actual_tokens: int # Number of tokens excluding padding. |
| 68 | max_query_len: int |
| 69 | query_start_loc: torch.Tensor |
| 70 | max_seq_len: int |
| 71 | seq_lens: torch.Tensor |
| 72 | block_table: torch.Tensor |
| 73 | slot_mapping: torch.Tensor |
| 74 | |
| 75 | seq_threshold_3D: int |
| 76 | num_par_softmax_segments: int |
| 77 | softmax_segm_output: torch.Tensor |
| 78 | softmax_segm_max: torch.Tensor |
| 79 | softmax_segm_expsum: torch.Tensor |
| 80 | |
| 81 | causal: bool | torch.Tensor |
| 82 | |
| 83 | # For cascade attention. |
| 84 | use_cascade: bool |
| 85 | common_prefix_len: int |
| 86 | cu_prefix_query_lens: torch.Tensor | None |
| 87 | prefix_kv_lens: torch.Tensor | None |
| 88 | suffix_kv_lens: torch.Tensor | None |
| 89 | |
| 90 | # Optional aot scheduling |
| 91 | scheduler_metadata: torch.Tensor | None = None |
| 92 | prefix_scheduler_metadata: torch.Tensor | None = None |
| 93 | mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None |
| 94 | mm_prefix_range_tensor: torch.Tensor | None = None |
| 95 | rswa_prefix_lens: torch.Tensor | None = None |
| 96 | rswa_window: int | None = None |
| 97 | |
| 98 | |
| 99 | class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]): |
| 100 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 101 | # Step-dependent fields reference persistent input buffers directly. |
| 102 | supports_draft_decode_metadata_update = True |
| 103 | |
| 104 | def __init__( |
| 105 | self, |
| 106 | kv_cache_spec: AttentionSpec, |
| 107 | layer_names: list[str], |
| 108 | vllm_config: VllmConfig, |
| 109 | device: torch.device, |
| 110 | ): |
| 111 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 112 | |
| 113 | self.block_size = kv_cache_spec.block_size |
| 114 | |
| 115 | model_config = vllm_config.model_config |
| 116 | # Compatible with models with non-uniform per-layer head counts. |
| 117 | self.num_heads_q = get_num_attention_heads_from_layers( |
| 118 | vllm_config, layer_names |
| 119 | ) or model_config.get_num_attention_heads(vllm_config.parallel_config) |
| 120 | self.num_heads_kv = model_config.get_num_kv_heads(vllm_config.parallel_config) |
| 121 | self.headdim = model_config.get_head_size() |
| 122 | |
| 123 | # Check if CUDA Graphs are enabled for decode |
| 124 | self.decode_cudagraph_enabled = ( |
| 125 | self.vllm_config.compilation_config.cudagraph_mode |
| 126 | in ( |
| 127 | CUDAGraphMode.FULL_AND_PIECEWISE, |
| 128 | CUDAGraphMode.FULL_DECODE_ONLY, |
| 129 | CUDAGraphMode.FULL, |
| 130 | ) |
| 131 | ) |
| 132 | |
| 133 | # The launch grid for the 2D kernel is defined as (num_q_blocks, num_heads_kv). |
| 134 | # A lower bound for num_q_blocks is the number of sequences. |
| 135 | # To ensure the minimum launch grid size is achieved, the number of sequences |
| 136 | # must be at least equal to the threshold below. |
| 137 | # If this threshold is not reached (i.e., the batch size is not large enough), |
| 138 | # the 3D kernel will be selected instead. |
| 139 | self.seq_threshold_3D = MIN_LAUNCH_GRID_SIZE_2D // self.num_heads_kv |
| 140 | |
| 141 | # Modify the threshold if needed. |
| 142 | if self.decode_cudagraph_enabled: |
| 143 | capture_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes |
| 144 | assert capture_sizes, "CUDA Graphs enabled but no capture sizes specified." |
| 145 | |
| 146 | # Select the CUDA Graph capture size closest to self.seq_threshold_3D |
| 147 | # as threshold. This ensures that each captured graph covers the |
| 148 | # correct execution path. |
| 149 | self.seq_threshold_3D = min( |
| 150 | capture_sizes, |
| 151 | key=lambda x: abs(x - self.seq_threshold_3D), |
| 152 | ) |
| 153 | |
| 154 | self.num_par_softmax_segments = NUM_PAR_SOFTMAX_SEGMENTS |
| 155 | # On SM120, batches whose 16-segment grid under-fills the SMs use 64. |
| 156 | self.max_seqs_64_segments = 0 |
| 157 | if current_platform.is_cuda() and current_platform.is_device_capability( |
| 158 | (12, 0) |
| 159 | ): |
| 160 | self.max_seqs_64_segments = (current_platform.num_compute_units() - 1) // ( |
| 161 | NUM_PAR_SOFTMAX_SEGMENTS * self.num_heads_kv |
| 162 | ) |
| 163 | max_num_tokens_3d = self.seq_threshold_3D |
| 164 | if self.max_seqs_64_segments > 0: |
| 165 | self.num_par_softmax_segments = 64 |
| 166 | # build() reuses this scratch at 16 segments with proportionally more rows. |
| 167 | max_num_tokens_3d = max( |
| 168 | min(self.max_seqs_64_segments, max_num_tokens_3d), |
| 169 | cdiv( |
| 170 | max_num_tokens_3d, |
| 171 | self.num_par_softmax_segments // NUM_PAR_SOFTMAX_SEGMENTS, |
| 172 | ), |
| 173 | ) |
| 174 | headdim_padded = next_power_of_2(self.headdim) |
| 175 | self.softmax_segm_output = torch.empty( |
| 176 | ( |
| 177 | max_num_tokens_3d, |
| 178 | self.num_heads_q, |
| 179 | self.num_par_softmax_segments, |
| 180 | headdim_padded, |
| 181 | ), |
| 182 | dtype=torch.float32, |
| 183 | device=device, |
| 184 | ) |
| 185 | self.softmax_segm_max = torch.empty( |
| 186 | (max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments), |
| 187 | dtype=torch.float32, |
| 188 | device=device, |
| 189 | ) |
| 190 | self.softmax_segm_expsum = torch.empty( |
| 191 | (max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments), |
| 192 | dtype=torch.float32, |
| 193 | device=device, |
| 194 | ) |
| 195 | self.rswa_window = model_config.rswa_window |
| 196 | self.persistent_rswa_prefix_lens: torch.Tensor | None = None |
| 197 | if self.rswa_window is not None: |
| 198 | self.persistent_rswa_prefix_lens = torch.empty( |
| 199 | vllm_config.scheduler_config.max_num_seqs, |
| 200 | dtype=torch.int32, |
| 201 | device=device, |
| 202 | ) |
| 203 | |
| 204 | def build_for_cudagraph_capture( |
| 205 | self, common_attn_metadata: CommonAttentionMetadata |
| 206 | ) -> TritonAttentionMetadata: |
| 207 | attn_metadata = self.build(0, common_attn_metadata) |
| 208 | # When doing full graph capture, setting seq_lens to |
| 209 | # max_model_len will cause graph capture to be extremely |
| 210 | # slow, so here we set it to 1. |
| 211 | attn_metadata.seq_lens.fill_(1) |
| 212 | return attn_metadata |
| 213 | |
| 214 | def build( |
| 215 | self, |
| 216 | common_prefix_len: int, |
| 217 | common_attn_metadata: CommonAttentionMetadata, |
| 218 | fast_build: bool = False, |
| 219 | ) -> TritonAttentionMetadata: |
| 220 | num_reqs = common_attn_metadata.num_reqs |
| 221 | num_actual_tokens = common_attn_metadata.num_actual_tokens |
| 222 | max_query_len = common_attn_metadata.max_query_len |
| 223 | |
| 224 | max_seq_len = common_attn_metadata.max_seq_len |
| 225 | query_start_loc = common_attn_metadata.query_start_loc |
| 226 | seq_lens = common_attn_metadata.seq_lens |
| 227 | block_table_tensor = common_attn_metadata.block_table_tensor |
| 228 | slot_mapping = common_attn_metadata.slot_mapping |
| 229 | |
| 230 | use_cascade = common_prefix_len > 0 |
| 231 | |
| 232 | if use_cascade: |
| 233 | cu_prefix_query_lens = torch.tensor( |
| 234 | [0, num_actual_tokens], dtype=torch.int32, device=self.device |
| 235 | ) |
| 236 | prefix_kv_lens = torch.tensor( |
| 237 | [common_prefix_len], dtype=torch.int32, device=self.device |
| 238 | ) |
| 239 | suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len |
| 240 | suffix_kv_lens = suffix_kv_lens.to(self.device) |
| 241 | else: |
| 242 | cu_prefix_query_lens = None |
| 243 | prefix_kv_lens = None |
| 244 | suffix_kv_lens = None |
| 245 | prefix_scheduler_metadata = None |
| 246 | |
| 247 | attn_metadata = TritonAttentionMetadata( |
| 248 | num_actual_tokens=num_actual_tokens, |
| 249 | max_query_len=max_query_len, |
| 250 | query_start_loc=query_start_loc, |
| 251 | max_seq_len=max_seq_len, |
| 252 | seq_lens=seq_lens, |
| 253 | block_table=block_table_tensor, |
| 254 | slot_mapping=slot_mapping, |
| 255 | causal=common_attn_metadata.causal, |
| 256 | use_cascade=use_cascade, |
| 257 | common_prefix_len=common_prefix_len, |
| 258 | cu_prefix_query_lens=cu_prefix_query_lens, |
| 259 | prefix_kv_lens=prefix_kv_lens, |
| 260 | suffix_kv_lens=suffix_kv_lens, |
| 261 | prefix_scheduler_metadata=prefix_scheduler_metadata, |
| 262 | seq_threshold_3D=self.seq_threshold_3D, |
| 263 | num_par_softmax_segments=self.num_par_softmax_segments, |
| 264 | softmax_segm_output=self.softmax_segm_output, |
| 265 | softmax_segm_max=self.softmax_segm_max, |
| 266 | softmax_segm_expsum=self.softmax_segm_expsum, |
| 267 | ) |
| 268 | if self.max_seqs_64_segments > 0 and ( |
| 269 | self.max_seqs_64_segments < seq_lens.shape[0] |
| 270 | or self.softmax_segm_max.shape[0] < num_actual_tokens |
| 271 | ): |
| 272 | # The 16-segment grid fills the SMs, or the 64-segment rows cannot |
| 273 | # hold every query token: same scratch, 16-segment layout. |
| 274 | segments = NUM_PAR_SOFTMAX_SEGMENTS |
| 275 | attn_metadata.num_par_softmax_segments = segments |
| 276 | attn_metadata.softmax_segm_output = self.softmax_segm_output.view( |
| 277 | -1, self.num_heads_q, segments, self.softmax_segm_output.shape[-1] |
| 278 | ) |
| 279 | attn_metadata.softmax_segm_max = self.softmax_segm_max.view( |
| 280 | -1, self.num_heads_q, segments |
| 281 | ) |
| 282 | attn_metadata.softmax_segm_expsum = self.softmax_segm_expsum.view( |
| 283 | -1, self.num_heads_q, segments |
| 284 | ) |
| 285 | |
| 286 | mm_ranges = common_attn_metadata.mm_req_doc_ranges |
| 287 | if mm_ranges is not None: |
| 288 | attn_metadata.mm_prefix_range = mm_ranges |
| 289 | attn_metadata.mm_prefix_range_tensor = compute_mm_prefix_range_tensor( |
| 290 | mm_ranges, num_reqs, seq_lens.device |
| 291 | ) |
| 292 | |
| 293 | rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens |
| 294 | if self.rswa_window is not None and rswa_prefix_lens is not None: |
| 295 | assert self.persistent_rswa_prefix_lens is not None |
| 296 | rswa_prefix_lens = rswa_prefix_lens.to( |
| 297 | device=self.device, dtype=torch.int32, non_blocking=True |
| 298 | ) |
| 299 | persistent_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs] |
| 300 | persistent_prefix_lens.copy_(rswa_prefix_lens[:num_reqs]) |
| 301 | attn_metadata.rswa_prefix_lens = persistent_prefix_lens |
| 302 | attn_metadata.rswa_window = self.rswa_window |
| 303 | |
| 304 | return attn_metadata |
| 305 | |
| 306 | def update_draft_decode_metadata(self, _metadata: TritonAttentionMetadata) -> None: |
| 307 | pass |
| 308 | |
| 309 | |
| 310 | class TritonAttentionBackend(AttentionBackend): |
| 311 | @classmethod |
| 312 | def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": |
| 313 | """Per-token-head modes pack inline fp32 scales after each head's |
| 314 | data, so the content is (data + one scale) per K/V side.""" |
| 315 | mode = spec.kv_quant_mode |
| 316 | if spec.state_content_bytes is not None or not mode.is_per_token_head: |
| 317 | return spec |
| 318 | hs_k, hs_v = spec.head_size, spec.head_size_v |
| 319 | if mode == KVQuantMode.INT4_PER_TOKEN_HEAD: |
| 320 | hs_k, hs_v = hs_k // 2, hs_v // 2 |
| 321 | scale_bytes = get_dtype_size(torch.float32) |
| 322 | content = (hs_k + hs_v) * get_dtype_size(spec.dtype) + 2 * scale_bytes |
| 323 | return replace(spec, state_content_bytes=content) |
| 324 | |
| 325 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 326 | torch.float16, |
| 327 | torch.bfloat16, |
| 328 | torch.float32, |
| 329 | ] |
| 330 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 331 | "auto", |
| 332 | "float16", |
| 333 | "bfloat16", |
| 334 | "fp8", |
| 335 | "fp8_e4m3", |
| 336 | "fp8_e5m2", |
| 337 | "int4_per_token_head", |
| 338 | "int8_per_token_head", |
| 339 | "fp8_per_token_head", |
| 340 | ] |
| 341 | |
| 342 | @staticmethod |
| 343 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 344 | return [MultipleOf(16)] |
| 345 | |
| 346 | @classmethod |
| 347 | def supports_block_size(cls, block_size: int | None) -> bool: |
| 348 | if block_size is None: |
| 349 | return True |
| 350 | return block_size % 16 == 0 |
| 351 | |
| 352 | forward_includes_kv_cache_update: bool = False |
| 353 | |
| 354 | @classmethod |
| 355 | def supports_non_causal(cls) -> bool: |
| 356 | return True |
| 357 | |
| 358 | @staticmethod |
| 359 | def get_name() -> str: |
| 360 | return "TRITON_ATTN" |
| 361 | |
| 362 | @classmethod |
| 363 | def supports_rswa(cls) -> bool: |
| 364 | return True |
| 365 | |
| 366 | @classmethod |
| 367 | def supports_sliding_window(cls) -> bool: |
| 368 | return True |
| 369 | |
| 370 | @classmethod |
| 371 | def supports_batch_invariance(cls) -> bool: |
| 372 | return True |
| 373 | |
| 374 | @staticmethod |
| 375 | def get_impl_cls() -> type["TritonAttentionImpl"]: |
| 376 | return TritonAttentionImpl |
| 377 | |
| 378 | @staticmethod |
| 379 | def use_cascade_attention(*args, **kwargs) -> bool: |
| 380 | return False |
| 381 | |
| 382 | @staticmethod |
| 383 | def get_builder_cls() -> type["TritonAttentionMetadataBuilder"]: |
| 384 | return TritonAttentionMetadataBuilder |
| 385 | |
| 386 | @classmethod |
| 387 | def supports_head_size(cls, head_size: int) -> bool: |
| 388 | return head_size >= 32 |
| 389 | |
| 390 | @classmethod |
| 391 | def supports_mm_prefix(cls) -> bool: |
| 392 | return True |
| 393 | |
| 394 | @classmethod |
| 395 | def supports_sink(cls) -> bool: |
| 396 | return True |
| 397 | |
| 398 | @classmethod |
| 399 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 400 | """TritonAttention supports all attention types.""" |
| 401 | return attn_type in ( |
| 402 | AttentionType.DECODER, |
| 403 | AttentionType.ENCODER, |
| 404 | AttentionType.ENCODER_ONLY, |
| 405 | AttentionType.ENCODER_DECODER, |
| 406 | ) |
| 407 | |
| 408 | @classmethod |
| 409 | def supports_alibi_sqrt(cls) -> bool: |
| 410 | return True |
| 411 | |
| 412 | @classmethod |
| 413 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 414 | return True |
| 415 | |
| 416 | |
| 417 | class TritonAttentionImpl(AttentionImpl): |
| 418 | # Per-token-head quant: scale views carved from inline head padding. |
| 419 | _k_scale_cache: torch.Tensor | None = None |
| 420 | _v_scale_cache: torch.Tensor | None = None |
| 421 | |
| 422 | def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None: |
| 423 | """Extract per-head scale views from the padded content dimension. |
| 424 | |
| 425 | The KV cache is packed as logical shape |
| 426 | ``(num_blocks, nkv, block_size, 2 * (hs + pad))`` where |
| 427 | ``pad = sizeof(float32) / sizeof(cache_dtype)``. The content dim holds |
| 428 | ``[K(hs) | K_scale(pad) | V(hs) | V_scale(pad)]`` per (head, slot); the |
| 429 | last ``pad`` elements of each half hold one float32 scale. We create |
| 430 | strided float32 views over those bytes. ``kv_cache`` must be the |
| 431 | packed logical tensor (call before any transpose), but may have HND or |
| 432 | NHD physical strides. |
| 433 | |
| 434 | Scale shape: ``(num_blocks, block_size, num_kv_heads)`` |
| 435 | """ |
| 436 | if self._k_scale_cache is not None: |
| 437 | return |
| 438 | from vllm.utils.torch_utils import get_dtype_size |
| 439 | |
| 440 | num_blocks, nkv, block_size, content = kv_cache.shape |
| 441 | dtype_sz = kv_cache.element_size() |
| 442 | scale_pad = get_dtype_size(torch.float32) // dtype_sz # e.g. 4 |
| 443 | padded_hs = content // 2 |
| 444 | hs = padded_hs - scale_pad |
| 445 | |
| 446 | raw = kv_cache.untyped_storage() |
| 447 | base_f32 = torch.tensor([], dtype=torch.float32, device=kv_cache.device).set_( |
| 448 | raw |
| 449 | ) |
| 450 | |
| 451 | def to_f32_units(elements: int) -> int: |
| 452 | nbytes = elements * dtype_sz |
| 453 | assert nbytes % 4 == 0 |
| 454 | return nbytes // 4 |
| 455 | |
| 456 | # Actual strides (in float32 units) from the tensor. The logical cache |
| 457 | # may be physically NHD, so do not assume C-contiguous HND layout. |
| 458 | strides = kv_cache.stride() |
| 459 | block_f32 = to_f32_units(strides[0]) |
| 460 | head_f32 = to_f32_units(strides[1]) |
| 461 | slot_f32 = to_f32_units(strides[2]) |
| 462 | # Scale sits at byte offset hs within each (K, then V) content half. |
| 463 | base_off_f32 = to_f32_units(kv_cache.storage_offset()) |
| 464 | k_scale_off_f32 = base_off_f32 + to_f32_units(hs) |
| 465 | v_scale_off_f32 = base_off_f32 + to_f32_units(padded_hs + hs) |
| 466 | |
| 467 | # K scales (first content half) |
| 468 | self._k_scale_cache = torch.as_strided( |
| 469 | base_f32, |
| 470 | size=(num_blocks, block_size, nkv), |
| 471 | stride=(block_f32, slot_f32, head_f32), |
| 472 | storage_offset=k_scale_off_f32, |
| 473 | ) |
| 474 | self._k_scale_cache.fill_(1.0) |
| 475 | |
| 476 | # V scales (second content half) |
| 477 | self._v_scale_cache = torch.as_strided( |
| 478 | base_f32, |
| 479 | size=(num_blocks, block_size, nkv), |
| 480 | stride=(block_f32, slot_f32, head_f32), |
| 481 | storage_offset=v_scale_off_f32, |
| 482 | ) |
| 483 | self._v_scale_cache.fill_(1.0) |
| 484 | |
| 485 | def fused_output_quant_supported(self, quant_key: QuantKey): |
| 486 | return quant_key == kFp8StaticTensorSym |
| 487 | |
| 488 | def __init__( |
| 489 | self, |
| 490 | num_heads: int, |
| 491 | head_size: int, |
| 492 | scale: float, |
| 493 | num_kv_heads: int, |
| 494 | alibi_slopes: list[float] | None, |
| 495 | sliding_window: int | None, |
| 496 | kv_cache_dtype: str, |
| 497 | logits_soft_cap: float | None = None, |
| 498 | attn_type: AttentionType = AttentionType.DECODER, |
| 499 | kv_sharing_target_layer_name: int | None = None, |
| 500 | sinks: torch.Tensor | None = None, |
| 501 | use_alibi_sqrt: bool = False, |
| 502 | chunk_lookback: int = -1, |
| 503 | ) -> None: |
| 504 | self.num_heads = num_heads |
| 505 | self.head_size = head_size |
| 506 | self.scale = float(scale) |
| 507 | self.num_kv_heads = num_kv_heads |
| 508 | if alibi_slopes is not None: |
| 509 | alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) |
| 510 | self.alibi_slopes = alibi_slopes |
| 511 | if sliding_window is None: |
| 512 | self.sliding_window = (-1, -1) |
| 513 | elif attn_type in (AttentionType.ENCODER, AttentionType.ENCODER_ONLY): |
| 514 | self.sliding_window = (sliding_window - 1, sliding_window - 1) |
| 515 | else: |
| 516 | self.sliding_window = (sliding_window - 1, 0) |
| 517 | self.kv_cache_dtype = kv_cache_dtype |
| 518 | if current_platform.is_cuda(): |
| 519 | cap = current_platform.get_device_capability() |
| 520 | cap_str = cap.as_version_str() if cap is not None else "unknown" |
| 521 | dev = current_platform.get_device_name() |
| 522 | if self.kv_cache_dtype.startswith("fp8") and not ( |
| 523 | current_platform.has_device_capability(89) |
| 524 | ): |
| 525 | suggested = ( |
| 526 | "float16" if (cap is None or cap.to_int() < 80) else "bfloat16" |
| 527 | ) |
| 528 | raise ValueError( |
| 529 | f"FP8 KV cache is not supported by the Triton attention backend " |
| 530 | f"on {dev} (compute capability {cap_str}); native FP8 (fp8e4nv) " |
| 531 | f"requires SM89+. Re-run with --kv-cache-dtype {suggested}." |
| 532 | ) |
| 533 | if self.kv_cache_dtype == "bfloat16" and not ( |
| 534 | current_platform.has_device_capability(80) |
| 535 | ): |
| 536 | raise ValueError( |
| 537 | f"bfloat16 KV cache is not supported on {dev} (compute capability " |
| 538 | f"{cap_str}); bfloat16 requires SM80+. Re-run with " |
| 539 | f"--kv-cache-dtype float16." |
| 540 | ) |
| 541 | if logits_soft_cap is None: |
| 542 | # In flash-attn, setting logits_soft_cap as 0 means no soft cap. |
| 543 | logits_soft_cap = 0 |
| 544 | self.logits_soft_cap = logits_soft_cap |
| 545 | self.kv_sharing_target_layer_name = kv_sharing_target_layer_name |
| 546 | |
| 547 | self.num_queries_per_kv = self.num_heads // self.num_kv_heads |
| 548 | |
| 549 | self.attn_type = attn_type |
| 550 | self.fp8_dtype = current_platform.fp8_dtype() |
| 551 | |
| 552 | self.sinks = sinks |
| 553 | if sinks is not None: |
| 554 | assert sinks.shape[0] == num_heads, ( |
| 555 | "Sinks must have the same number of heads as the number of " |
| 556 | f"heads in the layer. Sinks shape: {sinks.shape}, " |
| 557 | f"num_heads: {num_heads}." |
| 558 | ) |
| 559 | self.use_alibi_sqrt = use_alibi_sqrt |
| 560 | self.chunk_lookback = chunk_lookback |
| 561 | self.supports_quant_query_input = current_platform.is_cuda() |
| 562 | |
| 563 | self._kv_quant_mode = get_kv_quant_mode(kv_cache_dtype) |
| 564 | self._is_per_token_head_quant = self._kv_quant_mode.is_per_token_head |
| 565 | |
| 566 | # Enable tensor descriptors for Q/K/V load/store on platforms that |
| 567 | # benefit from HW 2D block reads (Intel XPU). The dead branch |
| 568 | # is eliminated at Triton compile time, so other platforms see |
| 569 | # zero cost when TD is off. |
| 570 | # |
| 571 | # ``VLLM_TRITON_USE_TD`` is tri-state: |
| 572 | # - unset (None): auto-select (TD on for XPU, off elsewhere), |
| 573 | # - ``1``: force TD on regardless of platform, |
| 574 | # - ``0``: force TD off regardless of platform (useful for A/B). |
| 575 | td_override = envs.VLLM_TRITON_USE_TD |
| 576 | if td_override is None: |
| 577 | self.use_td = current_platform.is_xpu() |
| 578 | else: |
| 579 | self.use_td = td_override |
| 580 | |
| 581 | def forward( |
| 582 | self, |
| 583 | layer: torch.nn.Module, |
| 584 | query: torch.Tensor, |
| 585 | key: torch.Tensor, |
| 586 | value: torch.Tensor, |
| 587 | kv_cache: torch.Tensor, |
| 588 | attn_metadata: TritonAttentionMetadata, |
| 589 | output: torch.Tensor, |
| 590 | output_scale: torch.Tensor | None = None, |
| 591 | output_block_scale: torch.Tensor | None = None, |
| 592 | ) -> torch.Tensor: |
| 593 | """Forward pass with Paged Attention impl. in Triton. |
| 594 | |
| 595 | Args: |
| 596 | layer: The attention layer, providing the q/k/v quantization scales. |
| 597 | query: shape = [num_tokens, num_heads, head_size] |
| 598 | key: shape = [num_tokens, num_kv_heads, head_size] |
| 599 | value: shape = [num_tokens, num_kv_heads, head_size] |
| 600 | kv_cache: shape = |
| 601 | [num_blocks, num_kv_heads, block_size, 2 * head_size] |
| 602 | attn_metadata: Metadata for attention. |
| 603 | output: Tensor that the attention result is written into. |
| 604 | output_scale: Scale for fused output quantization. |
| 605 | output_block_scale: Block scale for fused output quantization; |
| 606 | not supported by this backend. |
| 607 | |
| 608 | Returns: |
| 609 | shape = [num_tokens, num_heads * head_size] |
| 610 | |
| 611 | """ |
| 612 | if output_block_scale is not None: |
| 613 | raise NotImplementedError( |
| 614 | "fused block_scale output quantization is not yet supported" |
| 615 | " for TritonAttentionImpl" |
| 616 | ) |
| 617 | |
| 618 | if attn_metadata is None: |
| 619 | # Profiling run. |
| 620 | return output.fill_(0) |
| 621 | |
| 622 | assert attn_metadata.use_cascade is False |
| 623 | |
| 624 | # IMPORTANT! |
| 625 | # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in |
| 626 | # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead |
| 627 | # in this method. For example, `view` and `slice` (or `[:n]`) operations |
| 628 | # are surprisingly slow even in the case they do not invoke any GPU ops. |
| 629 | # Minimize the PyTorch ops in this method as much as possible. |
| 630 | # Whenever making a change in this method, please benchmark the |
| 631 | # performance to make sure it does not introduce any overhead. |
| 632 | |
| 633 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 634 | |
| 635 | # Handle encoder attention differently - no KV cache needed |
| 636 | if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 637 | # For encoder attention, |
| 638 | # we use direct Q, K, V tensors without caching |
| 639 | return self._forward_encoder_attention( |
| 640 | query[:num_actual_tokens], |
| 641 | key[:num_actual_tokens], |
| 642 | value[:num_actual_tokens], |
| 643 | output[:num_actual_tokens], |
| 644 | attn_metadata, |
| 645 | layer, |
| 646 | ) |
| 647 | |
| 648 | # KV cache arrives in logical (B, H, N, 2*hs) order. |
| 649 | # Per-token-head quantized KV cache: handled by the core unified |
| 650 | # kernel, which dequantizes per-(token, head) inline via constexpr |
| 651 | # branches (INT8 / FP8) and dispatches to the packed INT4 kernel. |
| 652 | if self._is_per_token_head_quant: |
| 653 | key_cache, value_cache = self._pth_key_value_caches(kv_cache) |
| 654 | k_scale_cache = self._k_scale_cache |
| 655 | v_scale_cache = self._v_scale_cache |
| 656 | q_descale = k_descale = v_descale = None |
| 657 | # FP8 per-tensor / auto path (original flow). |
| 658 | else: |
| 659 | kv_cache = kv_cache.transpose(1, 2) |
| 660 | hs = self.head_size |
| 661 | key_cache, value_cache = kv_cache.split(hs, dim=-1) |
| 662 | if ( |
| 663 | is_quantized_kv_cache(self.kv_cache_dtype) |
| 664 | and key_cache.dtype != self.fp8_dtype |
| 665 | ): |
| 666 | key_cache = key_cache.view(self.fp8_dtype) |
| 667 | value_cache = value_cache.view(self.fp8_dtype) |
| 668 | descale_shape = ( |
| 669 | attn_metadata.query_start_loc.shape[0] - 1, |
| 670 | key_cache.shape[2], |
| 671 | ) |
| 672 | q_descale = ( |
| 673 | layer._q_scale |
| 674 | if ( |
| 675 | self._kv_quant_mode == KVQuantMode.FP8_PER_TENSOR |
| 676 | and query.dtype == self.fp8_dtype |
| 677 | ) |
| 678 | else None |
| 679 | ) |
| 680 | k_descale = layer._k_scale.expand(descale_shape) |
| 681 | v_descale = layer._v_scale.expand(descale_shape) |
| 682 | k_scale_cache = None |
| 683 | v_scale_cache = None |
| 684 | |
| 685 | cu_seqlens_q = attn_metadata.query_start_loc |
| 686 | seqused_k = attn_metadata.seq_lens |
| 687 | max_seqlen_q = attn_metadata.max_query_len |
| 688 | max_seqlen_k = attn_metadata.max_seq_len |
| 689 | block_table = attn_metadata.block_table |
| 690 | |
| 691 | seq_threshold_3D = attn_metadata.seq_threshold_3D |
| 692 | num_par_softmax_segments = attn_metadata.num_par_softmax_segments |
| 693 | softmax_segm_output = attn_metadata.softmax_segm_output |
| 694 | softmax_segm_max = attn_metadata.softmax_segm_max |
| 695 | softmax_segm_expsum = attn_metadata.softmax_segm_expsum |
| 696 | |
| 697 | mm_prefix_range_tensor = attn_metadata.mm_prefix_range_tensor |
| 698 | |
| 699 | unified_attention( |
| 700 | q=query[:num_actual_tokens], |
| 701 | k=key_cache, |
| 702 | v=value_cache, |
| 703 | out=output[:num_actual_tokens], |
| 704 | cu_seqlens_q=cu_seqlens_q, |
| 705 | max_seqlen_q=max_seqlen_q, |
| 706 | seqused_k=seqused_k, |
| 707 | max_seqlen_k=max_seqlen_k, |
| 708 | softmax_scale=self.scale, |
| 709 | causal=attn_metadata.causal, |
| 710 | alibi_slopes=self.alibi_slopes, |
| 711 | use_alibi_sqrt=self.use_alibi_sqrt, |
| 712 | window_size=self.sliding_window, |
| 713 | block_table=block_table, |
| 714 | softcap=self.logits_soft_cap, |
| 715 | q_descale=q_descale, |
| 716 | k_descale=k_descale, |
| 717 | v_descale=v_descale, |
| 718 | seq_threshold_3D=seq_threshold_3D, |
| 719 | num_par_softmax_segments=num_par_softmax_segments, |
| 720 | softmax_segm_output=softmax_segm_output, |
| 721 | softmax_segm_max=softmax_segm_max, |
| 722 | softmax_segm_expsum=softmax_segm_expsum, |
| 723 | sinks=self.sinks, |
| 724 | output_scale=output_scale, |
| 725 | mm_prefix_range=mm_prefix_range_tensor, |
| 726 | rswa_prefix_lens=attn_metadata.rswa_prefix_lens, |
| 727 | rswa_window=attn_metadata.rswa_window, |
| 728 | kv_quant_mode=self._kv_quant_mode, |
| 729 | k_scale_cache=k_scale_cache, |
| 730 | v_scale_cache=v_scale_cache, |
| 731 | chunk_lookback=self.chunk_lookback, |
| 732 | use_td=self.use_td, |
| 733 | mm_prefix_clamp_sliding_window=getattr( |
| 734 | layer, "mm_prefix_clamp_sliding_window", False |
| 735 | ), |
| 736 | ) |
| 737 | |
| 738 | return output |
| 739 | |
| 740 | def _pth_key_value_caches( |
| 741 | self, kv_cache: torch.Tensor |
| 742 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 743 | """Per-token-head K/V cache views (ensures scale caches; FP8 retyped).""" |
| 744 | self._ensure_scale_caches(kv_cache) |
| 745 | padded_hs = kv_cache.shape[-1] // 2 |
| 746 | key_cache, value_cache = kv_cache.transpose(1, 2).split(padded_hs, dim=-1) |
| 747 | if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD: |
| 748 | key_cache = key_cache.view(self.fp8_dtype) |
| 749 | value_cache = value_cache.view(self.fp8_dtype) |
| 750 | return key_cache, value_cache |
| 751 | |
| 752 | def _forward_encoder_attention( |
| 753 | self, |
| 754 | query: torch.Tensor, |
| 755 | key: torch.Tensor, |
| 756 | value: torch.Tensor, |
| 757 | output: torch.Tensor, |
| 758 | attn_metadata: TritonAttentionMetadata, |
| 759 | layer: torch.nn.Module, |
| 760 | ) -> torch.Tensor: |
| 761 | """Forward pass for encoder attention without KV cache. |
| 762 | |
| 763 | Args: |
| 764 | query: shape = [num_encoder_tokens, num_heads, head_size] |
| 765 | key: shape = [num_encoder_tokens, num_kv_heads, head_size] |
| 766 | value: shape = [num_encoder_tokens, num_kv_heads, head_size] |
| 767 | output: shape = [num_encoder_tokens, num_heads, head_size] |
| 768 | attn_metadata: Encoder attention metadata |
| 769 | layer: The attention layer |
| 770 | |
| 771 | """ |
| 772 | # Quantized KV cache is not supported for encoder attention. |
| 773 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 774 | raise NotImplementedError( |
| 775 | "quantized KV cache is not supported for encoder attention" |
| 776 | ) |
| 777 | |
| 778 | # Use encoder-specific metadata for sequence information |
| 779 | query_start_loc = attn_metadata.query_start_loc |
| 780 | seq_lens = attn_metadata.seq_lens |
| 781 | max_query_len = attn_metadata.max_query_len |
| 782 | |
| 783 | # Call flash attention directly on Q, K, V tensors |
| 784 | context_attention_fwd( |
| 785 | q=query, |
| 786 | k=key, |
| 787 | v=value, |
| 788 | o=output, |
| 789 | b_start_loc=query_start_loc, |
| 790 | b_seq_len=seq_lens, |
| 791 | max_input_len=max_query_len, |
| 792 | is_causal=False, # Encoder attention is bidirectional |
| 793 | softmax_scale=self.scale, |
| 794 | sliding_window_q=self.sliding_window[0], |
| 795 | sliding_window_k=self.sliding_window[1], |
| 796 | sinks=self.sinks, |
| 797 | ) |
| 798 | return output |
| 799 | |
| 800 | def do_kv_cache_update( |
| 801 | self, |
| 802 | layer: AttentionLayer, |
| 803 | key: torch.Tensor, |
| 804 | value: torch.Tensor, |
| 805 | kv_cache: torch.Tensor, |
| 806 | slot_mapping: torch.Tensor, |
| 807 | ): |
| 808 | if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 809 | # For encoder attention, |
| 810 | # we use direct Q, K, V tensors without caching |
| 811 | return |
| 812 | # Reshape the input keys and values and store them in the cache. |
| 813 | if self._is_per_token_head_quant: |
| 814 | key_cache, value_cache = self._pth_key_value_caches(kv_cache) |
| 815 | k_scale_cache = self._k_scale_cache |
| 816 | v_scale_cache = self._v_scale_cache |
| 817 | triton_reshape_and_cache_flash_per_token_head_quant( |
| 818 | key, |
| 819 | value, |
| 820 | key_cache, |
| 821 | value_cache, |
| 822 | k_scale_cache, |
| 823 | v_scale_cache, |
| 824 | slot_mapping, |
| 825 | kv_quant_mode=self._kv_quant_mode, |
| 826 | ) |
| 827 | return |
| 828 | # For decoder and cross-attention, use KV cache as before. |
| 829 | # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) |
| 830 | key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) |
| 831 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 832 | key_cache = key_cache.view(self.fp8_dtype) |
| 833 | value_cache = value_cache.view(self.fp8_dtype) |
| 834 | triton_reshape_and_cache_flash( |
| 835 | key, |
| 836 | value, |
| 837 | key_cache, |
| 838 | value_cache, |
| 839 | slot_mapping, |
| 840 | self.kv_cache_dtype, |
| 841 | layer._k_scale, |
| 842 | layer._v_scale, |
| 843 | ) |
| 844 | |
| 845 | def fused_rope_kvcache_supported(self): |
| 846 | if self._is_per_token_head_quant: |
| 847 | return False |
| 848 | return rocm_aiter_ops.is_enabled() |
| 849 | |
| 850 | def do_rope_and_kv_cache_update( |
| 851 | self, |
| 852 | layer: AttentionLayer, |
| 853 | query: torch.Tensor, |
| 854 | key: torch.Tensor, |
| 855 | value: torch.Tensor, |
| 856 | positions: torch.Tensor, |
| 857 | cos_sin_cache: torch.Tensor, |
| 858 | is_neox: bool, |
| 859 | kv_cache: torch.Tensor, |
| 860 | layer_slot_mapping: torch.Tensor, |
| 861 | ): |
| 862 | # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) |
| 863 | key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) |
| 864 | flash_layout = True |
| 865 | |
| 866 | is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) |
| 867 | if is_fp8_kv_cache: |
| 868 | key_cache = key_cache.view(self.fp8_dtype) |
| 869 | value_cache = value_cache.view(self.fp8_dtype) |
| 870 | |
| 871 | rocm_aiter_ops.triton_rope_and_cache( |
| 872 | query, |
| 873 | key, |
| 874 | value, |
| 875 | positions, |
| 876 | cos_sin_cache, |
| 877 | is_neox, |
| 878 | key_cache, |
| 879 | value_cache, |
| 880 | layer_slot_mapping, |
| 881 | layer._k_scale, |
| 882 | layer._v_scale, |
| 883 | flash_layout, |
| 884 | is_fp8_kv_cache, |
| 885 | ) |