| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Reusable attention backend composition and routing policies.""" |
| 4 | |
| 5 | import functools |
| 6 | import math |
| 7 | from typing import ClassVar, Protocol, cast |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.v1.attention.backend import ( |
| 13 | AttentionBackend, |
| 14 | AttentionCGSupport, |
| 15 | AttentionImpl, |
| 16 | AttentionMetadataBuilder, |
| 17 | CommonAttentionMetadata, |
| 18 | ) |
| 19 | from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts |
| 20 | from vllm.v1.kv_cache_interface import KVCacheSpec |
| 21 | |
| 22 | |
| 23 | def requires_mm_prefix( |
| 24 | metadata: CommonAttentionMetadata, *, unclamped_window: bool = False |
| 25 | ) -> bool: |
| 26 | ranges = metadata.mm_req_doc_ranges |
| 27 | # With one query per request there are no future keys to unmask. This |
| 28 | # also keeps single-token prefills on the captured FlashInfer path. |
| 29 | if not ranges or (metadata.max_query_len <= 1 and not unclamped_window): |
| 30 | return False |
| 31 | seq_lens = metadata.seq_lens_cpu_upper_bound |
| 32 | assert seq_lens is not None |
| 33 | starts = metadata.query_start_loc_cpu |
| 34 | for req_idx, spans in ranges.items(): |
| 35 | query_len = int(starts[req_idx + 1] - starts[req_idx]) |
| 36 | if query_len == 0 or (query_len == 1 and not unclamped_window): |
| 37 | continue |
| 38 | end = int(seq_lens[req_idx]) |
| 39 | begin = end - query_len |
| 40 | # An unclamped image can also expose keys behind the sliding window. |
| 41 | if unclamped_window and any( |
| 42 | start < stop and start < end and stop >= begin for start, stop in spans |
| 43 | ): |
| 44 | return True |
| 45 | # Prefill lengths are exact. Historical image ranges must not divert |
| 46 | # later text/decode queries to the image backend. |
| 47 | if any( |
| 48 | start < stop and start < end - 1 and stop > begin for start, stop in spans |
| 49 | ): |
| 50 | return True |
| 51 | return False |
| 52 | |
| 53 | |
| 54 | def _has_unclamped_window(layers) -> bool: |
| 55 | return any( |
| 56 | getattr(layer, "sliding_window", None) is not None |
| 57 | and not getattr(layer, "mm_prefix_clamp_sliding_window", False) |
| 58 | for layer in layers |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | class CompositeAttentionRouting(Protocol): |
| 63 | """Select a child backend and declare routing constraints for composition.""" |
| 64 | |
| 65 | capture_variant: ClassVar[int] |
| 66 | """Child used during graph capture; replay must select the same child.""" |
| 67 | |
| 68 | def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: ... |
| 69 | |
| 70 | def select(self, metadata: CommonAttentionMetadata) -> int: |
| 71 | """Return 0 for the first child or 1 for the second child.""" |
| 72 | ... |
| 73 | |
| 74 | @staticmethod |
| 75 | def get_cudagraph_support( |
| 76 | vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec |
| 77 | ) -> AttentionCGSupport: |
| 78 | """Limit graph support to batches whose routing stays invariant.""" |
| 79 | ... |
| 80 | |
| 81 | @staticmethod |
| 82 | def variant_uses_mm_prefix(variant: int) -> bool: |
| 83 | """Whether this child must support the model's image-mask requirement.""" |
| 84 | ... |
| 85 | |
| 86 | |
| 87 | class MMPrefixAttentionRouting(CompositeAttentionRouting): |
| 88 | """Route image masks to variant zero and causal queries to variant one.""" |
| 89 | |
| 90 | def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: |
| 91 | layers = vllm_config.compilation_config.static_forward_context |
| 92 | self.unclamped_window = _has_unclamped_window( |
| 93 | layers[name] for name in layer_names |
| 94 | ) |
| 95 | |
| 96 | def select(self, metadata: CommonAttentionMetadata) -> int: |
| 97 | return int( |
| 98 | not requires_mm_prefix(metadata, unclamped_window=self.unclamped_window) |
| 99 | ) |
| 100 | |
| 101 | capture_variant: ClassVar[int] = 1 |
| 102 | |
| 103 | @staticmethod |
| 104 | def get_cudagraph_support( |
| 105 | vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec |
| 106 | ) -> AttentionCGSupport: |
| 107 | if _has_unclamped_window( |
| 108 | vllm_config.compilation_config.static_forward_context.values() |
| 109 | ): |
| 110 | return AttentionCGSupport.NEVER |
| 111 | return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE |
| 112 | |
| 113 | @staticmethod |
| 114 | def variant_uses_mm_prefix(variant: int) -> bool: |
| 115 | return variant == 0 |
| 116 | |
| 117 | |
| 118 | def _intersect_kernel_block_sizes( |
| 119 | first: type[AttentionBackend], |
| 120 | second: type[AttentionBackend], |
| 121 | kv_cache_spec: KVCacheSpec | None = None, |
| 122 | ): |
| 123 | from vllm.v1.attention.backend import MultipleOf |
| 124 | |
| 125 | result: list[int | MultipleOf] = [] |
| 126 | seen: set[tuple[type, int]] = set() |
| 127 | for lhs in first.get_supported_kernel_block_sizes(kv_cache_spec): |
| 128 | for rhs in second.get_supported_kernel_block_sizes(kv_cache_spec): |
| 129 | candidate: int | MultipleOf | None = None |
| 130 | if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf): |
| 131 | candidate = MultipleOf(math.lcm(lhs.base, rhs.base)) |
| 132 | elif isinstance(lhs, MultipleOf) and isinstance(rhs, int): |
| 133 | candidate = rhs if rhs % lhs.base == 0 else None |
| 134 | elif isinstance(lhs, int) and isinstance(rhs, MultipleOf): |
| 135 | candidate = lhs if lhs % rhs.base == 0 else None |
| 136 | elif lhs == rhs: |
| 137 | candidate = lhs |
| 138 | if candidate is None: |
| 139 | continue |
| 140 | value = candidate.base if isinstance(candidate, MultipleOf) else candidate |
| 141 | key = (type(candidate), value) |
| 142 | if key not in seen: |
| 143 | result.append(candidate) |
| 144 | seen.add(key) |
| 145 | return result |
| 146 | |
| 147 | |
| 148 | @functools.cache |
| 149 | def create_composite_attention_backend( |
| 150 | general_backend: type[AttentionBackend], |
| 151 | causal_backend: type[AttentionBackend], |
| 152 | *, |
| 153 | name: str, |
| 154 | backend_name: str, |
| 155 | module: str, |
| 156 | routing_policy: type[CompositeAttentionRouting], |
| 157 | head_sizes: tuple[int, ...] = (), |
| 158 | kernel_block_sizes: tuple[int, ...] = (), |
| 159 | device_major: int | None = None, |
| 160 | ) -> type[AttentionBackend]: |
| 161 | """Compose two backends with shared storage and an explicit routing policy.""" |
| 162 | if general_backend.full_cls_name() == causal_backend.full_cls_name(): |
| 163 | return general_backend |
| 164 | general_impl_cls = general_backend.get_impl_cls() |
| 165 | causal_impl_cls = causal_backend.get_impl_cls() |
| 166 | general_builder_cls = general_backend.get_builder_cls() |
| 167 | causal_builder_cls = causal_backend.get_builder_cls() |
| 168 | if ( |
| 169 | general_backend.forward_includes_kv_cache_update |
| 170 | or causal_backend.forward_includes_kv_cache_update |
| 171 | ): |
| 172 | raise ValueError("Composite backends require separate KV cache updates") |
| 173 | |
| 174 | class CompositeAttentionImpl(AttentionImpl): |
| 175 | def __init__(self, *args, **kwargs) -> None: |
| 176 | self.general_impl = cast(AttentionImpl, general_impl_cls(*args, **kwargs)) |
| 177 | try: |
| 178 | self.causal_impl = cast(AttentionImpl, causal_impl_cls(*args, **kwargs)) |
| 179 | except TypeError as e: |
| 180 | raise ValueError( |
| 181 | f"Causal attention backend {causal_backend.get_name()} " |
| 182 | f"does not accept the attention arguments required by " |
| 183 | f"the general backend {general_backend.get_name()}: {e}" |
| 184 | ) from e |
| 185 | for name in ( |
| 186 | "num_heads", |
| 187 | "num_kv_heads", |
| 188 | "head_size", |
| 189 | "scale", |
| 190 | "kv_cache_dtype", |
| 191 | ): |
| 192 | if hasattr(self.general_impl, name): |
| 193 | setattr(self, name, getattr(self.general_impl, name)) |
| 194 | self.supports_quant_query_input = ( |
| 195 | self.general_impl.supports_quant_query_input |
| 196 | and self.causal_impl.supports_quant_query_input |
| 197 | ) |
| 198 | self.supports_pcp = self.general_impl.supports_pcp |
| 199 | self.supports_dcp = self.causal_impl.supports_dcp |
| 200 | self.can_return_lse_for_decode = self.causal_impl.can_return_lse_for_decode |
| 201 | self.lse_base_on_e = self.causal_impl.lse_base_on_e |
| 202 | self.need_to_return_lse_for_decode = ( |
| 203 | self.causal_impl.need_to_return_lse_for_decode |
| 204 | ) |
| 205 | self.supports_mtp_with_cp_non_trivial_interleave_size = ( |
| 206 | self.causal_impl.supports_mtp_with_cp_non_trivial_interleave_size |
| 207 | ) |
| 208 | |
| 209 | def get_impl_for_metadata(self, attn_metadata): |
| 210 | return ( |
| 211 | self.causal_impl |
| 212 | if getattr(attn_metadata, "_attention_backend_variant", 0) |
| 213 | else self.general_impl |
| 214 | ) |
| 215 | |
| 216 | def get_impl_variants(self): |
| 217 | return self.general_impl, self.causal_impl |
| 218 | |
| 219 | def forward( |
| 220 | self, |
| 221 | layer, |
| 222 | query, |
| 223 | key, |
| 224 | value, |
| 225 | kv_cache, |
| 226 | attn_metadata, |
| 227 | output, |
| 228 | output_scale=None, |
| 229 | output_block_scale=None, |
| 230 | ): |
| 231 | impl = self.get_impl_for_metadata(attn_metadata) |
| 232 | return impl.forward( |
| 233 | layer, |
| 234 | query, |
| 235 | key, |
| 236 | value, |
| 237 | kv_cache, |
| 238 | attn_metadata, |
| 239 | output, |
| 240 | output_scale, |
| 241 | output_block_scale, |
| 242 | ) |
| 243 | |
| 244 | def process_weights_after_loading(self, act_dtype: torch.dtype): |
| 245 | for impl in self.get_impl_variants(): |
| 246 | impl.process_weights_after_loading(act_dtype) |
| 247 | |
| 248 | def fused_output_quant_supported(self, quant_key): |
| 249 | return all( |
| 250 | impl.fused_output_quant_supported(quant_key) |
| 251 | for impl in self.get_impl_variants() |
| 252 | ) |
| 253 | |
| 254 | def fused_qk_norm_rope_kvcache_supported(self): |
| 255 | return self.general_impl.fused_qk_norm_rope_kvcache_supported() |
| 256 | |
| 257 | def do_qk_norm_rope_kvcache_update(self, *args, **kwargs): |
| 258 | method = self.general_impl.do_qk_norm_rope_kvcache_update |
| 259 | return method(*args, **kwargs) |
| 260 | |
| 261 | def fused_rope_kvcache_supported(self): |
| 262 | return self.general_impl.fused_rope_kvcache_supported() |
| 263 | |
| 264 | def do_rope_and_kv_cache_update(self, *args, **kwargs): |
| 265 | method = self.general_impl.do_rope_and_kv_cache_update |
| 266 | return method(*args, **kwargs) |
| 267 | |
| 268 | def do_kv_cache_update(self, *args, **kwargs): |
| 269 | method = self.general_impl.do_kv_cache_update # type: ignore[attr-defined] |
| 270 | return method(*args, **kwargs) |
| 271 | |
| 272 | class CompositeAttentionMetadataBuilder(AttentionMetadataBuilder): |
| 273 | requires_block_table_width = ( |
| 274 | general_builder_cls.requires_block_table_width |
| 275 | or causal_builder_cls.requires_block_table_width |
| 276 | ) |
| 277 | |
| 278 | def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): |
| 279 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 280 | args = (kv_cache_spec, layer_names, vllm_config, device) |
| 281 | # Only the variant that asked for the extra ctor kwargs receives them. |
| 282 | self.general_builder = general_builder_cls( |
| 283 | *args, |
| 284 | **(kwargs if general_builder_cls.requires_block_table_width else {}), |
| 285 | ) |
| 286 | self.causal_builder = causal_builder_cls( |
| 287 | *args, |
| 288 | **(kwargs if causal_builder_cls.requires_block_table_width else {}), |
| 289 | ) |
| 290 | self.reorder_batch_threshold = min( |
| 291 | ( |
| 292 | builder.reorder_batch_threshold |
| 293 | for builder in self._builders |
| 294 | if builder.reorder_batch_threshold is not None |
| 295 | ), |
| 296 | default=None, |
| 297 | ) |
| 298 | self.routing = routing_policy(layer_names, vllm_config) |
| 299 | |
| 300 | def set_kernel_block_size(self, kernel_block_size): |
| 301 | super().set_kernel_block_size(kernel_block_size) |
| 302 | for builder in self._builders: |
| 303 | builder.set_kernel_block_size(kernel_block_size) |
| 304 | |
| 305 | @classmethod |
| 306 | def get_cudagraph_support(cls, vllm_config, kv_cache_spec): |
| 307 | return min( |
| 308 | general_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec), |
| 309 | causal_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec), |
| 310 | routing_policy.get_cudagraph_support(vllm_config, kv_cache_spec), |
| 311 | key=lambda support: support.value, |
| 312 | ) |
| 313 | |
| 314 | def _builder(self, common_attn_metadata): |
| 315 | backend_variant = self.routing.select(common_attn_metadata) |
| 316 | builder = self.causal_builder if backend_variant else self.general_builder |
| 317 | return builder, backend_variant |
| 318 | |
| 319 | def _build(self, method, common_attn_metadata, **kwargs): |
| 320 | builder, backend_variant = self._builder(common_attn_metadata) |
| 321 | metadata = getattr(builder, method)( |
| 322 | common_attn_metadata=common_attn_metadata, **kwargs |
| 323 | ) |
| 324 | metadata._attention_backend_variant = backend_variant |
| 325 | return metadata |
| 326 | |
| 327 | def build(self, common_prefix_len, common_attn_metadata, fast_build=False): |
| 328 | return self._build( |
| 329 | "build", |
| 330 | common_attn_metadata, |
| 331 | common_prefix_len=common_prefix_len, |
| 332 | fast_build=fast_build, |
| 333 | ) |
| 334 | |
| 335 | def build_for_cudagraph_capture(self, common_attn_metadata): |
| 336 | variant = routing_policy.capture_variant |
| 337 | metadata = self._builders[variant].build_for_cudagraph_capture( |
| 338 | common_attn_metadata |
| 339 | ) |
| 340 | metadata._attention_backend_variant = variant |
| 341 | return metadata |
| 342 | |
| 343 | @property |
| 344 | def _builders(self): |
| 345 | return self.general_builder, self.causal_builder |
| 346 | |
| 347 | # Forward the optional workspace protocol so that a wrapped builder |
| 348 | # which allocates one (FlashInfer) still joins the runner's cross-group |
| 349 | # sharing instead of allocating a second buffer behind the composite. |
| 350 | def _get_workspace_buffer(self): |
| 351 | for builder in self._builders: |
| 352 | if hasattr(builder, "_get_workspace_buffer"): |
| 353 | workspace = builder._get_workspace_buffer() |
| 354 | # The runner only calls the getter for the group that |
| 355 | # provides the buffer, so hand it to the sibling here. |
| 356 | self.set_workspace_buffer(workspace) |
| 357 | return workspace |
| 358 | return None |
| 359 | |
| 360 | def set_workspace_buffer(self, workspace_buffer): |
| 361 | for builder in self._builders: |
| 362 | if hasattr(builder, "set_workspace_buffer"): |
| 363 | builder.set_workspace_buffer(workspace_buffer) |
| 364 | |
| 365 | class CompositeAttentionBackend(AttentionBackend): |
| 366 | forward_includes_kv_cache_update = False |
| 367 | general_backend_cls = general_backend |
| 368 | causal_backend_cls = causal_backend |
| 369 | |
| 370 | @staticmethod |
| 371 | def get_name(): |
| 372 | return backend_name |
| 373 | |
| 374 | @staticmethod |
| 375 | def get_impl_cls(): |
| 376 | return CompositeAttentionImpl |
| 377 | |
| 378 | @staticmethod |
| 379 | def get_builder_cls(): |
| 380 | return CompositeAttentionMetadataBuilder |
| 381 | |
| 382 | @staticmethod |
| 383 | def get_supported_kernel_block_sizes( |
| 384 | kv_cache_spec: KVCacheSpec | None = None, |
| 385 | ): |
| 386 | if kernel_block_sizes: |
| 387 | return [ |
| 388 | size |
| 389 | for size in kernel_block_sizes |
| 390 | if all( |
| 391 | backend.supports_block_size(size) |
| 392 | for backend in (general_backend, causal_backend) |
| 393 | ) |
| 394 | ] |
| 395 | return _intersect_kernel_block_sizes( |
| 396 | general_backend, causal_backend, kv_cache_spec |
| 397 | ) |
| 398 | |
| 399 | @classmethod |
| 400 | def supports_block_size(cls, block_size): |
| 401 | return bool( |
| 402 | cls.get_supported_kernel_block_sizes() |
| 403 | ) and super().supports_block_size(block_size) |
| 404 | |
| 405 | @classmethod |
| 406 | def get_supported_head_sizes(cls): |
| 407 | return list(head_sizes) |
| 408 | |
| 409 | @classmethod |
| 410 | def supports_head_size(cls, head_size): |
| 411 | return (not head_sizes or head_size in head_sizes) and all( |
| 412 | backend.supports_head_size(head_size) |
| 413 | for backend in (general_backend, causal_backend) |
| 414 | ) |
| 415 | |
| 416 | @classmethod |
| 417 | def supports_compute_capability(cls, capability): |
| 418 | return (device_major is None or capability.major == device_major) and all( |
| 419 | backend.supports_compute_capability(capability) |
| 420 | for backend in (general_backend, causal_backend) |
| 421 | ) |
| 422 | |
| 423 | @classmethod |
| 424 | def supports_dtype(cls, dtype): |
| 425 | return all( |
| 426 | backend.supports_dtype(dtype) |
| 427 | for backend in (general_backend, causal_backend) |
| 428 | ) |
| 429 | |
| 430 | @classmethod |
| 431 | def supports_kv_cache_dtype(cls, dtype): |
| 432 | return all( |
| 433 | backend.supports_kv_cache_dtype(dtype) |
| 434 | for backend in (general_backend, causal_backend) |
| 435 | ) |
| 436 | |
| 437 | @classmethod |
| 438 | def supports_mm_prefix(cls): |
| 439 | return all( |
| 440 | not routing_policy.variant_uses_mm_prefix(variant) |
| 441 | or backend.supports_mm_prefix() |
| 442 | for variant, backend in enumerate((general_backend, causal_backend)) |
| 443 | ) |
| 444 | |
| 445 | @classmethod |
| 446 | def supports_sliding_window(cls): |
| 447 | return all( |
| 448 | backend.supports_sliding_window() |
| 449 | for backend in (general_backend, causal_backend) |
| 450 | ) |
| 451 | |
| 452 | @classmethod |
| 453 | def supports_rswa(cls): |
| 454 | return all( |
| 455 | backend.supports_rswa() for backend in (general_backend, causal_backend) |
| 456 | ) |
| 457 | |
| 458 | @classmethod |
| 459 | def supports_pcp(cls): |
| 460 | return False |
| 461 | |
| 462 | @classmethod |
| 463 | def supports_dcp(cls): |
| 464 | return False |
| 465 | |
| 466 | @classmethod |
| 467 | def supports_device_cpu_query_lens_mismatch(cls): |
| 468 | return False |
| 469 | |
| 470 | @classmethod |
| 471 | def supports_combination( |
| 472 | cls, |
| 473 | head_size, |
| 474 | dtype, |
| 475 | kv_cache_dtype, |
| 476 | block_size, |
| 477 | use_mla, |
| 478 | has_sink, |
| 479 | use_sparse, |
| 480 | use_mm_prefix, |
| 481 | device_capability, |
| 482 | ): |
| 483 | for variant, backend in enumerate((general_backend, causal_backend)): |
| 484 | reason = backend.supports_combination( |
| 485 | head_size, |
| 486 | dtype, |
| 487 | kv_cache_dtype, |
| 488 | block_size, |
| 489 | use_mla, |
| 490 | has_sink, |
| 491 | use_sparse, |
| 492 | use_mm_prefix and routing_policy.variant_uses_mm_prefix(variant), |
| 493 | device_capability, |
| 494 | ) |
| 495 | if reason: |
| 496 | return f"{backend.get_name()}: {reason}" |
| 497 | return None |
| 498 | |
| 499 | @classmethod |
| 500 | def customize_spec(cls, spec): |
| 501 | general_spec = general_backend.customize_spec(spec) |
| 502 | if general_spec != causal_backend.customize_spec(spec): |
| 503 | raise ValueError("Composite backends require the same KV cache spec") |
| 504 | return general_spec |
| 505 | |
| 506 | @classmethod |
| 507 | def supported_kv_cache_layouts(cls): |
| 508 | return tuple( |
| 509 | get_supported_kv_cache_layouts((general_backend, causal_backend)) |
| 510 | ) |
| 511 | |
| 512 | CompositeAttentionBackend.__name__ = name |
| 513 | CompositeAttentionBackend.__qualname__ = name |
| 514 | CompositeAttentionBackend.__module__ = module |
| 515 | return CompositeAttentionBackend |