| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from dataclasses import dataclass, replace |
| 6 | from enum import Enum |
| 7 | from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar |
| 8 | |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | |
| 12 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 13 | kFp8Dynamic64Sym, |
| 14 | kFp8Dynamic128Sym, |
| 15 | kFp8StaticTensorSym, |
| 16 | kNvfp4Dynamic, |
| 17 | ) |
| 18 | |
| 19 | if TYPE_CHECKING: |
| 20 | from vllm.config import VllmConfig |
| 21 | from vllm.config.cache import CacheDType |
| 22 | from vllm.model_executor.layers.linear import ColumnParallelLinear |
| 23 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 24 | from vllm.platforms.interface import DeviceCapability |
| 25 | from vllm.v1.hisparse.runtime import HiSparseCacheHandle |
| 26 | from vllm.v1.kv_cache_interface import ( |
| 27 | AttentionSpec, |
| 28 | KVCacheLayout, |
| 29 | KVCacheSpec, |
| 30 | KVQuantMode, |
| 31 | ) |
| 32 | |
| 33 | from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode |
| 34 | |
| 35 | |
| 36 | class AttentionType(str, Enum): |
| 37 | """Attention type. |
| 38 | Use string to be compatible with `torch.compile`. |
| 39 | """ |
| 40 | |
| 41 | DECODER = "decoder" |
| 42 | """Decoder attention between previous layer Q/K/V.""" |
| 43 | ENCODER = "encoder" |
| 44 | """Encoder attention between previous layer Q/K/V for encoder-decoder.""" |
| 45 | ENCODER_ONLY = "encoder_only" |
| 46 | """Encoder attention between previous layer Q/K/V.""" |
| 47 | ENCODER_DECODER = "encoder_decoder" |
| 48 | """Attention between dec. Q and enc. K/V for encoder-decoder.""" |
| 49 | |
| 50 | |
| 51 | class MultipleOf: |
| 52 | base: int |
| 53 | |
| 54 | def __init__(self, base: int): |
| 55 | self.base = base |
| 56 | |
| 57 | def __repr__(self) -> str: |
| 58 | return f"MultipleOf({self.base})" |
| 59 | |
| 60 | |
| 61 | class AttentionBackend(ABC): |
| 62 | """Abstract class for attention backends.""" |
| 63 | |
| 64 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] |
| 65 | supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [ |
| 66 | "auto", |
| 67 | "float16", |
| 68 | "bfloat16", |
| 69 | ] |
| 70 | |
| 71 | # Does attention's forward() include kv cache update? |
| 72 | forward_includes_kv_cache_update: bool = True |
| 73 | |
| 74 | @staticmethod |
| 75 | def get_supported_kernel_block_sizes( |
| 76 | kv_cache_spec: "KVCacheSpec | None" = None, |
| 77 | ) -> list[int | MultipleOf]: |
| 78 | return [MultipleOf(1)] |
| 79 | |
| 80 | @staticmethod |
| 81 | @abstractmethod |
| 82 | def get_name() -> str: |
| 83 | raise NotImplementedError |
| 84 | |
| 85 | @staticmethod |
| 86 | @abstractmethod |
| 87 | def get_impl_cls() -> type["AttentionImplBase"]: |
| 88 | raise NotImplementedError |
| 89 | |
| 90 | @staticmethod |
| 91 | @abstractmethod |
| 92 | def get_builder_cls(): # -> Type["AttentionMetadataBuilder"]: |
| 93 | raise NotImplementedError |
| 94 | |
| 95 | @classmethod |
| 96 | def full_cls_name(cls) -> tuple[str, str]: |
| 97 | return (cls.__module__, cls.__qualname__) |
| 98 | |
| 99 | @classmethod |
| 100 | def get_supported_head_sizes(cls) -> list[int]: |
| 101 | return [] |
| 102 | |
| 103 | @classmethod |
| 104 | def supports_head_size(cls, head_size: int) -> bool: |
| 105 | supported_head_sizes = cls.get_supported_head_sizes() |
| 106 | return (not supported_head_sizes) or head_size in supported_head_sizes |
| 107 | |
| 108 | @classmethod |
| 109 | def supports_dtype(cls, dtype: torch.dtype) -> bool: |
| 110 | return dtype in cls.supported_dtypes |
| 111 | |
| 112 | @classmethod |
| 113 | def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool: |
| 114 | if kv_cache_dtype is None: |
| 115 | return True |
| 116 | return (not cls.supported_kv_cache_dtypes) or ( |
| 117 | kv_cache_dtype in cls.supported_kv_cache_dtypes |
| 118 | ) |
| 119 | |
| 120 | @classmethod |
| 121 | def supports_block_size(cls, block_size: int | None) -> bool: |
| 122 | if block_size is None: |
| 123 | return True |
| 124 | |
| 125 | supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes() |
| 126 | if not supported_kernel_block_sizes: |
| 127 | return True |
| 128 | |
| 129 | for supported_size in supported_kernel_block_sizes: |
| 130 | if isinstance(supported_size, MultipleOf): |
| 131 | supported_size = supported_size.base |
| 132 | # With hybrid_blocks feature, the framework-level block size |
| 133 | # only needs to be a multiple of the kernel's requirement, |
| 134 | # even if the kernel requires a fixed block_size. |
| 135 | if block_size % supported_size == 0: |
| 136 | return True |
| 137 | return False |
| 138 | |
| 139 | @classmethod |
| 140 | def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": |
| 141 | """Adjust the layer's KV cache spec for this backend's kernels. Used when the |
| 142 | kernels want KV packed in a specific way. |
| 143 | |
| 144 | NOTE: temporary compatibility API. Today the Attention layer builds the spec |
| 145 | from the model config and the backend only gets to adjust it post-hoc; the end |
| 146 | state is for the backend to build and return the spec directly, at which point |
| 147 | this hook goes away. |
| 148 | |
| 149 | (see: https://github.com/vllm-project/vllm/issues/42449) |
| 150 | """ |
| 151 | return spec |
| 152 | |
| 153 | @classmethod |
| 154 | def get_preferred_block_size(cls, default_block_size: int) -> int: |
| 155 | supported_sizes = cls.get_supported_kernel_block_sizes() |
| 156 | if not supported_sizes: |
| 157 | return default_block_size |
| 158 | |
| 159 | if cls.supports_block_size(default_block_size): |
| 160 | return default_block_size |
| 161 | |
| 162 | return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes) |
| 163 | |
| 164 | @classmethod |
| 165 | def is_mla(cls) -> bool: |
| 166 | return False |
| 167 | |
| 168 | @classmethod |
| 169 | def supports_sink(cls) -> bool: |
| 170 | return False |
| 171 | |
| 172 | @classmethod |
| 173 | def supports_alibi_sqrt(cls) -> bool: |
| 174 | return False |
| 175 | |
| 176 | @classmethod |
| 177 | def supports_mm_prefix(cls) -> bool: |
| 178 | return False |
| 179 | |
| 180 | @classmethod |
| 181 | def is_sparse(cls) -> bool: |
| 182 | return False |
| 183 | |
| 184 | @classmethod |
| 185 | def supports_per_head_quant_scales(cls) -> bool: |
| 186 | return False |
| 187 | |
| 188 | @classmethod |
| 189 | def supports_sliding_window(cls) -> bool: |
| 190 | return False |
| 191 | |
| 192 | @classmethod |
| 193 | def supports_rswa(cls) -> bool: |
| 194 | return False |
| 195 | |
| 196 | @classmethod |
| 197 | def supports_non_causal(cls) -> bool: |
| 198 | """Check if backend supports non-causal (bidirectional) attention |
| 199 | for decoder models. |
| 200 | |
| 201 | Unlike ENCODER_ONLY attention type which implies a different |
| 202 | execution model, this refers to non-causal attention within the |
| 203 | standard paged-KV-cache decoder path. |
| 204 | """ |
| 205 | return False |
| 206 | |
| 207 | @classmethod |
| 208 | def supports_batch_invariance(cls) -> bool: |
| 209 | return False |
| 210 | |
| 211 | @classmethod |
| 212 | def supports_kv_connector(cls) -> bool: |
| 213 | return True |
| 214 | |
| 215 | @classmethod |
| 216 | def supports_device_cpu_query_lens_mismatch(cls) -> bool: |
| 217 | """Whether this backend can run a batch whose device query_start_loc disagrees |
| 218 | with the CPU one; backends that plan off the CPU query lengths must opt out. |
| 219 | |
| 220 | Currently only verification requests are affected: adaptive verification trims |
| 221 | their drafts on device. On the CPU the draft budget is evenly distributed across |
| 222 | requests, so the total draft budget, the decode/prefill split point and the CPU |
| 223 | prefill query lengths all stay correct. |
| 224 | |
| 225 | SSM backends opt out: their recurrent-state planning is built from the CPU |
| 226 | per-request boundaries, which the trimmed batch no longer matches. |
| 227 | """ |
| 228 | return not cls.is_ssm() |
| 229 | |
| 230 | @classmethod |
| 231 | def supports_pcp(cls) -> bool: |
| 232 | try: |
| 233 | return cls.get_impl_cls().supports_pcp |
| 234 | except NotImplementedError: |
| 235 | return False |
| 236 | |
| 237 | @classmethod |
| 238 | def supports_dcp(cls) -> bool: |
| 239 | try: |
| 240 | return cls.get_impl_cls().supports_dcp |
| 241 | except NotImplementedError: |
| 242 | return False |
| 243 | |
| 244 | @classmethod |
| 245 | def supports_non_causal_dcp(cls) -> bool: |
| 246 | builder_cls = cls.get_builder_cls() |
| 247 | return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False)) |
| 248 | |
| 249 | @classmethod |
| 250 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 251 | """Check if backend supports a given attention type. |
| 252 | |
| 253 | By default, only supports decoder attention. |
| 254 | Backends should override this to support other attention types. |
| 255 | """ |
| 256 | return attn_type == AttentionType.DECODER |
| 257 | |
| 258 | @classmethod |
| 259 | def supports_compute_capability(cls, capability: "DeviceCapability") -> bool: |
| 260 | return True |
| 261 | |
| 262 | @classmethod |
| 263 | def supports_combination( |
| 264 | cls, |
| 265 | head_size: int, |
| 266 | dtype: torch.dtype, |
| 267 | kv_cache_dtype: "CacheDType | None", |
| 268 | block_size: int | None, |
| 269 | use_mla: bool, |
| 270 | has_sink: bool, |
| 271 | use_sparse: bool, |
| 272 | use_mm_prefix: bool, |
| 273 | device_capability: "DeviceCapability", |
| 274 | ) -> str | None: |
| 275 | return None |
| 276 | |
| 277 | @classmethod |
| 278 | def validate_configuration( |
| 279 | cls, |
| 280 | head_size: int, |
| 281 | dtype: torch.dtype, |
| 282 | kv_cache_dtype: "CacheDType | None", |
| 283 | block_size: int | None, |
| 284 | use_mla: bool, |
| 285 | has_sink: bool, |
| 286 | use_sparse: bool, |
| 287 | use_mm_prefix: bool, |
| 288 | use_per_head_quant_scales: bool, |
| 289 | device_capability: "DeviceCapability", |
| 290 | attn_type: str, |
| 291 | has_sliding_window: bool = False, |
| 292 | use_non_causal: bool = False, |
| 293 | use_batch_invariant: bool = False, |
| 294 | use_kv_connector: bool = False, |
| 295 | use_pcp: bool = False, |
| 296 | use_adaptive_verification: bool = False, |
| 297 | use_dcp: bool = False, |
| 298 | use_rswa: bool = False, |
| 299 | ) -> list[str]: |
| 300 | invalid_reasons = [] |
| 301 | if not cls.supports_head_size(head_size): |
| 302 | invalid_reasons.append("head_size not supported") |
| 303 | if not cls.supports_dtype(dtype): |
| 304 | invalid_reasons.append("dtype not supported") |
| 305 | if not cls.supports_kv_cache_dtype(kv_cache_dtype): |
| 306 | invalid_reasons.append("kv_cache_dtype not supported") |
| 307 | if not cls.supports_block_size(block_size): |
| 308 | invalid_reasons.append("block_size not supported") |
| 309 | if use_mm_prefix and not cls.supports_mm_prefix(): |
| 310 | invalid_reasons.append( |
| 311 | "partial multimodal token full attention not supported" |
| 312 | ) |
| 313 | if use_mla != cls.is_mla(): |
| 314 | if use_mla: |
| 315 | invalid_reasons.append("MLA not supported") |
| 316 | else: |
| 317 | invalid_reasons.append("non-MLA not supported") |
| 318 | if has_sink and not cls.supports_sink(): |
| 319 | invalid_reasons.append("attention sinks not supported") |
| 320 | if use_sparse != cls.is_sparse(): |
| 321 | if use_sparse: |
| 322 | invalid_reasons.append("sparse not supported") |
| 323 | else: |
| 324 | invalid_reasons.append("non-sparse not supported") |
| 325 | if use_per_head_quant_scales and not cls.supports_per_head_quant_scales(): |
| 326 | invalid_reasons.append("per-head quant scales not supported") |
| 327 | if not cls.supports_compute_capability(device_capability): |
| 328 | invalid_reasons.append("compute capability not supported") |
| 329 | if not cls.supports_attn_type(attn_type): |
| 330 | invalid_reasons.append(f"attention type {attn_type} not supported") |
| 331 | if has_sliding_window and not cls.supports_sliding_window(): |
| 332 | invalid_reasons.append("sliding window not supported") |
| 333 | if use_rswa and not cls.supports_rswa(): |
| 334 | invalid_reasons.append("R-SWA not supported") |
| 335 | if use_non_causal and not cls.supports_non_causal(): |
| 336 | invalid_reasons.append("non-causal attention not supported") |
| 337 | if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp(): |
| 338 | invalid_reasons.append("non-causal MLA attention with DCP not supported") |
| 339 | if use_batch_invariant and not cls.supports_batch_invariance(): |
| 340 | invalid_reasons.append("batch invariance not supported") |
| 341 | if use_kv_connector and not cls.supports_kv_connector(): |
| 342 | invalid_reasons.append("KV connector not supported") |
| 343 | if use_pcp and not cls.supports_pcp(): |
| 344 | invalid_reasons.append("PCP not supported") |
| 345 | if use_dcp and not cls.supports_dcp(): |
| 346 | invalid_reasons.append("DCP not supported") |
| 347 | if ( |
| 348 | use_adaptive_verification |
| 349 | and not cls.supports_device_cpu_query_lens_mismatch() |
| 350 | ): |
| 351 | invalid_reasons.append( |
| 352 | "device-cpu query lens mismatch not supported, " |
| 353 | "this is needed for adaptive verification" |
| 354 | ) |
| 355 | combination_reason = cls.supports_combination( |
| 356 | head_size, |
| 357 | dtype, |
| 358 | kv_cache_dtype, |
| 359 | block_size, |
| 360 | use_mla, |
| 361 | has_sink, |
| 362 | use_sparse, |
| 363 | use_mm_prefix, |
| 364 | device_capability, |
| 365 | ) |
| 366 | if combination_reason is not None: |
| 367 | invalid_reasons.append(combination_reason) |
| 368 | return invalid_reasons |
| 369 | |
| 370 | @classmethod |
| 371 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: |
| 372 | """Layouts this backend's kernels can consume, most preferred first, or |
| 373 | None when the kernels consume any layout and express no preference.""" |
| 374 | return None |
| 375 | |
| 376 | @classmethod |
| 377 | def is_ssm(cls) -> bool: |
| 378 | return False |
| 379 | |
| 380 | |
| 381 | class AttentionMetadata: |
| 382 | pass |
| 383 | |
| 384 | |
| 385 | T = TypeVar("T", bound=AttentionMetadata) |
| 386 | |
| 387 | |
| 388 | @dataclass |
| 389 | class CommonAttentionMetadata: |
| 390 | """Per-batch attention metadata, shared across layers and backends. |
| 391 | AttentionMetadataBuilder instances use it to construct per-layer metadata. |
| 392 | |
| 393 | For many of the tensors we keep both GPU and CPU versions. |
| 394 | """ |
| 395 | |
| 396 | query_start_loc: torch.Tensor |
| 397 | query_start_loc_cpu: torch.Tensor |
| 398 | """(batch_size + 1,), the start location of each request in query Tensor""" |
| 399 | |
| 400 | seq_lens: torch.Tensor |
| 401 | """(batch_size,), the number of computed tokens for each request""" |
| 402 | |
| 403 | num_reqs: int |
| 404 | """Number of requests""" |
| 405 | # TODO(lucas): rename to num_tokens since it may be padded and this is misleading |
| 406 | num_actual_tokens: int |
| 407 | """Total number of tokens in batch""" |
| 408 | max_query_len: int |
| 409 | """Longest query in batch""" |
| 410 | max_seq_len: int |
| 411 | """Longest context length (may be an upper bound)""" |
| 412 | |
| 413 | block_table_tensor: torch.Tensor |
| 414 | slot_mapping: torch.Tensor |
| 415 | |
| 416 | causal: bool | torch.Tensor = True |
| 417 | |
| 418 | # Needed by FastPrefillAttentionBuilder |
| 419 | logits_indices_padded: torch.Tensor | None = None |
| 420 | num_logits_indices: int | None = None |
| 421 | max_logits_per_req: int | None = None |
| 422 | |
| 423 | # Needed by CrossAttentionBuilder |
| 424 | encoder_seq_lens: torch.Tensor | None = None |
| 425 | encoder_seq_lens_cpu: np.ndarray | None = None |
| 426 | |
| 427 | dcp_local_seq_lens: torch.Tensor | None = None |
| 428 | """Sequence lengths of the local rank in decode context parallelism world""" |
| 429 | |
| 430 | dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None |
| 431 | """(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it |
| 432 | holds, on every row, the largest shard any DCP rank has of the row's |
| 433 | whole request, identical on every PCP rank, so the sparse backends pad |
| 434 | their KV gather to it.""" |
| 435 | |
| 436 | positions: torch.Tensor | None = None |
| 437 | """(num_actual_tokens,) token positions. Optional; set when the caller |
| 438 | has positions available so that builders can pre-compute position-dependent |
| 439 | sparse metadata for DeepSeek V4 C128A layers.""" |
| 440 | |
| 441 | is_prefilling: torch.Tensor | None = None |
| 442 | """(batch_size,) bool tensor: True if request is still in prefill phase |
| 443 | (num_computed_tokens < num_prompt_tokens). Used by some backends to |
| 444 | distinguish actual decodes from short extends.""" |
| 445 | |
| 446 | seq_lens_cpu_upper_bound: torch.Tensor | None = None |
| 447 | """(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows |
| 448 | and for all rows outside async spec decode; optimistic for async-spec |
| 449 | decode rows (assumes every draft was accepted). Not safe for kernels |
| 450 | that need exact per-row context lengths on decode rows.""" |
| 451 | |
| 452 | req_idx: np.ndarray | None = None |
| 453 | """(batch_size,) index of each row's request in the runner's request |
| 454 | table. Rows of one request are adjacent, so equal neighbours are PCP |
| 455 | chunks sharing one KV context.""" |
| 456 | |
| 457 | mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None |
| 458 | """PrefixLM bidirectional ranges for multimodal tokens. Maps |
| 459 | request index to list of (start, end) token position ranges |
| 460 | where bidirectional attention should apply. None for text-only |
| 461 | batches or non-PrefixLM models. A request's ranges must not overlap.""" |
| 462 | |
| 463 | rswa_prefix_lens: torch.Tensor | None = None |
| 464 | """(batch_size,) per-request prefix length (prompt/image token count) for |
| 465 | Reference Sliding Window Attention (R-SWA). Tokens with logical index below |
| 466 | this stay globally visible; later (generated) tokens additionally see a |
| 467 | fixed sliding window. None disables R-SWA. The attention backend copies this |
| 468 | into its own persistent buffer and reads ``rswa_window`` from model config.""" |
| 469 | |
| 470 | replayssm_decode_base_cpu: torch.Tensor | None = None |
| 471 | """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed |
| 472 | at the current decode run's last full-state write. write_pos counts from |
| 473 | here, so a preemption-resumed request re-anchors past the prompt boundary.""" |
| 474 | |
| 475 | _num_computed_tokens_cache: torch.Tensor | None = None |
| 476 | _token_to_req_indices_cache: torch.Tensor | None = None |
| 477 | |
| 478 | def batch_size(self) -> int: |
| 479 | return self.seq_lens.shape[0] |
| 480 | |
| 481 | def naive_query_lens(self) -> torch.Tensor: |
| 482 | """Naive because it assumes that query ends where the next query starts.""" |
| 483 | return self.query_start_loc[1:] - self.query_start_loc[:-1] |
| 484 | |
| 485 | def replace(self, **kwargs) -> "CommonAttentionMetadata": |
| 486 | return replace(self, **kwargs) |
| 487 | |
| 488 | def compute_num_computed_tokens(self) -> torch.Tensor: |
| 489 | """Compute num_computed_tokens on device (seq_lens - query_lens).""" |
| 490 | if self._num_computed_tokens_cache is None: |
| 491 | query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1] |
| 492 | self._num_computed_tokens_cache = self.seq_lens - query_lens |
| 493 | return self._num_computed_tokens_cache |
| 494 | |
| 495 | def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor: |
| 496 | """Build or reuse the per-token request index mapping.""" |
| 497 | num_tokens = self.num_actual_tokens |
| 498 | if self._token_to_req_indices_cache is not None: |
| 499 | assert self._token_to_req_indices_cache.device == buffer.device |
| 500 | assert self._token_to_req_indices_cache.dtype == torch.int32 |
| 501 | assert self._token_to_req_indices_cache.shape[0] >= num_tokens |
| 502 | return self._token_to_req_indices_cache[:num_tokens] |
| 503 | |
| 504 | # Built from the device query_start_loc: adaptive verification decides the |
| 505 | # per-request draft split on device, so the CPU copy carries the right total |
| 506 | # but not the right per-request boundaries. Padding requests have a query |
| 507 | # length of zero and are skipped by the device boundary search. |
| 508 | num_mapped_tokens = int(self.query_start_loc_cpu[-1]) |
| 509 | from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel |
| 510 | |
| 511 | num_output_tokens = max(num_mapped_tokens, num_tokens) |
| 512 | assert buffer.shape[0] >= num_output_tokens |
| 513 | _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)]( |
| 514 | self.query_start_loc, |
| 515 | buffer, |
| 516 | self.query_start_loc.shape[0] - 1, |
| 517 | num_mapped_tokens, |
| 518 | num_output_tokens, |
| 519 | num_warps=4, |
| 520 | ) |
| 521 | self._token_to_req_indices_cache = buffer[:num_output_tokens] |
| 522 | return self._token_to_req_indices_cache[:num_tokens] |
| 523 | |
| 524 | # TODO(lucas): remove once we have FULL-CG spec-decode support |
| 525 | def unpadded( |
| 526 | self, num_actual_tokens: int, num_actual_reqs: int |
| 527 | ) -> "CommonAttentionMetadata": |
| 528 | maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None |
| 529 | return CommonAttentionMetadata( |
| 530 | query_start_loc=self.query_start_loc[: num_actual_reqs + 1], |
| 531 | query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1], |
| 532 | seq_lens=self.seq_lens[:num_actual_reqs], |
| 533 | num_reqs=num_actual_reqs, |
| 534 | num_actual_tokens=num_actual_tokens, |
| 535 | max_query_len=self.max_query_len, |
| 536 | max_seq_len=self.max_seq_len, |
| 537 | block_table_tensor=self.block_table_tensor[:num_actual_reqs], |
| 538 | slot_mapping=self.slot_mapping[:num_actual_tokens], |
| 539 | causal=self.causal[:num_actual_reqs] |
| 540 | if isinstance(self.causal, torch.Tensor) |
| 541 | else self.causal, |
| 542 | logits_indices_padded=self.logits_indices_padded, |
| 543 | num_logits_indices=self.num_logits_indices, |
| 544 | max_logits_per_req=self.max_logits_per_req, |
| 545 | encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), |
| 546 | encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu), |
| 547 | dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), |
| 548 | dcp_local_seq_lens_cpu_upper_bound=maybe_slice_reqs( |
| 549 | self.dcp_local_seq_lens_cpu_upper_bound |
| 550 | ), |
| 551 | seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound), |
| 552 | is_prefilling=maybe_slice_reqs(self.is_prefilling), |
| 553 | req_idx=maybe_slice_reqs(self.req_idx), |
| 554 | rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), |
| 555 | replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), |
| 556 | ) |
| 557 | |
| 558 | |
| 559 | M = TypeVar("M") |
| 560 | |
| 561 | |
| 562 | class AttentionCGSupport(Enum): |
| 563 | """Constants for the cudagraph support of the attention backend |
| 564 | Here we do not consider the cascade attention, as currently |
| 565 | it is never cudagraph supported.""" |
| 566 | |
| 567 | ALWAYS = 3 |
| 568 | """Cudagraph always supported; supports mixed-prefill-decode""" |
| 569 | UNIFORM_BATCH = 2 |
| 570 | """Cudagraph supported for batches the only contain query lengths that are |
| 571 | the same, this can be used for spec-decode |
| 572 | i.e. "decodes" are 1 + num_speculative_tokens""" |
| 573 | UNIFORM_SINGLE_TOKEN_DECODE = 1 |
| 574 | """Cudagraph supported for batches the only contain query_len==1 decodes""" |
| 575 | NEVER = 0 |
| 576 | """NO cudagraph support""" |
| 577 | |
| 578 | |
| 579 | def max_decode_query_len(vllm_config: "VllmConfig") -> int: |
| 580 | """Widest request a spec-as-decode builder treats as a decode. |
| 581 | |
| 582 | On model runner V2 this is a verification request, 1 + |
| 583 | num_speculative_tokens: no V2 speculator builds a wider query. The reorder |
| 584 | threshold and the code that mirrors the builders' decode/prefill split all |
| 585 | read this, so they cannot drift apart. |
| 586 | """ |
| 587 | speculative_config = getattr(vllm_config, "speculative_config", None) |
| 588 | if speculative_config is None or speculative_config.num_speculative_tokens is None: |
| 589 | return 1 |
| 590 | num_speculative_tokens = speculative_config.num_speculative_tokens |
| 591 | max_query_len = 1 + num_speculative_tokens |
| 592 | if speculative_config.parallel_drafting and not vllm_config.use_v2_model_runner: |
| 593 | # Model runner V1 only; remove with it. Its parallel drafter appends up |
| 594 | # to num_speculative_tokens mask slots to each verified request. |
| 595 | max_query_len += num_speculative_tokens |
| 596 | return max_query_len |
| 597 | |
| 598 | |
| 599 | class AttentionMetadataBuilder(ABC, Generic[M]): |
| 600 | # Does this backend/builder support CUDA Graphs for attention (default: no). |
| 601 | # Do not access directly. Call get_cudagraph_support() instead. |
| 602 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER |
| 603 | # Does this backend/builder reorder the batch? |
| 604 | # If not, set this to None. Otherwise set it to the query |
| 605 | # length that will be pulled into the front of the batch. |
| 606 | reorder_batch_threshold: int | None = None |
| 607 | # Does this backend/builder support updating the block table in existing |
| 608 | # metadata |
| 609 | supports_update_block_table: bool = False |
| 610 | # Whether the builder constructor requires the block-table width. |
| 611 | requires_block_table_width: ClassVar[bool] = False |
| 612 | # Whether all step-dependent draft decode metadata can be updated in place, |
| 613 | # allowing one metadata build to be reused across autoregressive draft steps. |
| 614 | supports_draft_decode_metadata_update: bool = False |
| 615 | |
| 616 | @abstractmethod |
| 617 | def __init__( |
| 618 | self, |
| 619 | kv_cache_spec: "KVCacheSpec", |
| 620 | layer_names: list[str], |
| 621 | vllm_config: "VllmConfig", |
| 622 | device: torch.device, |
| 623 | ): |
| 624 | self.kv_cache_spec = kv_cache_spec |
| 625 | self.layer_names = layer_names |
| 626 | self.vllm_config = vllm_config |
| 627 | self.device = device |
| 628 | self.kernel_block_size: int | None = None |
| 629 | |
| 630 | def set_kernel_block_size(self, kernel_block_size: int) -> None: |
| 631 | self.kernel_block_size = kernel_block_size |
| 632 | |
| 633 | @classmethod |
| 634 | def get_cudagraph_support( |
| 635 | cls: type["AttentionMetadataBuilder"], |
| 636 | vllm_config: "VllmConfig", |
| 637 | kv_cache_spec: "KVCacheSpec", |
| 638 | ) -> AttentionCGSupport: |
| 639 | """Get the cudagraph support level of this builder class.""" |
| 640 | return cls._cudagraph_support |
| 641 | |
| 642 | @classmethod |
| 643 | def get_varlen_cudagraph_max_query_len( |
| 644 | cls: type["AttentionMetadataBuilder"], |
| 645 | vllm_config: "VllmConfig", |
| 646 | kv_cache_spec: "KVCacheSpec", |
| 647 | ) -> int | None: |
| 648 | """Get the largest per-request query length L of the variable-length |
| 649 | decode batches a FULL cudagraph of this builder class can replay. |
| 650 | |
| 651 | The graph must replay any decode batch in which every real request has |
| 652 | between 1 and L query tokens and every padding request has 0. Lengths |
| 653 | come from the device query_start_loc; host metadata carries only the |
| 654 | token count and an upper bound on the per-request length. Batches with |
| 655 | a prefill never replay these graphs. Whether host metadata may |
| 656 | understate device query lengths at all is a separate backend question; |
| 657 | see supports_device_cpu_query_lens_mismatch(). |
| 658 | |
| 659 | Returns: |
| 660 | L, or None for builders reporting ALWAYS, which replay any batch, |
| 661 | and for builders that cannot replay variable-length batches. |
| 662 | |
| 663 | """ |
| 664 | return None |
| 665 | |
| 666 | def _init_reorder_batch_threshold( |
| 667 | self, |
| 668 | reorder_batch_threshold: int | None = 1, |
| 669 | supports_spec_as_decode: bool = False, |
| 670 | supports_dcp_with_varlen: bool = False, |
| 671 | ) -> None: |
| 672 | self.reorder_batch_threshold = reorder_batch_threshold |
| 673 | if self.reorder_batch_threshold is not None and supports_spec_as_decode: |
| 674 | # If the backend supports spec-as-decode kernels, then we can set |
| 675 | # the reorder_batch_threshold based on the number of speculative |
| 676 | # tokens from the config. |
| 677 | speculative_config = self.vllm_config.speculative_config |
| 678 | if ( |
| 679 | speculative_config is not None |
| 680 | and speculative_config.num_speculative_tokens is not None |
| 681 | ): |
| 682 | self.reorder_batch_threshold = max( |
| 683 | self.reorder_batch_threshold, |
| 684 | max_decode_query_len(self.vllm_config), |
| 685 | ) |
| 686 | |
| 687 | if ( |
| 688 | self.vllm_config.parallel_config.decode_context_parallel_size > 1 |
| 689 | and not supports_dcp_with_varlen |
| 690 | ): |
| 691 | self.reorder_batch_threshold = 1 |
| 692 | |
| 693 | @abstractmethod |
| 694 | def build( |
| 695 | self, |
| 696 | common_prefix_len: int, |
| 697 | common_attn_metadata: CommonAttentionMetadata, |
| 698 | fast_build: bool = False, |
| 699 | ) -> M: |
| 700 | """Central method that builds attention metadata. |
| 701 | Some builders (MLA) require reorder_batch to be called prior to build. |
| 702 | |
| 703 | Args: |
| 704 | common_prefix_len: The length of the common prefix of the batch. |
| 705 | common_attn_metadata: The common attention metadata. |
| 706 | fast_build: The meta-data will prioritize speed of building over |
| 707 | then speed at execution. Can be used for spec-decode where the |
| 708 | result of a build call may only be used for few layers/iters. |
| 709 | |
| 710 | """ |
| 711 | raise NotImplementedError |
| 712 | |
| 713 | def update_block_table( |
| 714 | self, |
| 715 | metadata: M, |
| 716 | blk_table: torch.Tensor, |
| 717 | slot_mapping: torch.Tensor, |
| 718 | ) -> M: |
| 719 | """Update the block table for the attention metadata. |
| 720 | Faster when theres multiple kv-cache groups that create virtually the |
| 721 | same metadata but just with different block tables. |
| 722 | |
| 723 | Only needs to be implemented if supports_update_block_table is True. |
| 724 | """ |
| 725 | raise NotImplementedError |
| 726 | |
| 727 | def build_for_cudagraph_capture( |
| 728 | self, common_attn_metadata: CommonAttentionMetadata |
| 729 | ) -> M: |
| 730 | """Build attention metadata for CUDA graph capture. Uses build by default. |
| 731 | Subclasses that override this method should call self.build or |
| 732 | super().build_for_cudagraph_capture. |
| 733 | """ |
| 734 | return self.build( |
| 735 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 736 | ) |
| 737 | |
| 738 | def build_for_drafting( |
| 739 | self, |
| 740 | common_attn_metadata: CommonAttentionMetadata, |
| 741 | draft_index: int, |
| 742 | ) -> M: |
| 743 | """Build attention metadata for draft model. Uses build by default. |
| 744 | |
| 745 | Args: |
| 746 | common_attn_metadata: The common attention metadata. |
| 747 | draft_index: The index of the current draft operation. |
| 748 | When speculating a chain of tokens, this index refers to the |
| 749 | draft attempt for the i-th token. |
| 750 | For tree-based attention, this index instead refers to the |
| 751 | draft attempt for the i-th level in the tree of tokens. |
| 752 | |
| 753 | """ |
| 754 | return self.build( |
| 755 | common_prefix_len=0, |
| 756 | common_attn_metadata=common_attn_metadata, |
| 757 | fast_build=True, |
| 758 | ) |
| 759 | |
| 760 | def update_draft_decode_metadata(self, metadata: M) -> None: |
| 761 | """Update step-dependent draft decode metadata in place. |
| 762 | |
| 763 | The fused draft loop may call this method during full CUDA graph |
| 764 | capture. CUDA graph replay does not run this Python method, so |
| 765 | implementations must emit capture-safe operations and keep replayed |
| 766 | tensor state in persistent storage. |
| 767 | """ |
| 768 | raise NotImplementedError |
| 769 | |
| 770 | def use_cascade_attention( |
| 771 | self, |
| 772 | common_prefix_len: int, |
| 773 | query_lens: np.ndarray, |
| 774 | num_query_heads: int, |
| 775 | num_kv_heads: int, |
| 776 | use_alibi: bool, |
| 777 | use_sliding_window: bool, |
| 778 | use_local_attention: bool, |
| 779 | num_sms: int, |
| 780 | dcp_world_size: int, |
| 781 | ) -> bool: |
| 782 | return False |
| 783 | |
| 784 | |
| 785 | class AttentionLayer(Protocol): |
| 786 | _q_scale: torch.Tensor |
| 787 | _k_scale: torch.Tensor |
| 788 | _k_scale_cpu: torch.Tensor |
| 789 | _v_scale: torch.Tensor |
| 790 | _v_scale_cpu: torch.Tensor |
| 791 | _q_scale_float: float |
| 792 | _k_scale_float: float |
| 793 | _v_scale_float: float |
| 794 | _prob_scale: torch.Tensor |
| 795 | |
| 796 | def forward( |
| 797 | self, |
| 798 | query: torch.Tensor, |
| 799 | key: torch.Tensor, |
| 800 | value: torch.Tensor, |
| 801 | kv_cache: torch.Tensor, |
| 802 | attn_metadata: AttentionMetadata, |
| 803 | ) -> torch.Tensor: ... |
| 804 | |
| 805 | |
| 806 | class AttentionImplBase(ABC, Generic[T]): |
| 807 | """Base class for attention implementations. |
| 808 | |
| 809 | Contains common attributes and initialization logic shared by both |
| 810 | standard AttentionImpl and MLAAttentionImpl. Does not define a forward |
| 811 | method - subclasses define their own forward interfaces. |
| 812 | """ |
| 813 | |
| 814 | # Whether this impl uses a sparse (top-k) attention path. Used by MLA to |
| 815 | # route between the dense-MHA prefill and sparse-MQA paths. |
| 816 | is_sparse: ClassVar[bool] = False |
| 817 | |
| 818 | # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse |
| 819 | # impls without one run the top-k MQA path for all requests. |
| 820 | supports_dense_mha_prefill: ClassVar[bool] = True |
| 821 | |
| 822 | # Required attributes that all impls should have |
| 823 | num_heads: int |
| 824 | head_size: int |
| 825 | scale: float |
| 826 | |
| 827 | # Whether the attention impl can return the softmax lse for decode. |
| 828 | # Some features like decode context parallelism require the softmax lse. |
| 829 | can_return_lse_for_decode: bool = False |
| 830 | |
| 831 | # Base of the logarithm used by this backend when returning softmax lse. |
| 832 | # True => natural log (lse = ln(sum(exp(qk)))) |
| 833 | # -- e.g. Triton MLA, FlashAttention, FlashMLA, Cutlass MLA |
| 834 | # False => base 2 (lse = log2(sum(exp(qk)))) |
| 835 | # -- e.g. FlashInfer trtllm-gen MLA |
| 836 | # The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in |
| 837 | # vllm/v1/attention/ops/dcp.py) branches on this via its IS_BASE_E |
| 838 | # constexpr; getting it wrong silently corrupts the cross-shard |
| 839 | # softmax denominator. |
| 840 | lse_base_on_e: bool = True |
| 841 | |
| 842 | # Whether the attention impl supports Prefill Context Parallelism. |
| 843 | supports_pcp: bool = False |
| 844 | # Whether the attention impl supports Decode Context Parallelism. |
| 845 | supports_dcp: bool = False |
| 846 | # Whether the attention impl(or ops) supports MTP |
| 847 | # when cp_kv_cache_interleave_size > 1 |
| 848 | supports_mtp_with_cp_non_trivial_interleave_size: bool = False |
| 849 | |
| 850 | # some attention backends might not always want to return lse |
| 851 | # even if they can return lse (for efficiency reasons) |
| 852 | need_to_return_lse_for_decode: bool = False |
| 853 | |
| 854 | # Whether this attention implementation supports pre-quantized query input. |
| 855 | # When True, the attention layer will quantize queries before passing them |
| 856 | # to this backend, allowing torch.compile to fuse the quantization with |
| 857 | # previous operations. This is typically supported when using FP8 KV cache |
| 858 | # with compatible attention kernels (e.g., TRT-LLM). |
| 859 | # Subclasses should set this in __init__. |
| 860 | # TODO add support to more backends: |
| 861 | # https://github.com/vllm-project/vllm/issues/25584 |
| 862 | supports_quant_query_input: bool = False |
| 863 | |
| 864 | dcp_world_size: int |
| 865 | dcp_rank: int |
| 866 | |
| 867 | pcp_world_size: int |
| 868 | pcp_rank: int |
| 869 | |
| 870 | total_cp_world_size: int |
| 871 | total_cp_rank: int |
| 872 | |
| 873 | def __new__(cls, *args, **kwargs): |
| 874 | # use __new__ so that all subclasses will call this |
| 875 | self = super().__new__(cls) |
| 876 | from vllm.config import get_current_vllm_config_or_none |
| 877 | from vllm.distributed.parallel_state import get_dcp_world_size_and_rank |
| 878 | |
| 879 | # Replicated drafts run at DCP=1 inside a DCP target's process group. |
| 880 | config = get_current_vllm_config_or_none() |
| 881 | self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank( |
| 882 | config is None or config.parallel_config.decode_context_parallel_size > 1 |
| 883 | ) |
| 884 | try: |
| 885 | from vllm.distributed.parallel_state import get_pcp_group |
| 886 | |
| 887 | self.pcp_world_size = get_pcp_group().world_size |
| 888 | self.pcp_rank = get_pcp_group().rank_in_group |
| 889 | except AssertionError: |
| 890 | self.pcp_world_size = 1 |
| 891 | self.pcp_rank = 0 |
| 892 | self.total_cp_world_size = self.dcp_world_size |
| 893 | self.total_cp_rank = self.dcp_rank |
| 894 | |
| 895 | self.need_to_return_lse_for_decode = ( |
| 896 | self.dcp_world_size > 1 and self.can_return_lse_for_decode |
| 897 | ) |
| 898 | return self |
| 899 | |
| 900 | def process_weights_after_loading(self, act_dtype: torch.dtype): |
| 901 | pass |
| 902 | |
| 903 | def get_impl_variants(self) -> tuple["AttentionImplBase", ...]: |
| 904 | return (self,) |
| 905 | |
| 906 | def prepare_for_batch(self, attn_metadata: T | None) -> None: |
| 907 | """Prepare implementation-specific state for the current batch.""" |
| 908 | |
| 909 | |
| 910 | class AttentionImpl(AttentionImplBase[T], Generic[T]): |
| 911 | """Standard attention implementation with forward method.""" |
| 912 | |
| 913 | kv_cache_dtype: str |
| 914 | |
| 915 | @property |
| 916 | def kv_quant_mode(self) -> "KVQuantMode": |
| 917 | """Return the KV cache quantization mode for this layer.""" |
| 918 | return get_kv_quant_mode(self.kv_cache_dtype) |
| 919 | |
| 920 | @abstractmethod |
| 921 | def __init__( |
| 922 | self, |
| 923 | num_heads: int, |
| 924 | head_size: int, |
| 925 | scale: float, |
| 926 | num_kv_heads: int | None = None, |
| 927 | alibi_slopes: list[float] | None = None, |
| 928 | sliding_window: int | None = None, |
| 929 | kv_cache_dtype: str = "auto", |
| 930 | logits_soft_cap: float | None = None, |
| 931 | attn_type: str = AttentionType.DECODER, |
| 932 | kv_sharing_target_layer_name: str | None = None, |
| 933 | ) -> None: |
| 934 | raise NotImplementedError |
| 935 | |
| 936 | @abstractmethod |
| 937 | def forward( |
| 938 | self, |
| 939 | layer: AttentionLayer, |
| 940 | query: torch.Tensor, |
| 941 | key: torch.Tensor, |
| 942 | value: torch.Tensor, |
| 943 | kv_cache: torch.Tensor, |
| 944 | attn_metadata: T, |
| 945 | output: torch.Tensor, |
| 946 | output_scale: torch.Tensor | None = None, |
| 947 | output_block_scale: torch.Tensor | None = None, |
| 948 | ) -> torch.Tensor: |
| 949 | raise NotImplementedError |
| 950 | |
| 951 | def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool: |
| 952 | """Does this attention implementation support fused output quantization. |
| 953 | This is used by the AttnFusionPass to only fuse output quantization |
| 954 | onto implementations that support it. |
| 955 | |
| 956 | Args: |
| 957 | quant_key: QuantKey object that describes the quantization op |
| 958 | |
| 959 | Returns: |
| 960 | is fusion supported for this type of quantization |
| 961 | |
| 962 | """ |
| 963 | return False |
| 964 | |
| 965 | def fused_qk_norm_rope_kvcache_supported(self): |
| 966 | """Does this attention implementation support fused QKNorm+RoPE+KVCache fusion. |
| 967 | This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops |
| 968 | with the RoPE ops and the KV cache update for implementations that support it. |
| 969 | """ |
| 970 | return False |
| 971 | |
| 972 | def fused_qk_norm_mrope_kvcache_supported(self): |
| 973 | """Whether this implementation supports fused QKNorm+MRoPE+KVCache.""" |
| 974 | return False |
| 975 | |
| 976 | def fused_rope_kvcache_supported(self): |
| 977 | """Does this attention implementation support RoPE+KVCache fusion. |
| 978 | This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops |
| 979 | with the KV cache update for implementations that support it. |
| 980 | """ |
| 981 | return False |
| 982 | |
| 983 | def do_qk_norm_rope_kvcache_update( |
| 984 | self, |
| 985 | layer: AttentionLayer, |
| 986 | qkv: torch.Tensor, |
| 987 | q_out: torch.Tensor, |
| 988 | k_out: torch.Tensor, |
| 989 | positions: torch.Tensor, |
| 990 | q_weight: torch.Tensor, |
| 991 | k_weight: torch.Tensor, |
| 992 | rms_norm_eps: float, |
| 993 | cos_sin_cache: torch.Tensor, |
| 994 | is_neox: bool, |
| 995 | kv_cache: torch.Tensor, |
| 996 | layer_slot_mapping: torch.Tensor, |
| 997 | ): |
| 998 | """If `fused_qk_norm_rope_kvcache_supported` returns True, this method |
| 999 | will be called by the fused custom op. Applies QK-norm + RoPE and |
| 1000 | writes K/V to the KV cache. Results are written to the pre-allocated |
| 1001 | q_out and k_out tensors; V is split from QKV at the graph level. |
| 1002 | """ |
| 1003 | raise NotImplementedError |
| 1004 | |
| 1005 | def do_qk_norm_mrope_kvcache_update( |
| 1006 | self, |
| 1007 | layer: AttentionLayer, |
| 1008 | qkv: torch.Tensor, |
| 1009 | q_out: torch.Tensor, |
| 1010 | positions: torch.Tensor, |
| 1011 | q_weight: torch.Tensor, |
| 1012 | k_weight: torch.Tensor, |
| 1013 | rms_norm_eps: float, |
| 1014 | cos_sin_cache: torch.Tensor, |
| 1015 | is_neox: bool, |
| 1016 | mrope_section: tuple[int, int, int], |
| 1017 | is_interleaved: bool, |
| 1018 | rotary_dim: int, |
| 1019 | kv_cache: torch.Tensor, |
| 1020 | layer_slot_mapping: torch.Tensor, |
| 1021 | ): |
| 1022 | """Apply QK-norm and MRoPE, then write K/V to the cache.""" |
| 1023 | raise NotImplementedError |
| 1024 | |
| 1025 | def do_rope_and_kv_cache_update( |
| 1026 | self, |
| 1027 | layer: AttentionLayer, |
| 1028 | query: torch.Tensor, |
| 1029 | key: torch.Tensor, |
| 1030 | value: torch.Tensor, |
| 1031 | positions: torch.Tensor, |
| 1032 | cos_sin_cache: torch.Tensor, |
| 1033 | is_neox: bool, |
| 1034 | kv_cache: torch.Tensor, |
| 1035 | layer_slot_mapping: torch.Tensor, |
| 1036 | ): |
| 1037 | """If `fused_rope_kvcache_supported` returns True, this method will be called |
| 1038 | by torch.ops.vllm.fused_rope_and_unified_kv_cache_update |
| 1039 | to perform the inplace RoPE and KV cache update. |
| 1040 | """ |
| 1041 | raise NotImplementedError |
| 1042 | |
| 1043 | |
| 1044 | class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): |
| 1045 | """MLA attention implementation with forward_mqa and forward_mha methods.""" |
| 1046 | |
| 1047 | hisparse_cache: "HiSparseCacheHandle | None" = None |
| 1048 | supports_pcp: bool = True |
| 1049 | |
| 1050 | @abstractmethod |
| 1051 | def __init__( |
| 1052 | self, |
| 1053 | num_heads: int, |
| 1054 | head_size: int, |
| 1055 | scale: float, |
| 1056 | num_kv_heads: int, |
| 1057 | alibi_slopes: list[float] | None, |
| 1058 | sliding_window: int | None, |
| 1059 | kv_cache_dtype: str, |
| 1060 | logits_soft_cap: float | None, |
| 1061 | attn_type: str, |
| 1062 | kv_sharing_target_layer_name: str | None, |
| 1063 | # MLA Specific Arguments |
| 1064 | q_lora_rank: int | None, |
| 1065 | kv_lora_rank: int, |
| 1066 | qk_nope_head_dim: int, |
| 1067 | qk_rope_head_dim: int, |
| 1068 | qk_head_dim: int, |
| 1069 | v_head_dim: int, |
| 1070 | kv_b_proj: "ColumnParallelLinear", |
| 1071 | indexer: object | None = None, |
| 1072 | q_pad_num_heads: int | None = None, |
| 1073 | ) -> None: |
| 1074 | raise NotImplementedError |
| 1075 | |
| 1076 | def forward_mha( |
| 1077 | self, |
| 1078 | q: torch.Tensor, |
| 1079 | kv_c_normed: torch.Tensor, |
| 1080 | k_pe: torch.Tensor, |
| 1081 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1082 | attn_metadata: T, |
| 1083 | k_scale: torch.Tensor, |
| 1084 | output: torch.Tensor, |
| 1085 | output_scale: torch.Tensor | None = None, |
| 1086 | ) -> None: |
| 1087 | """MHA-style prefill forward pass.""" |
| 1088 | raise NotImplementedError |
| 1089 | |
| 1090 | @abstractmethod |
| 1091 | def forward_mqa( |
| 1092 | self, |
| 1093 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 1094 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1095 | attn_metadata: T, |
| 1096 | layer: AttentionLayer, |
| 1097 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 1098 | """MQA-style decode forward pass.""" |
| 1099 | raise NotImplementedError |
| 1100 | |
| 1101 | def fused_output_quant_supported(self, quant_key: "QuantKey"): |
| 1102 | """Does this attention implementation support fused output quantization. |
| 1103 | Since MLA quantization is done manually in forward_impl (common code), |
| 1104 | all MLA backends support it by default. |
| 1105 | """ |
| 1106 | return quant_key in ( |
| 1107 | kFp8StaticTensorSym, |
| 1108 | kNvfp4Dynamic, |
| 1109 | kFp8Dynamic128Sym, |
| 1110 | kFp8Dynamic64Sym, |
| 1111 | ) |
| 1112 | |
| 1113 | def do_kv_cache_update( |
| 1114 | self, |
| 1115 | kv_c_normed: torch.Tensor, |
| 1116 | k_pe: torch.Tensor, |
| 1117 | kv_cache: torch.Tensor, |
| 1118 | slot_mapping: torch.Tensor, |
| 1119 | kv_cache_dtype: str, |
| 1120 | k_scale: torch.Tensor, |
| 1121 | ) -> None: |
| 1122 | if kv_cache.numel() == 0: |
| 1123 | return |
| 1124 | from vllm import _custom_ops as ops |
| 1125 | |
| 1126 | ops.concat_and_cache_mla( |
| 1127 | kv_c_normed, |
| 1128 | k_pe.squeeze(1), |
| 1129 | kv_cache, |
| 1130 | slot_mapping.flatten(), |
| 1131 | kv_cache_dtype=kv_cache_dtype, |
| 1132 | scale=k_scale, |
| 1133 | ) |
| 1134 | |
| 1135 | |
| 1136 | def subclass_attention_backend( |
| 1137 | name_prefix: str, |
| 1138 | attention_backend_cls: type[AttentionBackend], |
| 1139 | builder_cls: type[AttentionMetadataBuilder[M]], |
| 1140 | ) -> type[AttentionBackend]: |
| 1141 | """Return a new subclass where `get_builder_cls` returns `builder_cls`.""" |
| 1142 | name: str = name_prefix + attention_backend_cls.__name__ # type: ignore |
| 1143 | |
| 1144 | return type( |
| 1145 | name, (attention_backend_cls,), {"get_builder_cls": lambda: builder_cls} |
| 1146 | ) |
| 1147 | |
| 1148 | |
| 1149 | def subclass_attention_backend_with_overrides( |
| 1150 | name_prefix: str, |
| 1151 | attention_backend_cls: type[AttentionBackend], |
| 1152 | overrides: dict[str, Any], |
| 1153 | ) -> type[AttentionBackend]: |
| 1154 | name: str = name_prefix + attention_backend_cls.__name__ # type: ignore |
| 1155 | return type(name, (attention_backend_cls,), overrides) |