| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention layer with FlexAttention.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Callable |
| 7 | from dataclasses import dataclass |
| 8 | from functools import cached_property |
| 9 | from typing import ClassVar, NamedTuple |
| 10 | |
| 11 | import torch |
| 12 | import torch._dynamo.decorators |
| 13 | import torch.nn.functional as F |
| 14 | from torch.nn.attention.flex_attention import ( |
| 15 | AuxRequest, |
| 16 | BlockMask, |
| 17 | _mask_mod_signature, |
| 18 | _score_mod_signature, |
| 19 | and_masks, |
| 20 | create_block_mask, |
| 21 | flex_attention, |
| 22 | or_masks, |
| 23 | ) |
| 24 | |
| 25 | import vllm.envs as envs |
| 26 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 27 | from vllm.config.cache import CacheDType |
| 28 | from vllm.logger import init_logger |
| 29 | from vllm.model_executor.layers.attention import Attention |
| 30 | from vllm.platforms import current_platform |
| 31 | from vllm.utils.math_utils import cdiv |
| 32 | from vllm.utils.torch_utils import ( |
| 33 | async_tensor_h2d, |
| 34 | is_quantized_kv_cache, |
| 35 | is_torch_equal_or_newer, |
| 36 | ) |
| 37 | from vllm.v1.attention.backend import ( |
| 38 | AttentionBackend, |
| 39 | AttentionCGSupport, |
| 40 | AttentionImpl, |
| 41 | AttentionMetadataBuilder, |
| 42 | AttentionType, |
| 43 | CommonAttentionMetadata, |
| 44 | MultipleOf, |
| 45 | ) |
| 46 | from vllm.v1.kv_cache_interface import ( |
| 47 | AttentionSpec, |
| 48 | EncoderOnlyAttentionSpec, |
| 49 | KVCacheLayout, |
| 50 | ) |
| 51 | |
| 52 | logger = init_logger(__name__) |
| 53 | |
| 54 | torch._dynamo.config.recompile_limit = 16 |
| 55 | create_block_mask_compiled = torch.compile( |
| 56 | create_block_mask, fullgraph=True, mode="reduce-overhead" |
| 57 | ) |
| 58 | flex_attention_compiled = torch.compile(flex_attention, fullgraph=True) |
| 59 | |
| 60 | |
| 61 | def _offsets_to_doc_ids_tensor( |
| 62 | offsets_cpu: torch.Tensor, device: torch.device |
| 63 | ) -> torch.Tensor: |
| 64 | # Build on CPU (so `repeat_interleave` doesn't force a GPU->CPU sync to |
| 65 | # learn the data-dependent output length) and upload non-blocking. |
| 66 | counts = offsets_cpu[1:] - offsets_cpu[:-1] |
| 67 | doc_ids = torch.repeat_interleave( |
| 68 | torch.arange(len(counts), dtype=torch.int32), counts |
| 69 | ) |
| 70 | return async_tensor_h2d(doc_ids, device=device) |
| 71 | |
| 72 | |
| 73 | def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int): |
| 74 | difference = (multiple - (x.shape[dim] % multiple)) % multiple |
| 75 | if difference == 0: |
| 76 | return x |
| 77 | |
| 78 | dim = dim if dim >= 0 else x.ndim + dim |
| 79 | pad_list = [] |
| 80 | |
| 81 | for i in range(x.ndim - 1, dim - 1, -1): |
| 82 | if i == dim: |
| 83 | pad_list.extend([0, difference]) |
| 84 | else: |
| 85 | pad_list.extend([0, 0]) |
| 86 | |
| 87 | return F.pad(x, pad_list, mode="constant", value=0) |
| 88 | |
| 89 | |
| 90 | class FlexAttentionBackend(AttentionBackend): |
| 91 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 92 | torch.float16, |
| 93 | torch.bfloat16, |
| 94 | torch.float32, |
| 95 | ] |
| 96 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 97 | "auto", |
| 98 | "float16", |
| 99 | "bfloat16", |
| 100 | ] |
| 101 | |
| 102 | forward_includes_kv_cache_update: bool = False |
| 103 | |
| 104 | @staticmethod |
| 105 | def get_name() -> str: |
| 106 | return "FLEX_ATTENTION" |
| 107 | |
| 108 | @classmethod |
| 109 | def supports_rswa(cls) -> bool: |
| 110 | return True |
| 111 | |
| 112 | @classmethod |
| 113 | def supports_sliding_window(cls) -> bool: |
| 114 | return True |
| 115 | |
| 116 | @classmethod |
| 117 | def supports_non_causal(cls) -> bool: |
| 118 | return True |
| 119 | |
| 120 | @classmethod |
| 121 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 122 | """FlexAttention supports both decoder and encoder-only attention.""" |
| 123 | return attn_type in (AttentionType.DECODER, AttentionType.ENCODER_ONLY) |
| 124 | |
| 125 | @classmethod |
| 126 | def supports_batch_invariance(cls) -> bool: |
| 127 | return True |
| 128 | |
| 129 | @classmethod |
| 130 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 131 | # Flex flattens the (B, N) block/token axes into a single token dim, which |
| 132 | # only LBNHC's strides allow as a zero-copy view of the layer cache. |
| 133 | return (KVCacheLayout.LBNHC,) |
| 134 | |
| 135 | @classmethod |
| 136 | def supports_mm_prefix(cls) -> bool: |
| 137 | """FlexAttention supports full attention for image tokens.""" |
| 138 | return True |
| 139 | |
| 140 | @staticmethod |
| 141 | def get_impl_cls() -> type["FlexAttentionImpl"]: |
| 142 | return FlexAttentionImpl |
| 143 | |
| 144 | @staticmethod |
| 145 | def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]: |
| 146 | return FlexAttentionMetadataBuilder |
| 147 | |
| 148 | @staticmethod |
| 149 | def use_cascade_attention(*args, **kwargs) -> bool: |
| 150 | return False |
| 151 | |
| 152 | @classmethod |
| 153 | def get_supported_head_sizes(cls) -> list[int]: |
| 154 | return [] |
| 155 | |
| 156 | @staticmethod |
| 157 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 158 | return [MultipleOf(16)] |
| 159 | |
| 160 | |
| 161 | # @torch.compile(fullgraph=True, mode="reduce-overhead") |
| 162 | def physical_to_logical_mapping( |
| 163 | block_table: torch.Tensor, |
| 164 | seq_lens: torch.Tensor, |
| 165 | block_size: int, |
| 166 | total_blocks: int, |
| 167 | ) -> torch.Tensor: |
| 168 | """Creates an inverse mapping from physical block locations to logical indices. |
| 169 | |
| 170 | The original block_table maps from logical blocks to physical locations: |
| 171 | |
| 172 | Logical to Physical (Original block_table): |
| 173 | ┌───────────────────────────────────────────┐ |
| 174 | │ Request 0: │ |
| 175 | │ │ |
| 176 | │ Logical Blocks: 0 1 2 3 4 5 6 7 │ |
| 177 | │ │ │ │ │ │ │ │ │ │ |
| 178 | │ v v v v v v v v │ |
| 179 | │ Physical Blocks: 3 5 1 7 4 2 0 6 │ |
| 180 | └───────────────────────────────────────────┘ |
| 181 | |
| 182 | This function creates the inverse mapping: |
| 183 | |
| 184 | Physical to Logical (Inverse mapping): |
| 185 | ┌───────────────────────────────────────────┐ |
| 186 | │ Request 0: │ |
| 187 | │ │ |
| 188 | │ Physical Blocks: 0 1 2 3 4 5 6 7 │ |
| 189 | │ │ │ │ │ │ │ │ │ │ |
| 190 | │ v v v v v v v v │ |
| 191 | │ Logical Blocks: 6 2 5 0 4 1 7 3 │ |
| 192 | └───────────────────────────────────────────┘ |
| 193 | |
| 194 | If multiple logical blocks map to the same physical block, |
| 195 | this function returns the latest (maximum) logical block index. |
| 196 | |
| 197 | If a physical block is not mapped to by any logical block, |
| 198 | its value in the result will be -1. |
| 199 | |
| 200 | IMPORTANT: Garbage Value Protection |
| 201 | ──────────────────────────────────── |
| 202 | The block_table tensor may contain garbage values in unused positions |
| 203 | (beyond the actual sequence length). For example, if a sequence only |
| 204 | needs 3 blocks but the table has space for 8: |
| 205 | |
| 206 | block_table[0] = [10, 25, 7, 999, 1234, 888, ...] |
| 207 | ^^^^^^^^^^^^^^^^^^^^ |
| 208 | garbage values |
| 209 | |
| 210 | These garbage values can cause issues because: |
| 211 | 1. They may map to valid physical blocks by coincidence |
| 212 | 2. The scatter_ operation will assign them logical indices |
| 213 | 3. Later attention computations may incorrectly access these blocks |
| 214 | |
| 215 | To prevent this, we use seq_lens and block_size to mask out unused |
| 216 | entries, ensuring only valid block references are processed. |
| 217 | |
| 218 | IMPORTANT: Reused physical blocks (sliding-window / hybrid attention) |
| 219 | ──────────────────────────────────────────────────────────────────── |
| 220 | For some attention types, physical cache blocks can be reused over time. |
| 221 | This can cause the same physical block id to appear multiple times in a row |
| 222 | of `block_table` at different logical block indices. In that case, only the |
| 223 | latest logical block index corresponds to the current contents of that |
| 224 | physical block. Therefore, the inverse mapping must pick the maximum logical |
| 225 | block index for each physical block id. |
| 226 | |
| 227 | Args: |
| 228 | block_table: Tensor of shape [max_reqs, max_num_blocks] |
| 229 | mapping logical blocks to physical locations. May contain |
| 230 | garbage values in unused positions. |
| 231 | seq_lens: Tensor of sequence lengths for each request. Used to |
| 232 | determine how many blocks are actually needed per sequence. |
| 233 | block_size: Size of each block in tokens. Used with seq_lens to |
| 234 | compute the number of valid blocks per sequence. |
| 235 | total_blocks: Total number of physical blocks available |
| 236 | |
| 237 | Returns: |
| 238 | A tensor of shape [max_reqs, total_blocks] where each entry |
| 239 | physical_to_logical[req_id, physical_block] contains the logical |
| 240 | block index for that physical block, or -1 if unused. |
| 241 | |
| 242 | """ |
| 243 | max_reqs, max_num_blocks = block_table.shape |
| 244 | device = block_table.device |
| 245 | |
| 246 | physical_to_logical = torch.full( |
| 247 | (max_reqs, total_blocks), -1, dtype=torch.long, device=device |
| 248 | ) |
| 249 | |
| 250 | # Only process valid blocks to avoid garbage values |
| 251 | num_blocks_per_seq: torch.Tensor = cdiv(seq_lens, block_size) |
| 252 | mask = ( |
| 253 | torch.arange(max_num_blocks, device=device)[None, :] |
| 254 | < num_blocks_per_seq[:, None] |
| 255 | ) |
| 256 | |
| 257 | valid_block_table = torch.where(mask, block_table, 0) |
| 258 | valid_logical_indices = torch.where( |
| 259 | mask, torch.arange(max_num_blocks, device=device)[None, :], 0 |
| 260 | ) |
| 261 | |
| 262 | physical_to_logical.scatter_reduce_( |
| 263 | -1, valid_block_table.to(torch.int64), valid_logical_indices, reduce="amax" |
| 264 | ) |
| 265 | # NB - Seems like block 0 is always empty so we reset it manually |
| 266 | physical_to_logical[:, 0] = -1 |
| 267 | return physical_to_logical |
| 268 | |
| 269 | |
| 270 | def unique_static_unsorted( |
| 271 | x: torch.Tensor, |
| 272 | *, |
| 273 | M: int, # maximum positive value (0 is “skip me”) |
| 274 | dim: int = -1, # axis along which to deduplicate |
| 275 | ignored_val: int = 0, # value to ignore |
| 276 | pad_val: int = -1, # sentinel for unused slots |
| 277 | ) -> torch.Tensor: |
| 278 | """- Keeps the first occurrence of each non-zero value while preserving order, |
| 279 | then left-packs those uniques and fills the rest with `pad_val`. |
| 280 | - Returns (packed, keep_mask) with the *same shape* as `x`. |
| 281 | - Requires that all values be in the range [0, M] |
| 282 | - Skips ignored_val |
| 283 | |
| 284 | Works on CPU or GPU, no Python loops, O(B·N) time / O(B·M) memory. |
| 285 | |
| 286 | Example: |
| 287 | x =[3, 1, 0, 1, 2], M=3, ignored_val=0 => [3, 1, 2, -1, -1] |
| 288 | |
| 289 | """ |
| 290 | if not (-1 <= pad_val <= M): |
| 291 | raise ValueError("`pad_val` must lie in [-1, M]") |
| 292 | |
| 293 | # ── move `dim` to the end so we can treat tensor as [B, N] ────────── |
| 294 | dim = dim % x.ndim |
| 295 | x_perm = x.movedim(dim, -1) # shape [..., N] |
| 296 | B, N = x_perm.numel() // x_perm.shape[-1], x_perm.shape[-1] |
| 297 | x_flat = x_perm.reshape(B, N) # [B, N] |
| 298 | |
| 299 | device = x.device |
| 300 | idx = torch.arange(N, device=device).expand(B, N) # per-row indices |
| 301 | |
| 302 | # ── build first-occurrence table for every v ∈ [0, M] ─────────────── |
| 303 | first_idx = torch.full((B, M + 1), N, device=device) # “∞” |
| 304 | # scatter_reduce_: first_idx[b, v] = min(first_idx[b, v], i) for each i |
| 305 | first_idx.scatter_reduce_(1, x_flat, idx, reduce="amin") |
| 306 | |
| 307 | # ── keep mask: first occurrence *and* value ≠ 0 ───────────────────── |
| 308 | keep = (x_flat != ignored_val) & (idx == first_idx.gather(1, x_flat)) # [B, N] |
| 309 | |
| 310 | # ── left-pack uniques into a fresh tensor ─────────────────────────── |
| 311 | # Route non-kept entries to a garbage slot at column N so we can do a |
| 312 | # single scatter rather than using torch.nonzero (which would force a |
| 313 | # GPU->CPU sync to enumerate kept positions). |
| 314 | dest_pos = torch.cumsum(keep.to(torch.long), dim=1) - 1 # where to go |
| 315 | dest_pos = torch.where(keep, dest_pos, N) |
| 316 | packed_extended = torch.full((B, N + 1), pad_val, device=device, dtype=x_flat.dtype) |
| 317 | packed_flat = packed_extended.scatter_(1, dest_pos, x_flat)[:, :N] |
| 318 | |
| 319 | # ── restore original layout ───────────────────────────────────────── |
| 320 | packed = packed_flat.reshape(x_perm.shape).movedim(-1, dim) |
| 321 | return packed |
| 322 | |
| 323 | |
| 324 | def causal_mask_mod( |
| 325 | b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor |
| 326 | ): |
| 327 | return q_idx >= kv_idx |
| 328 | |
| 329 | |
| 330 | def bidirectional_mask_mod( |
| 331 | b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor |
| 332 | ): |
| 333 | return q_idx >= 0 |
| 334 | |
| 335 | |
| 336 | # Type alias for the block sparsity hint callable signature. |
| 337 | _block_sparsity_hint_signature = Callable[ |
| 338 | [torch.Tensor, torch.Tensor, int], torch.Tensor |
| 339 | ] |
| 340 | |
| 341 | |
| 342 | class BlockSparsityHint(NamedTuple): |
| 343 | """This prunes KV blocks from the BlockMask before the flex_attention kernel |
| 344 | is invoked, so that blocks that are fully masked never get loaded. |
| 345 | Use this with custom mask_mods that are sparse to avoid |
| 346 | the kernel iterating over all KV blocks unnecessarily. |
| 347 | |
| 348 | Attributes: |
| 349 | hint_fn: (q_block_idx [num_tokens, 1], kv_block_idx [1, num_kv_blocks], |
| 350 | block_size int) -> bool Tensor [num_tokens, num_kv_blocks]. |
| 351 | Returns True for block pairs that may contain non-masked elements. |
| 352 | |
| 353 | """ |
| 354 | |
| 355 | hint_fn: _block_sparsity_hint_signature |
| 356 | |
| 357 | |
| 358 | def copy_to_persistent(dst, src): |
| 359 | sliced = dst[tuple(slice(0, s) for s in src.shape)] |
| 360 | sliced.copy_(src) |
| 361 | return sliced |
| 362 | |
| 363 | |
| 364 | @dataclass |
| 365 | class FlexAttentionMetadata: |
| 366 | causal: bool |
| 367 | num_actual_tokens: int # Number of tokens excluding padding. |
| 368 | max_query_len: int |
| 369 | query_start_loc: torch.Tensor |
| 370 | # CPU-resident copy of query_start_loc used to derive doc_ids without a |
| 371 | # GPU->CPU sync from repeat_interleave's data-dependent output size. |
| 372 | query_start_loc_cpu: torch.Tensor |
| 373 | max_seq_len: int |
| 374 | seq_lens: torch.Tensor |
| 375 | block_table: torch.Tensor |
| 376 | slot_mapping: torch.Tensor |
| 377 | |
| 378 | use_cascade: bool |
| 379 | common_prefix_len: int |
| 380 | cu_prefix_query_lens: torch.Tensor | None |
| 381 | prefix_kv_lens: torch.Tensor | None |
| 382 | suffix_kv_lens: torch.Tensor | None |
| 383 | |
| 384 | # Block info |
| 385 | total_cache_tokens: int |
| 386 | block_size: int |
| 387 | max_possible_sequence_length: int |
| 388 | num_reqs: int |
| 389 | physical_to_logical: torch.Tensor |
| 390 | decode_offset: torch.Tensor |
| 391 | num_blocks_per_seq: torch.Tensor |
| 392 | persistent_kv_indices: torch.Tensor |
| 393 | persistent_kv_num_blocks: torch.Tensor |
| 394 | persistent_doc_ids: torch.Tensor |
| 395 | |
| 396 | # For logging. |
| 397 | num_input_tokens: int = 0 # Number of tokens including padding. |
| 398 | |
| 399 | # Flex Metadata |
| 400 | num_blocks = 0 |
| 401 | block_mask: BlockMask | None = None |
| 402 | score_mod: _score_mod_signature | None = None |
| 403 | logical_mask_mod: _mask_mod_signature = causal_mask_mod |
| 404 | uses_paged_kv: bool = True |
| 405 | doc_ids: torch.Tensor | None = None |
| 406 | direct_build: bool = True |
| 407 | q_block_size: int = 16 |
| 408 | kv_block_size: int = 16 |
| 409 | transformed_score_mod: _score_mod_signature | None = None |
| 410 | sliding_window: int | None = None |
| 411 | mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None |
| 412 | block_sparsity_hint: BlockSparsityHint | None = None |
| 413 | # Reference Sliding Window Attention (R-SWA): per-request prefix length |
| 414 | # (prompt/image tokens stay globally visible) plus a sliding window over |
| 415 | # generated tokens. Both must be set to enable. |
| 416 | rswa_prefix_lens: torch.Tensor | None = None |
| 417 | rswa_window: int | None = None |
| 418 | |
| 419 | @cached_property |
| 420 | def logical_block_ids(self): |
| 421 | return torch.arange( |
| 422 | cdiv(self.max_seq_len, self.block_size), |
| 423 | device=self.block_table.device, |
| 424 | dtype=torch.long, |
| 425 | ) |
| 426 | |
| 427 | def _convert_physical_to_logical( |
| 428 | self, |
| 429 | request_lookup: torch.Tensor, |
| 430 | q_idx: torch.Tensor, |
| 431 | physical_kv_idx: torch.Tensor, |
| 432 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 433 | """Convert physical indices to logical indices for both query and kv. |
| 434 | |
| 435 | NB is_within_lower_bound: do sequences start on block_boundaries? |
| 436 | |
| 437 | Returns: |
| 438 | tuple of (is_valid, logical_q_idx, logical_kv_idx) |
| 439 | |
| 440 | """ |
| 441 | # Map query indices to corresponding request indices |
| 442 | q_req = request_lookup[q_idx] |
| 443 | |
| 444 | # Convert physical KV indices to logical indices |
| 445 | physical_kv_block = physical_kv_idx // self.block_size |
| 446 | physical_kv_offset = physical_kv_idx % self.block_size |
| 447 | logical_block_idx = self.physical_to_logical[q_req, physical_kv_block] |
| 448 | logical_kv_idx = logical_block_idx * self.block_size + physical_kv_offset |
| 449 | |
| 450 | # Determine valid kv indices |
| 451 | live_block = logical_block_idx >= 0 |
| 452 | within_upper_bound = logical_kv_idx < self.seq_lens[q_req] |
| 453 | within_lower_bound = logical_kv_idx >= 0 |
| 454 | is_valid = live_block & within_upper_bound & within_lower_bound |
| 455 | |
| 456 | # Convert physical query indices to logical indices |
| 457 | local_q_idx = q_idx - self.query_start_loc[q_req] |
| 458 | logical_q_idx = local_q_idx + self.decode_offset[q_req] |
| 459 | |
| 460 | return is_valid, logical_q_idx, logical_kv_idx |
| 461 | |
| 462 | def get_paged_mask_mod(self) -> _mask_mod_signature: |
| 463 | """Creates the mask_mod function for FlexAttention. |
| 464 | |
| 465 | This function creates the combined mask mod function that handles: |
| 466 | 1. The paged attention block mapping |
| 467 | 2. The mapping from packed query sequences to logical query entries |
| 468 | |
| 469 | It also by defaults adds the decoding offset to the query indices. |
| 470 | With this info we create the "logical" indices that are passed to |
| 471 | mask_mod functions. This allows mask mod functions to be agnostic to |
| 472 | layout of the query and key/value tensors. |
| 473 | """ |
| 474 | assert self.doc_ids is not None |
| 475 | |
| 476 | def final_mask_mod( |
| 477 | b: torch.Tensor, |
| 478 | h: torch.Tensor, |
| 479 | q_idx: torch.Tensor, |
| 480 | physical_kv_idx: torch.Tensor, |
| 481 | ) -> torch.Tensor: |
| 482 | (is_valid, logical_q_idx, logical_kv_idx) = ( |
| 483 | self._convert_physical_to_logical(self.doc_ids, q_idx, physical_kv_idx) |
| 484 | ) |
| 485 | return is_valid & self.logical_mask_mod(b, h, logical_q_idx, logical_kv_idx) |
| 486 | |
| 487 | return final_mask_mod |
| 488 | |
| 489 | def get_bidirectional_mask_mod(self) -> _mask_mod_signature: |
| 490 | """Creates the encoder mask_mod function for FlexAttention. |
| 491 | |
| 492 | Since the encoder bidirectional attention doesn't run with |
| 493 | KV cache, this function creates a mask based on the |
| 494 | packed query sequences. |
| 495 | """ |
| 496 | # Create a lookup mapping from query indices -> request number |
| 497 | request_lookup = _offsets_to_doc_ids_tensor( |
| 498 | self.query_start_loc_cpu, self.query_start_loc.device |
| 499 | ) |
| 500 | |
| 501 | def final_mask_mod( |
| 502 | b: torch.Tensor, |
| 503 | h: torch.Tensor, |
| 504 | q_idx: torch.Tensor, |
| 505 | kv_idx: torch.Tensor, |
| 506 | ) -> torch.Tensor: |
| 507 | return request_lookup[q_idx] == request_lookup[kv_idx] |
| 508 | |
| 509 | return final_mask_mod |
| 510 | |
| 511 | def get_sliding_window_mask_mod(self) -> _mask_mod_signature: |
| 512 | """Creates the sliding window mask_mod function for FlexAttention. |
| 513 | |
| 514 | Note that the sliding window mask here is bidirectional, we need |
| 515 | to mask it with the bidirectional/causal mask for encoder/decoder. |
| 516 | """ |
| 517 | if self.sliding_window is None: |
| 518 | raise ValueError("sliding_window must be set for sliding window attention") |
| 519 | |
| 520 | def sliding_window_mask_mod( |
| 521 | b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor |
| 522 | ): |
| 523 | return torch.abs(q_idx - kv_idx) < self.sliding_window |
| 524 | |
| 525 | def final_mask_mod( |
| 526 | b: torch.Tensor, |
| 527 | h: torch.Tensor, |
| 528 | q_idx: torch.Tensor, |
| 529 | physical_kv_idx: torch.Tensor, |
| 530 | ) -> torch.Tensor: |
| 531 | (is_valid, logical_q_idx, logical_kv_idx) = ( |
| 532 | self._convert_physical_to_logical(self.doc_ids, q_idx, physical_kv_idx) |
| 533 | ) |
| 534 | return torch.where( |
| 535 | is_valid, |
| 536 | sliding_window_mask_mod(b, h, logical_q_idx, logical_kv_idx), |
| 537 | False, |
| 538 | ) |
| 539 | |
| 540 | return final_mask_mod if self.uses_paged_kv else sliding_window_mask_mod |
| 541 | |
| 542 | def get_prefix_lm_mask_mod(self) -> _mask_mod_signature: |
| 543 | """Creates the prefix LM mask_mod function for FlexAttention.""" |
| 544 | assert self.doc_ids is not None |
| 545 | request_lookup = self.doc_ids |
| 546 | |
| 547 | def prefix_lm_mask_mod( |
| 548 | b: torch.Tensor, |
| 549 | h: torch.Tensor, |
| 550 | cu_q_idx: torch.Tensor, |
| 551 | q_idx: torch.Tensor, |
| 552 | kv_idx: torch.Tensor, |
| 553 | ): |
| 554 | mask = torch.zeros_like(q_idx, dtype=torch.bool) |
| 555 | for req, doc_range_lst in (self.mm_prefix_range or {}).items(): |
| 556 | req_mask = request_lookup[cu_q_idx] == req |
| 557 | for start, end in doc_range_lst: |
| 558 | doc_mask_q = (q_idx >= start) & (q_idx <= end) |
| 559 | doc_mask_kv = (kv_idx >= start) & (kv_idx <= end) |
| 560 | mask = mask | (req_mask & doc_mask_q & doc_mask_kv) |
| 561 | return mask |
| 562 | |
| 563 | def final_mask_mod( |
| 564 | b: torch.Tensor, |
| 565 | h: torch.Tensor, |
| 566 | q_idx: torch.Tensor, |
| 567 | physical_kv_idx: torch.Tensor, |
| 568 | ) -> torch.Tensor: |
| 569 | (is_valid, logical_q_idx, logical_kv_idx) = ( |
| 570 | self._convert_physical_to_logical(self.doc_ids, q_idx, physical_kv_idx) |
| 571 | ) |
| 572 | return torch.where( |
| 573 | is_valid, |
| 574 | prefix_lm_mask_mod(b, h, q_idx, logical_q_idx, logical_kv_idx), |
| 575 | False, |
| 576 | ) |
| 577 | |
| 578 | return final_mask_mod |
| 579 | |
| 580 | def get_rswa_mask_mod(self) -> _mask_mod_signature: |
| 581 | """Creates the Reference Sliding Window Attention (R-SWA) mask_mod. |
| 582 | |
| 583 | R-SWA keeps the whole prefix (image + prompt tokens, i.e. logical index |
| 584 | ``< prefix_len``) globally visible while generated tokens additionally |
| 585 | attend a fixed sliding window of recent tokens. This term is combined |
| 586 | with the base causal mask via logical AND, so it only ever *removes* |
| 587 | far-away generated tokens that fall outside the window and outside the |
| 588 | prefix. |
| 589 | """ |
| 590 | assert self.doc_ids is not None |
| 591 | assert self.rswa_prefix_lens is not None |
| 592 | assert self.rswa_window is not None |
| 593 | doc_ids = self.doc_ids |
| 594 | prefix_lens = self.rswa_prefix_lens |
| 595 | window = self.rswa_window |
| 596 | |
| 597 | def rswa_mask_mod( |
| 598 | q_req: torch.Tensor, |
| 599 | logical_q_idx: torch.Tensor, |
| 600 | logical_kv_idx: torch.Tensor, |
| 601 | ) -> torch.Tensor: |
| 602 | prefix_len = prefix_lens[q_req] |
| 603 | in_prefix = logical_kv_idx < prefix_len |
| 604 | in_window = (logical_q_idx - logical_kv_idx) < window |
| 605 | return in_prefix | in_window |
| 606 | |
| 607 | def final_mask_mod( |
| 608 | b: torch.Tensor, |
| 609 | h: torch.Tensor, |
| 610 | q_idx: torch.Tensor, |
| 611 | physical_kv_idx: torch.Tensor, |
| 612 | ) -> torch.Tensor: |
| 613 | (is_valid, logical_q_idx, logical_kv_idx) = ( |
| 614 | self._convert_physical_to_logical(doc_ids, q_idx, physical_kv_idx) |
| 615 | ) |
| 616 | q_req = doc_ids[q_idx] |
| 617 | return torch.where( |
| 618 | is_valid, |
| 619 | rswa_mask_mod(q_req, logical_q_idx, logical_kv_idx), |
| 620 | False, |
| 621 | ) |
| 622 | |
| 623 | return final_mask_mod |
| 624 | |
| 625 | def get_mask_mod(self): |
| 626 | # Stage-1: initialize the base mask_mod |
| 627 | # (causal mask for decoder or bidirectional mask for encoder) |
| 628 | if self.uses_paged_kv: |
| 629 | mask_mod = self.get_paged_mask_mod() |
| 630 | else: |
| 631 | mask_mod = self.get_bidirectional_mask_mod() |
| 632 | # stage-2: add external mask_mod for special attention during |
| 633 | # forwarding runtime to create the combined mask_mod. |
| 634 | if self.sliding_window is not None: |
| 635 | # Add sliding window mask for sliding window attention |
| 636 | sliding_window_mask_mod = self.get_sliding_window_mask_mod() |
| 637 | mask_mod = and_masks(mask_mod, sliding_window_mask_mod) |
| 638 | if self.mm_prefix_range: |
| 639 | # Add prefix LM mask for vision-language prefix LM attention |
| 640 | prefix_lm_mask_mod = self.get_prefix_lm_mask_mod() |
| 641 | mask_mod = or_masks(mask_mod, prefix_lm_mask_mod) |
| 642 | if self.rswa_window is not None and self.rswa_prefix_lens is not None: |
| 643 | # Reference Sliding Window Attention: AND with the base causal mask |
| 644 | # (prefix stays global, generated tokens use a sliding window). |
| 645 | mask_mod = and_masks(mask_mod, self.get_rswa_mask_mod()) |
| 646 | return mask_mod |
| 647 | |
| 648 | def get_transformed_score_mod(self) -> _score_mod_signature | None: |
| 649 | """Creates the transformed score_mod function for FlexAttention. |
| 650 | |
| 651 | This function wraps the user's score_mod to handle physical-to-logical |
| 652 | index conversion, similar to how get_mask_mod works for mask functions. |
| 653 | """ |
| 654 | if self.score_mod is None: |
| 655 | return None |
| 656 | |
| 657 | # Create a lookup mapping from query indices -> request number |
| 658 | request_lookup = _offsets_to_doc_ids_tensor( |
| 659 | self.query_start_loc_cpu, self.query_start_loc.device |
| 660 | ) |
| 661 | user_score_mod = self.score_mod |
| 662 | |
| 663 | def transformed_score_mod( |
| 664 | score: torch.Tensor, |
| 665 | b: torch.Tensor, |
| 666 | h: torch.Tensor, |
| 667 | q_idx: torch.Tensor, |
| 668 | physical_kv_idx: torch.Tensor, |
| 669 | ) -> torch.Tensor: |
| 670 | (is_valid, logical_q_idx, logical_kv_idx) = ( |
| 671 | self._convert_physical_to_logical( |
| 672 | request_lookup, q_idx, physical_kv_idx |
| 673 | ) |
| 674 | ) |
| 675 | |
| 676 | return torch.where( |
| 677 | is_valid, |
| 678 | user_score_mod( |
| 679 | score, b, h, logical_q_idx, logical_kv_idx, physical_q=q_idx |
| 680 | ), |
| 681 | -float("inf"), |
| 682 | ) |
| 683 | |
| 684 | return transformed_score_mod |
| 685 | |
| 686 | def _build_block_mask_direct(self) -> BlockMask: |
| 687 | """Direct block mask construction for paged KV cache attention. |
| 688 | |
| 689 | This method constructs the block mask directly using |
| 690 | BlockMask.from_kv_blocks which is much more efficient than the |
| 691 | generic create_block_mask approach. |
| 692 | |
| 693 | The direct path works as follows: |
| 694 | 1. For each query token, fetch blocks from block_table using max_seq_len |
| 695 | and exclude out of sliding window blocks if needed. |
| 696 | (this fetches more blocks than needed for shorter sequences) |
| 697 | 2. Group query tokens into chunks of q_block_size |
| 698 | 3. For each group, deduplicate the blocks using unique_static_unsorted |
| 699 | 4. Create BlockMask using the deduplicated block indices |
| 700 | |
| 701 | Over-estimation occurs when a group of q_block_size tokens contains |
| 702 | multiple sequence IDs (doc_ids). In this case, we fetch ALL blocks for |
| 703 | each sequence represented in the group, even though individual query |
| 704 | tokens may only need a subset of those blocks based on causal masking |
| 705 | and their position. |
| 706 | |
| 707 | """ |
| 708 | page_to_block_ratio = self.kv_block_size // self.block_size |
| 709 | if page_to_block_ratio != 1: |
| 710 | raise ValueError( |
| 711 | f"FlexAttention currently requires the cache block size " |
| 712 | f"({self.block_size}) to be equal to the kv_block_size " |
| 713 | f"({self.kv_block_size}). Please check your model's " |
| 714 | f"configuration." |
| 715 | ) |
| 716 | |
| 717 | used_pages = self.block_table[ |
| 718 | self.doc_ids, : cdiv(self.max_seq_len, self.block_size) |
| 719 | ] |
| 720 | |
| 721 | # block_table slots beyond each request's seq_len may contain garbage |
| 722 | # physical page ids (see physical_to_logical_mapping). With batched |
| 723 | # decode, max_seq_len is the batch max while shorter requests still |
| 724 | # index all columns up to that max unless masked here. |
| 725 | num_blocks = self.num_blocks_per_seq[self.doc_ids] |
| 726 | past_seq = self.logical_block_ids[None, :] >= num_blocks[:, None] |
| 727 | used_pages.masked_fill_(past_seq, 0) |
| 728 | |
| 729 | custom_hint = self.block_sparsity_hint is not None |
| 730 | use_rswa = self.rswa_window is not None and self.rswa_prefix_lens is not None |
| 731 | needs_per_q_pruning = ( |
| 732 | self.causal or self.sliding_window or custom_hint or use_rswa |
| 733 | ) |
| 734 | |
| 735 | if needs_per_q_pruning: |
| 736 | device = used_pages.device |
| 737 | assert self.doc_ids is not None |
| 738 | token_indices = torch.arange( |
| 739 | self.doc_ids.shape[0], device=device, dtype=torch.long |
| 740 | ) |
| 741 | logical_q_idx = ( |
| 742 | token_indices |
| 743 | - self.query_start_loc[self.doc_ids] |
| 744 | + self.decode_offset[self.doc_ids] |
| 745 | ) |
| 746 | block_starts = self.logical_block_ids * self.block_size |
| 747 | block_ends = block_starts + self.block_size |
| 748 | |
| 749 | if self.causal: |
| 750 | future_blocks = block_starts[None, :] > logical_q_idx[:, None] |
| 751 | used_pages.masked_fill_(future_blocks, 0) |
| 752 | |
| 753 | if self.sliding_window: |
| 754 | assert self.sliding_window is not None |
| 755 | min_kv_idx = torch.clamp( |
| 756 | logical_q_idx - (self.sliding_window - 1), min=0 |
| 757 | ) |
| 758 | min_block_idx = min_kv_idx // self.block_size |
| 759 | sliding_mask = self.logical_block_ids >= min_block_idx[:, None] |
| 760 | used_pages.masked_fill_(~sliding_mask, 0) |
| 761 | if use_rswa: |
| 762 | # R-SWA keeps prefix KV globally visible and applies a sliding |
| 763 | # window over generated tokens. Prune blocks that fall entirely |
| 764 | # in the "hole" between prefix_len and the current window so |
| 765 | # FlexAttention does not gather invalid paged-KV slots (this |
| 766 | # mirrors uniform sliding-window block pruning above). |
| 767 | assert self.rswa_prefix_lens is not None |
| 768 | assert self.rswa_window is not None |
| 769 | prefix_len = self.rswa_prefix_lens[self.doc_ids] |
| 770 | min_kv_window = torch.maximum( |
| 771 | prefix_len, |
| 772 | logical_q_idx - (self.rswa_window - 1), |
| 773 | ) |
| 774 | in_gap = (block_starts[None, :] >= prefix_len[:, None]) & ( |
| 775 | block_ends[None, :] <= min_kv_window[:, None] |
| 776 | ) |
| 777 | used_pages.masked_fill_(in_gap, 0) |
| 778 | if custom_hint: |
| 779 | assert self.block_sparsity_hint is not None |
| 780 | q_block_idx = logical_q_idx // self.block_size |
| 781 | hint_mask = self.block_sparsity_hint.hint_fn( |
| 782 | q_block_idx[:, None], |
| 783 | self.logical_block_ids[None, :], |
| 784 | self.block_size, |
| 785 | ) |
| 786 | used_pages.masked_fill_(~hint_mask, 0) |
| 787 | |
| 788 | used_pages_padded = pad_to_multiple( |
| 789 | used_pages, multiple=self.q_block_size, dim=0 |
| 790 | ) |
| 791 | used_pages_padded = used_pages_padded.reshape( |
| 792 | used_pages_padded.shape[0] // self.q_block_size, -1 |
| 793 | ) |
| 794 | used_pages_padded = used_pages_padded // page_to_block_ratio |
| 795 | kv_indices = unique_static_unsorted( |
| 796 | (used_pages_padded.long()), M=self.num_blocks |
| 797 | ).to(torch.int32) |
| 798 | kv_indices = copy_to_persistent(self.persistent_kv_indices, kv_indices) |
| 799 | |
| 800 | kv_num_blocks = (kv_indices >= 0).sum(dim=-1).to(torch.int32) |
| 801 | kv_num_blocks = copy_to_persistent(self.persistent_kv_num_blocks, kv_num_blocks) |
| 802 | |
| 803 | block_mask_kwargs = { |
| 804 | "seq_lengths": (self.num_actual_tokens, self.total_cache_tokens), |
| 805 | "kv_num_blocks": kv_num_blocks[None, None], |
| 806 | "kv_indices": kv_indices[None, None], |
| 807 | "full_kv_num_blocks": None, |
| 808 | "full_kv_indices": None, |
| 809 | "BLOCK_SIZE": (self.q_block_size, self.kv_block_size), |
| 810 | "mask_mod": self.mask_mod, |
| 811 | } |
| 812 | |
| 813 | # compute_q_blocks parameter is available in PyTorch 2.9+ |
| 814 | if is_torch_equal_or_newer("2.9.0.dev0"): |
| 815 | block_mask_kwargs["compute_q_blocks"] = False |
| 816 | return BlockMask.from_kv_blocks(**block_mask_kwargs) |
| 817 | |
| 818 | def build_block_mask(self) -> BlockMask: |
| 819 | mask_mod = self.get_mask_mod() |
| 820 | kv_len = ( |
| 821 | self.total_cache_tokens if self.uses_paged_kv else self.num_actual_tokens |
| 822 | ) |
| 823 | return create_block_mask_compiled( |
| 824 | mask_mod, |
| 825 | None, |
| 826 | None, |
| 827 | self.num_actual_tokens, |
| 828 | kv_len, |
| 829 | device=self.block_table.device, |
| 830 | BLOCK_SIZE=(self.q_block_size, self.kv_block_size), |
| 831 | ) |
| 832 | |
| 833 | def __post_init__(self): |
| 834 | assert self.use_cascade is False, "Not implemented yet." |
| 835 | assert self.common_prefix_len == 0, "Not implemented yet." |
| 836 | assert self.cu_prefix_query_lens is None, "Not implemented yet." |
| 837 | assert self.prefix_kv_lens is None, "Not implemented yet." |
| 838 | assert self.suffix_kv_lens is None, "Not implemented yet." |
| 839 | # Create a lookup mapping from query indices -> request number |
| 840 | self.doc_ids = _offsets_to_doc_ids_tensor( |
| 841 | self.query_start_loc_cpu, self.query_start_loc.device |
| 842 | ) |
| 843 | self.doc_ids = copy_to_persistent(self.persistent_doc_ids, self.doc_ids) |
| 844 | self.num_blocks = self.total_cache_tokens // self.block_size |
| 845 | |
| 846 | self.mask_mod = self.get_mask_mod() |
| 847 | self.transformed_score_mod = self.get_transformed_score_mod() |
| 848 | |
| 849 | |
| 850 | class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadata]): |
| 851 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 852 | |
| 853 | def __init__( |
| 854 | self, |
| 855 | kv_cache_spec: AttentionSpec, |
| 856 | layer_names: list[str], |
| 857 | vllm_config: VllmConfig, |
| 858 | device: torch.device, |
| 859 | ): |
| 860 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 861 | |
| 862 | self.model_config = vllm_config.model_config |
| 863 | self.parallel_config = vllm_config.parallel_config |
| 864 | self.cache_config = vllm_config.cache_config |
| 865 | |
| 866 | self.num_heads_q = self.model_config.get_num_attention_heads( |
| 867 | self.parallel_config |
| 868 | ) |
| 869 | self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config) |
| 870 | self.headdim = self.model_config.get_head_size() |
| 871 | self.block_size = kv_cache_spec.block_size |
| 872 | self.kv_cache_spec = kv_cache_spec |
| 873 | supports_small_blocks = is_torch_equal_or_newer("2.9.0.dev0") |
| 874 | uses_paged_kv = not isinstance(kv_cache_spec, EncoderOnlyAttentionSpec) |
| 875 | self.direct_build: bool = supports_small_blocks |
| 876 | |
| 877 | self.q_block_size, self.kv_block_size = self._get_block_sizes( |
| 878 | vllm_config.attention_config, |
| 879 | supports_small_blocks, |
| 880 | self.block_size, |
| 881 | uses_paged_kv, |
| 882 | ) |
| 883 | |
| 884 | if self.direct_build and self.kv_block_size != self.block_size: |
| 885 | self.direct_build = False |
| 886 | |
| 887 | self.max_model_len = self.model_config.max_model_len |
| 888 | max_num_seqs = vllm_config.scheduler_config.max_num_seqs |
| 889 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 890 | self.max_num_query_groups = cdiv(max_num_batched_tokens, self.q_block_size) |
| 891 | max_num_pages_per_seq = cdiv(self.max_model_len, self.block_size) |
| 892 | self.max_num_kv_indices = self.q_block_size * max_num_pages_per_seq |
| 893 | # R-SWA uses q_block_size=1 so block lists are not merged across requests |
| 894 | # in a q-group (mixed-length batches otherwise gather foreign paged-KV). |
| 895 | self.max_num_rswa_query_groups = max_num_batched_tokens |
| 896 | # +1 sentinel column: the flex-attention kernel's get_offset_for_next_block |
| 897 | # always prefetches kv_indices[q, kv_num_blocks] (one past the last valid |
| 898 | # entry) to compute the jump offset for the next loop iteration. When |
| 899 | # kv_num_blocks[q] == W (every page of the sequence is live), that prefetch |
| 900 | # reads column W of the persistent buffer. Without the extra column this |
| 901 | # would land on stale data from a previous step (the buffer is wider than W |
| 902 | # but is never fully zeroed), producing an out-of-bounds K/V pointer and a |
| 903 | # CUDA illegal memory access. Allocating W_max+1 columns and initialising |
| 904 | # the whole buffer to -1 ensures the sentinel slot is always safe to read. |
| 905 | self.max_num_rswa_kv_indices = max_num_pages_per_seq + 1 |
| 906 | self.persistent_kv_num_blocks = torch.empty( |
| 907 | self.max_num_query_groups, dtype=torch.int32, device=device |
| 908 | ) |
| 909 | self.persistent_rswa_kv_num_blocks = torch.empty( |
| 910 | self.max_num_rswa_query_groups, dtype=torch.int32, device=device |
| 911 | ) |
| 912 | self.persistent_offset_tensor = torch.empty( |
| 913 | max_num_seqs, dtype=torch.int32, device=device |
| 914 | ) |
| 915 | self.persistent_query_start_loc = torch.empty( |
| 916 | max_num_seqs + 1, dtype=torch.int32, device=device |
| 917 | ) |
| 918 | self.persistent_seq_lens = torch.empty( |
| 919 | max_num_seqs, dtype=torch.int32, device=device |
| 920 | ) |
| 921 | # Persistent buffer for R-SWA per-request prefix lengths so the device |
| 922 | # address stays stable across steps (required for CUDA graph replay). |
| 923 | self.rswa_window: int | None = self.model_config.rswa_window |
| 924 | self.persistent_rswa_prefix_lens: torch.Tensor | None = None |
| 925 | if self.rswa_window is not None: |
| 926 | self.persistent_rswa_prefix_lens = torch.empty( |
| 927 | max_num_seqs, dtype=torch.int32, device=device |
| 928 | ) |
| 929 | self.persistent_doc_ids = torch.empty( |
| 930 | max_num_batched_tokens, dtype=torch.int32, device=device |
| 931 | ) |
| 932 | |
| 933 | # initialize later when we can access block_table |
| 934 | self.persistent_physical_to_logical = None |
| 935 | self.persistent_kv_indices = None |
| 936 | self.persistent_rswa_kv_indices = None |
| 937 | |
| 938 | self.custom_logical_mask_mod: _mask_mod_signature | None = None |
| 939 | if self._uses_full_cudagraphs(): |
| 940 | layers = get_layers_from_vllm_config( |
| 941 | vllm_config, Attention, self.layer_names |
| 942 | ) |
| 943 | self.custom_logical_mask_mod = self._maybe_get_custom_mask_mod(layers) |
| 944 | |
| 945 | @staticmethod |
| 946 | def _get_block_sizes( |
| 947 | attn_cfg, |
| 948 | supports_small_blocks: bool, |
| 949 | cache_block_size: int, |
| 950 | uses_paged_kv: bool, |
| 951 | ) -> tuple[int, int]: |
| 952 | # Small blocks are efficient with the direct mask builder, which maps |
| 953 | # logical blocks to paged KV cache blocks without materializing and |
| 954 | # sorting a generic block mask. Encoder-only attention has no paged KV |
| 955 | # cache and therefore cannot use that path. Keeping its generic path at |
| 956 | # 16-token blocks creates very large mask-conversion graphs and makes |
| 957 | # Inductor compilation dominate the first request. |
| 958 | use_small_blocks = supports_small_blocks and uses_paged_kv |
| 959 | q_block_size = 16 if use_small_blocks else 128 |
| 960 | kv_block_size = cache_block_size if use_small_blocks else 128 |
| 961 | |
| 962 | q_block_size = attn_cfg.flex_attn_q_block_size or q_block_size |
| 963 | if (q_block_size & (q_block_size - 1)) != 0 or ( |
| 964 | attn_cfg.flex_attn_block_m is not None |
| 965 | and q_block_size % attn_cfg.flex_attn_block_m != 0 |
| 966 | ): |
| 967 | raise ValueError( |
| 968 | f"flex_attn_q_block_size must be a power of 2 " |
| 969 | f"and divisible by flex_attn_block_m, got " |
| 970 | f"{q_block_size}, {attn_cfg.flex_attn_block_m}" |
| 971 | ) |
| 972 | |
| 973 | kv_block_size = attn_cfg.flex_attn_kv_block_size or kv_block_size |
| 974 | if (kv_block_size & (kv_block_size - 1)) != 0 or ( |
| 975 | attn_cfg.flex_attn_block_n is not None |
| 976 | and kv_block_size % attn_cfg.flex_attn_block_n != 0 |
| 977 | ): |
| 978 | raise ValueError( |
| 979 | f"flex_attn_kv_block_size must be a power of 2 " |
| 980 | f"and divisible by flex_attn_block_n, got " |
| 981 | f"{kv_block_size}, {attn_cfg.flex_attn_block_n}" |
| 982 | ) |
| 983 | |
| 984 | return q_block_size, kv_block_size |
| 985 | |
| 986 | def build_for_cudagraph_capture( |
| 987 | self, common_attn_metadata: CommonAttentionMetadata |
| 988 | ) -> FlexAttentionMetadata: |
| 989 | # Use actual max_seq_len (not max_model_len) to avoid torch.compile |
| 990 | # recompilation during CUDA graph capture. |
| 991 | assert common_attn_metadata.seq_lens_cpu_upper_bound is not None |
| 992 | common_attn_metadata.max_seq_len = int( |
| 993 | common_attn_metadata.seq_lens_cpu_upper_bound.max().item() |
| 994 | ) |
| 995 | return self.build( |
| 996 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 997 | ) |
| 998 | |
| 999 | def _maybe_get_custom_mask_mod(self, layers) -> _mask_mod_signature | None: |
| 1000 | mask_mods = { |
| 1001 | getattr(layer, "logical_mask_mod", None) for layer in layers.values() |
| 1002 | } |
| 1003 | if len(mask_mods) > 1: |
| 1004 | raise ValueError( |
| 1005 | f"Found differing mask mods {mask_mods}, " |
| 1006 | "cannot use alternating mask mods w/ full CUDA graphs" |
| 1007 | ) |
| 1008 | return next(iter(mask_mods), None) |
| 1009 | |
| 1010 | def _uses_full_cudagraphs(self) -> bool: |
| 1011 | mode = self.vllm_config.compilation_config.cudagraph_mode |
| 1012 | return mode is not None and mode.has_full_cudagraphs() |
| 1013 | |
| 1014 | def build( |
| 1015 | self, |
| 1016 | common_prefix_len: int, |
| 1017 | common_attn_metadata: CommonAttentionMetadata, |
| 1018 | fast_build: bool = False, |
| 1019 | ) -> FlexAttentionMetadata: |
| 1020 | num_reqs = common_attn_metadata.num_reqs |
| 1021 | num_actual_tokens = common_attn_metadata.num_actual_tokens |
| 1022 | max_query_len = common_attn_metadata.max_query_len |
| 1023 | |
| 1024 | max_seq_len = common_attn_metadata.max_seq_len |
| 1025 | query_start_loc = common_attn_metadata.query_start_loc |
| 1026 | query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu |
| 1027 | seq_lens = common_attn_metadata.seq_lens |
| 1028 | block_table_tensor = common_attn_metadata.block_table_tensor |
| 1029 | slot_mapping = common_attn_metadata.slot_mapping |
| 1030 | num_blocks_per_seq = cdiv(seq_lens, self.block_size) |
| 1031 | |
| 1032 | use_cascade = common_prefix_len > 0 |
| 1033 | cu_prefix_query_lens = None |
| 1034 | prefix_kv_lens = None |
| 1035 | suffix_kv_lens = None |
| 1036 | if use_cascade: |
| 1037 | raise NotImplementedError( |
| 1038 | "Cascade prefix attention is not yet implemented " |
| 1039 | "for FlexAttention backend" |
| 1040 | ) |
| 1041 | |
| 1042 | block_size = self.kv_cache_spec.block_size |
| 1043 | max_possible_seq_len = self.model_config.max_model_len |
| 1044 | num_gpu_blocks = self.cache_config.num_gpu_blocks |
| 1045 | |
| 1046 | assert num_gpu_blocks is not None, ( |
| 1047 | "FlexAttention requires num_gpu_blocks to be set" |
| 1048 | ) |
| 1049 | total_cache_tokens = num_gpu_blocks * block_size |
| 1050 | |
| 1051 | inverse_block_table = physical_to_logical_mapping( |
| 1052 | block_table_tensor, seq_lens, block_size, num_gpu_blocks |
| 1053 | ) |
| 1054 | if self.persistent_physical_to_logical is None: |
| 1055 | max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs |
| 1056 | self.persistent_physical_to_logical = torch.empty( |
| 1057 | max_num_seqs, |
| 1058 | num_gpu_blocks, |
| 1059 | dtype=torch.long, |
| 1060 | device=self.device, |
| 1061 | ) |
| 1062 | |
| 1063 | if self.persistent_kv_indices is None: |
| 1064 | self.persistent_kv_indices = torch.empty( |
| 1065 | self.max_num_query_groups, |
| 1066 | self.max_num_kv_indices, |
| 1067 | dtype=torch.int32, |
| 1068 | device=self.device, |
| 1069 | ) |
| 1070 | if self.persistent_rswa_kv_indices is None: |
| 1071 | # Initialise to -1 so the +1 sentinel column (see max_num_rswa_kv_indices) |
| 1072 | # is always a safe pad value for the flex kernel's prefetch. |
| 1073 | self.persistent_rswa_kv_indices = torch.full( |
| 1074 | (self.max_num_rswa_query_groups, self.max_num_rswa_kv_indices), |
| 1075 | fill_value=-1, |
| 1076 | dtype=torch.int32, |
| 1077 | device=self.device, |
| 1078 | ) |
| 1079 | |
| 1080 | use_rswa = self.rswa_window is not None |
| 1081 | q_block_size = 1 if use_rswa else self.q_block_size |
| 1082 | persistent_kv_indices = ( |
| 1083 | self.persistent_rswa_kv_indices if use_rswa else self.persistent_kv_indices |
| 1084 | ) |
| 1085 | persistent_kv_num_blocks = ( |
| 1086 | self.persistent_rswa_kv_num_blocks |
| 1087 | if use_rswa |
| 1088 | else self.persistent_kv_num_blocks |
| 1089 | ) |
| 1090 | |
| 1091 | copy_to_persistent(self.persistent_physical_to_logical, inverse_block_table) |
| 1092 | |
| 1093 | offset_tensor = common_attn_metadata.compute_num_computed_tokens() |
| 1094 | copy_to_persistent(self.persistent_offset_tensor, offset_tensor) |
| 1095 | copy_to_persistent(self.persistent_query_start_loc, query_start_loc) |
| 1096 | copy_to_persistent(self.persistent_seq_lens, seq_lens) |
| 1097 | |
| 1098 | rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens |
| 1099 | if use_rswa and rswa_prefix_lens is not None: |
| 1100 | assert self.persistent_rswa_prefix_lens is not None |
| 1101 | copy_to_persistent(self.persistent_rswa_prefix_lens, rswa_prefix_lens) |
| 1102 | rswa_prefix_lens = self.persistent_rswa_prefix_lens |
| 1103 | |
| 1104 | uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec) |
| 1105 | logical_mask_mod = ( |
| 1106 | bidirectional_mask_mod |
| 1107 | if uses_paged_kv and not common_attn_metadata.causal |
| 1108 | else causal_mask_mod |
| 1109 | ) |
| 1110 | |
| 1111 | sliding_window = None |
| 1112 | if self._uses_full_cudagraphs(): |
| 1113 | if self.custom_logical_mask_mod is not None: |
| 1114 | logical_mask_mod = self.custom_logical_mask_mod |
| 1115 | sliding_window = getattr(self.kv_cache_spec, "sliding_window", None) |
| 1116 | |
| 1117 | out = FlexAttentionMetadata( |
| 1118 | causal=common_attn_metadata.causal, |
| 1119 | logical_mask_mod=logical_mask_mod, |
| 1120 | sliding_window=sliding_window, |
| 1121 | num_actual_tokens=num_actual_tokens, |
| 1122 | max_query_len=max_query_len, |
| 1123 | # Mask closures only index active requests, so retain full-capacity |
| 1124 | # tensor shapes to avoid recompiling when the batch shrinks to one. |
| 1125 | query_start_loc=self.persistent_query_start_loc, |
| 1126 | query_start_loc_cpu=query_start_loc_cpu, |
| 1127 | max_seq_len=max_seq_len, |
| 1128 | seq_lens=self.persistent_seq_lens, |
| 1129 | block_table=block_table_tensor, |
| 1130 | slot_mapping=slot_mapping, |
| 1131 | use_cascade=use_cascade, |
| 1132 | common_prefix_len=common_prefix_len, |
| 1133 | cu_prefix_query_lens=cu_prefix_query_lens, |
| 1134 | prefix_kv_lens=prefix_kv_lens, |
| 1135 | suffix_kv_lens=suffix_kv_lens, |
| 1136 | block_size=block_size, |
| 1137 | max_possible_sequence_length=max_possible_seq_len, |
| 1138 | num_reqs=num_reqs, |
| 1139 | physical_to_logical=self.persistent_physical_to_logical, |
| 1140 | total_cache_tokens=total_cache_tokens, |
| 1141 | decode_offset=self.persistent_offset_tensor, |
| 1142 | num_blocks_per_seq=num_blocks_per_seq, |
| 1143 | uses_paged_kv=uses_paged_kv, |
| 1144 | # FIXME(Isotr0py): direct build has issue to build bidirectional |
| 1145 | # attention block mask for encoder-only models, disable it temporarily. |
| 1146 | # see: https://github.com/vllm-project/vllm/pull/27329#issuecomment-3431484053 |
| 1147 | direct_build=self.direct_build and uses_paged_kv, |
| 1148 | q_block_size=q_block_size, |
| 1149 | kv_block_size=self.kv_block_size, |
| 1150 | persistent_kv_indices=persistent_kv_indices, |
| 1151 | persistent_kv_num_blocks=persistent_kv_num_blocks, |
| 1152 | persistent_doc_ids=self.persistent_doc_ids, |
| 1153 | mm_prefix_range=common_attn_metadata.mm_req_doc_ranges, |
| 1154 | rswa_prefix_lens=rswa_prefix_lens, |
| 1155 | rswa_window=self.rswa_window, |
| 1156 | ) |
| 1157 | |
| 1158 | # Pre-build block_mask so it is ready before CUDA graph capture. |
| 1159 | # Without this, the lazy build in forward() would run non-graph-safe |
| 1160 | # ops (e.g. torch.nonzero) inside capture. |
| 1161 | if out.block_mask is None: |
| 1162 | if out.direct_build: |
| 1163 | out.block_mask = out._build_block_mask_direct() |
| 1164 | else: |
| 1165 | out.block_mask = out.build_block_mask() |
| 1166 | |
| 1167 | return out |
| 1168 | |
| 1169 | def use_cascade_attention(self, *args, **kwargs) -> bool: |
| 1170 | return False |
| 1171 | |
| 1172 | |
| 1173 | class FlexAttentionImpl(AttentionImpl): |
| 1174 | sliding_window: int | None |
| 1175 | alibi_slopes: torch.Tensor | None |
| 1176 | logits_soft_cap: float | None |
| 1177 | mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None |
| 1178 | logical_mask_mod: _mask_mod_signature | None = None |
| 1179 | block_sparsity_hint: BlockSparsityHint | None = None |
| 1180 | |
| 1181 | def __init__( |
| 1182 | self, |
| 1183 | num_heads: int, |
| 1184 | head_size: int, |
| 1185 | scale: float, |
| 1186 | num_kv_heads: int, |
| 1187 | alibi_slopes: list[float] | None, |
| 1188 | sliding_window: int | None, |
| 1189 | kv_cache_dtype: str, |
| 1190 | logits_soft_cap: float | None = None, |
| 1191 | attn_type: AttentionType = AttentionType.DECODER, |
| 1192 | kv_sharing_target_layer_name: str | None = None, |
| 1193 | block_m: int | None = None, |
| 1194 | block_n: int | None = None, |
| 1195 | **kwargs, |
| 1196 | ) -> None: |
| 1197 | self.num_heads = num_heads |
| 1198 | self.head_size = head_size |
| 1199 | self.scale = float(scale) |
| 1200 | self.num_kv_heads = num_kv_heads |
| 1201 | self.attn_type = attn_type |
| 1202 | |
| 1203 | if attn_type not in (AttentionType.ENCODER_ONLY, AttentionType.DECODER): |
| 1204 | raise NotImplementedError( |
| 1205 | f"FlexAttention does not support {attn_type} attention" |
| 1206 | ) |
| 1207 | |
| 1208 | if alibi_slopes is not None: |
| 1209 | raise NotImplementedError( |
| 1210 | "FlexAttention does not support alibi slopes yet." |
| 1211 | ) |
| 1212 | else: |
| 1213 | self.alibi_slopes = None |
| 1214 | |
| 1215 | self.sliding_window = sliding_window |
| 1216 | |
| 1217 | self.kv_cache_dtype = kv_cache_dtype |
| 1218 | self.logits_soft_cap = logits_soft_cap |
| 1219 | if self.logits_soft_cap is not None: |
| 1220 | raise NotImplementedError( |
| 1221 | "FlexAttention does not support logits soft cap yet." |
| 1222 | ) |
| 1223 | |
| 1224 | assert self.num_heads % self.num_kv_heads == 0 |
| 1225 | self.num_queries_per_kv = self.num_heads // self.num_kv_heads |
| 1226 | |
| 1227 | if kv_sharing_target_layer_name is not None: |
| 1228 | raise NotImplementedError("FlexAttention does not support kv sharing yet.") |
| 1229 | |
| 1230 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 1231 | raise NotImplementedError( |
| 1232 | "FlexAttention does not support quantized kv-cache. Yet" |
| 1233 | ) |
| 1234 | |
| 1235 | self.block_m = 16 if envs.VLLM_BATCH_INVARIANT else None |
| 1236 | self.block_n = 16 if envs.VLLM_BATCH_INVARIANT else None |
| 1237 | |
| 1238 | if block_m is not None: |
| 1239 | self.block_m = block_m |
| 1240 | if block_n is not None: |
| 1241 | self.block_n = block_n |
| 1242 | |
| 1243 | # Optional post-attention epilogue transform |
| 1244 | self.out_transform = kwargs.get("out_transform") |
| 1245 | |
| 1246 | @staticmethod |
| 1247 | def view_as_4d(tensor: torch.Tensor) -> torch.Tensor: |
| 1248 | """View a 3d tensor as 4D.""" |
| 1249 | if tensor.ndim == 4: |
| 1250 | return tensor |
| 1251 | assert tensor.ndim == 3 |
| 1252 | return tensor[None, :, :, :] |
| 1253 | |
| 1254 | def do_kv_cache_update( |
| 1255 | self, |
| 1256 | layer: torch.nn.Module, |
| 1257 | key: torch.Tensor, |
| 1258 | value: torch.Tensor, |
| 1259 | kv_cache: torch.Tensor, |
| 1260 | slot_mapping: torch.Tensor, |
| 1261 | ) -> None: |
| 1262 | if self.attn_type == AttentionType.ENCODER_ONLY: |
| 1263 | return |
| 1264 | |
| 1265 | # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) |
| 1266 | key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) |
| 1267 | torch.ops._C_cache_ops.reshape_and_cache_flash( |
| 1268 | key, |
| 1269 | value, |
| 1270 | key_cache, |
| 1271 | value_cache, |
| 1272 | slot_mapping, |
| 1273 | self.kv_cache_dtype, |
| 1274 | layer._k_scale, |
| 1275 | layer._v_scale, |
| 1276 | ) |
| 1277 | |
| 1278 | def forward( |
| 1279 | self, |
| 1280 | layer: torch.nn.Module, |
| 1281 | query: torch.Tensor, |
| 1282 | key: torch.Tensor, |
| 1283 | value: torch.Tensor, |
| 1284 | kv_cache: torch.Tensor, |
| 1285 | attn_metadata: FlexAttentionMetadata, |
| 1286 | output: torch.Tensor, |
| 1287 | output_scale: torch.Tensor | None = None, |
| 1288 | output_block_scale: torch.Tensor | None = None, |
| 1289 | ) -> torch.Tensor: |
| 1290 | """Forward pass with FLexAttention. |
| 1291 | |
| 1292 | Args: |
| 1293 | layer: The attention layer, providing the q/k/v quantization scales. |
| 1294 | query: shape = [num_tokens, num_heads, head_size] |
| 1295 | key: shape = [num_tokens, num_kv_heads, head_size] |
| 1296 | value: shape = [num_tokens, num_kv_heads, head_size] |
| 1297 | kv_cache: shape = |
| 1298 | [num_blocks, num_kv_heads, block_size, 2 * head_size] |
| 1299 | attn_metadata: Metadata for attention. |
| 1300 | output: Tensor that the attention result is written into. |
| 1301 | output_scale: Scale for fused output quantization; not supported |
| 1302 | by this backend. |
| 1303 | output_block_scale: Block scale for fused output quantization; |
| 1304 | not supported by this backend. |
| 1305 | |
| 1306 | Returns: |
| 1307 | shape = [num_tokens, num_heads * head_size] |
| 1308 | |
| 1309 | """ |
| 1310 | if output_scale is not None or output_block_scale is not None: |
| 1311 | raise NotImplementedError( |
| 1312 | "fused output quantization is not yet supported for FlexAttentionImpl" |
| 1313 | ) |
| 1314 | |
| 1315 | enable_gqa = self.num_kv_heads != self.num_heads |
| 1316 | |
| 1317 | if attn_metadata is None: |
| 1318 | # Profiling run. |
| 1319 | return output.fill_(0) |
| 1320 | # query = self.view_as_4d(query).permute(0, 2, 1, 3) |
| 1321 | # return torch.empty_like(query) |
| 1322 | |
| 1323 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 1324 | |
| 1325 | needs_rebuild_block_mask = False |
| 1326 | if attn_metadata.sliding_window != self.sliding_window: |
| 1327 | attn_metadata.sliding_window = self.sliding_window |
| 1328 | if attn_metadata.direct_build: |
| 1329 | # update mask mod in attention metadata |
| 1330 | attn_metadata.mask_mod = attn_metadata.get_mask_mod() |
| 1331 | needs_rebuild_block_mask = True |
| 1332 | |
| 1333 | if self.mm_prefix_range != getattr(attn_metadata, "mm_prefix_range", None): |
| 1334 | self.mm_prefix_range = attn_metadata.mm_prefix_range |
| 1335 | attn_metadata.mask_mod = attn_metadata.get_mask_mod() |
| 1336 | needs_rebuild_block_mask = True |
| 1337 | |
| 1338 | layer_mask_mod = getattr(layer, "logical_mask_mod", None) |
| 1339 | if ( |
| 1340 | layer_mask_mod is not None |
| 1341 | and attn_metadata.logical_mask_mod is not layer_mask_mod |
| 1342 | ): |
| 1343 | attn_metadata.logical_mask_mod = layer_mask_mod |
| 1344 | attn_metadata.mask_mod = attn_metadata.get_mask_mod() |
| 1345 | needs_rebuild_block_mask = True |
| 1346 | |
| 1347 | layer_hint = getattr(layer, "block_sparsity_hint", None) |
| 1348 | if ( |
| 1349 | layer_hint is not None |
| 1350 | and attn_metadata.block_sparsity_hint is not layer_hint |
| 1351 | ): |
| 1352 | attn_metadata.block_sparsity_hint = layer_hint |
| 1353 | needs_rebuild_block_mask = True |
| 1354 | |
| 1355 | if needs_rebuild_block_mask or attn_metadata.block_mask is None: |
| 1356 | if attn_metadata.direct_build: |
| 1357 | attn_metadata.block_mask = attn_metadata._build_block_mask_direct() |
| 1358 | else: |
| 1359 | attn_metadata.block_mask = attn_metadata.build_block_mask() |
| 1360 | |
| 1361 | if self.attn_type == AttentionType.ENCODER_ONLY: |
| 1362 | query, key_tensor, value_tensor = map( |
| 1363 | lambda x: self.view_as_4d(x).permute(0, 2, 1, 3), |
| 1364 | (query, key, value), |
| 1365 | ) |
| 1366 | |
| 1367 | query = query[:, :, :num_actual_tokens, :] |
| 1368 | if (key_tensor.size(-2) > num_actual_tokens) or ( |
| 1369 | value_tensor.size(-2) > num_actual_tokens |
| 1370 | ): |
| 1371 | # In the encoder-only model with torch.compile, |
| 1372 | # qkv might be padded, which might cause exception. |
| 1373 | # see: https://github.com/vllm-project/vllm/pull/24872#discussion_r2353252290 |
| 1374 | key_tensor = key_tensor[:, :, :num_actual_tokens, :] |
| 1375 | value_tensor = value_tensor[:, :, :num_actual_tokens, :] |
| 1376 | |
| 1377 | else: |
| 1378 | assert self.attn_type == AttentionType.DECODER |
| 1379 | kv_cache = kv_cache.transpose(1, 2) |
| 1380 | hs = self.head_size |
| 1381 | key_cache, value_cache = kv_cache.split(hs, dim=-1) |
| 1382 | |
| 1383 | # Flatten (num_blocks, block_size) into a single token dim. |
| 1384 | key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size) |
| 1385 | value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size) |
| 1386 | query, key_tensor, value_tensor = map( |
| 1387 | lambda x: self.view_as_4d(x).permute(0, 2, 1, 3), |
| 1388 | (query, key_cache, value_cache), |
| 1389 | ) |
| 1390 | |
| 1391 | query = query[:, :, :num_actual_tokens, :] |
| 1392 | |
| 1393 | # Doesn't work for now -> constraint violation |
| 1394 | # torch._dynamo.try_mark_dynamic(query, 2) |
| 1395 | |
| 1396 | assert attn_metadata.block_mask is not None |
| 1397 | block_m, block_n = attn_metadata.block_mask.BLOCK_SIZE |
| 1398 | |
| 1399 | kernel_options = get_kernel_options( |
| 1400 | query, block_m, block_n, attn_metadata.direct_build |
| 1401 | ) |
| 1402 | |
| 1403 | if self.block_m is not None: |
| 1404 | kernel_options["BLOCK_M"] = self.block_m |
| 1405 | if self.block_n is not None: |
| 1406 | kernel_options["BLOCK_N"] = self.block_n |
| 1407 | if envs.VLLM_BATCH_INVARIANT: |
| 1408 | kernel_options["IS_DIVISIBLE"] = False |
| 1409 | out = flex_attention_compiled( |
| 1410 | query, |
| 1411 | key_tensor, |
| 1412 | value_tensor, |
| 1413 | attn_metadata.transformed_score_mod, |
| 1414 | attn_metadata.block_mask, |
| 1415 | self.scale, |
| 1416 | enable_gqa=enable_gqa, |
| 1417 | kernel_options=kernel_options, |
| 1418 | return_aux=AuxRequest(lse=True) if self.out_transform is not None else None, |
| 1419 | ) |
| 1420 | |
| 1421 | if self.out_transform is not None: |
| 1422 | out, aux = out |
| 1423 | out = self.out_transform(out, aux.lse) |
| 1424 | |
| 1425 | # Flex doesn't have an out variant today, rely on epilogue fusion |
| 1426 | out = out.permute(0, 2, 1, 3).squeeze(0) |
| 1427 | output[:num_actual_tokens, :, :].copy_(out) |
| 1428 | return output |
| 1429 | |
| 1430 | |
| 1431 | def get_kernel_options( |
| 1432 | query, block_m, block_n, use_direct_build: bool |
| 1433 | ) -> dict[str, int | bool]: |
| 1434 | kernel_options: dict[str, int | bool] = { |
| 1435 | "FORCE_USE_FLEX_ATTENTION": True, |
| 1436 | } |
| 1437 | |
| 1438 | def ensure_divisible(candidate: int, block_size: int) -> int: |
| 1439 | """Pick a kernel block size that divides the logical block.""" |
| 1440 | if block_size <= 0: |
| 1441 | return candidate |
| 1442 | candidate = min(candidate, block_size) |
| 1443 | if candidate <= 0: |
| 1444 | return block_size |
| 1445 | if block_size % candidate == 0: |
| 1446 | return candidate |
| 1447 | |
| 1448 | candidate = math.gcd(candidate, block_size) |
| 1449 | if candidate <= 1: |
| 1450 | return block_size |
| 1451 | return candidate |
| 1452 | |
| 1453 | if use_direct_build: |
| 1454 | kernel_options["BLOCK_M"] = block_m |
| 1455 | kernel_options["BLOCK_N"] = block_n |
| 1456 | return kernel_options |
| 1457 | else: |
| 1458 | preferred_block = 32 if query.dtype == torch.float32 else 64 |
| 1459 | block_lower_bound = 16 |
| 1460 | |
| 1461 | block_m_candidate = ensure_divisible(preferred_block, block_m) |
| 1462 | block_n_candidate = ensure_divisible(preferred_block, block_n) |
| 1463 | |
| 1464 | if torch.cuda.is_available(): |
| 1465 | device_props = torch.cuda.get_device_properties() |
| 1466 | # ROCm doesn't expose shared_memory_per_block_optin attribute |
| 1467 | # AMD GPUs typically have 64KB LDS (Local Data Share) per workgroup |
| 1468 | if hasattr(device_props, "shared_memory_per_block_optin"): |
| 1469 | max_shared_memory = device_props.shared_memory_per_block_optin |
| 1470 | elif current_platform.is_rocm(): |
| 1471 | # ROCm fallback: use 64KB |
| 1472 | max_shared_memory = 65536 |
| 1473 | else: |
| 1474 | raise RuntimeError( |
| 1475 | "Unable to determine shared memory size on this hardware." |
| 1476 | ) |
| 1477 | |
| 1478 | if max_shared_memory < 144 * 1024: |
| 1479 | block_m_candidate = ensure_divisible( |
| 1480 | max(1, block_m_candidate // 2), block_m |
| 1481 | ) |
| 1482 | block_n_candidate = ensure_divisible( |
| 1483 | max(1, block_n_candidate // 2), block_n |
| 1484 | ) |
| 1485 | |
| 1486 | block_m_candidate = max(block_m_candidate, block_lower_bound) |
| 1487 | block_n_candidate = max(block_n_candidate, block_lower_bound) |
| 1488 | |
| 1489 | kernel_options["BLOCK_M"] = block_m_candidate |
| 1490 | kernel_options["BLOCK_N"] = block_n_candidate |
| 1491 | |
| 1492 | return kernel_options |