| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Paged side-cache ownership and metadata for Qwen4Exp QSA. |
| 4 | |
| 5 | Each QSA layer keeps a fixed circular buffer of raw index keys (the |
| 6 | compressor state) and one compressed key. MRoPE models pack exact three-axis |
| 7 | positions beside the raw keys; text models derive group positions from |
| 8 | logical positions. The compressor state uses one block per request, while |
| 9 | the compressed owner uses ``MLAAttentionSpec.tokens_per_state`` so its block |
| 10 | table follows the main KV-cache lifecycle. Their physical tensor storage is |
| 11 | shared by the generic cache-layout planner. |
| 12 | """ |
| 13 | |
| 14 | import math |
| 15 | from dataclasses import dataclass |
| 16 | from functools import cache |
| 17 | from typing import ClassVar |
| 18 | |
| 19 | import torch |
| 20 | from torch import nn |
| 21 | |
| 22 | from vllm.config import CacheConfig, VllmConfig |
| 23 | from vllm.config.cache import CacheDType |
| 24 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 25 | from vllm.platforms import current_platform |
| 26 | from vllm.triton_utils import HAS_TRITON, tl, triton |
| 27 | from vllm.utils.math_utils import cdiv |
| 28 | from vllm.v1.attention.backend import ( |
| 29 | AttentionBackend, |
| 30 | AttentionCGSupport, |
| 31 | AttentionMetadata, |
| 32 | AttentionMetadataBuilder, |
| 33 | CommonAttentionMetadata, |
| 34 | ) |
| 35 | from vllm.v1.attention.backends.utils import ( |
| 36 | PAD_SLOT_ID, |
| 37 | split_decodes_and_prefills, |
| 38 | ) |
| 39 | from vllm.v1.kv_cache_interface import ( |
| 40 | AttentionSpec, |
| 41 | CircularBufferSpec, |
| 42 | KVCacheLayout, |
| 43 | KVCacheSpec, |
| 44 | MLAAttentionSpec, |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: |
| 49 | """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" |
| 50 | if positions.ndim == 1: |
| 51 | positions = positions.unsqueeze(0).expand(3, -1) |
| 52 | elif positions.ndim != 2 or positions.shape[0] not in (1, 3): |
| 53 | raise ValueError("QSA RoPE positions must be [tokens] or [1|3, tokens]") |
| 54 | if positions.shape[0] == 1: |
| 55 | positions = positions.expand(3, -1) |
| 56 | return positions.transpose(0, 1).unsqueeze(1).to(torch.int64) |
| 57 | |
| 58 | |
| 59 | def _logical_positions( |
| 60 | query_start_loc: torch.Tensor, |
| 61 | seq_lens: torch.Tensor, |
| 62 | token_to_req: torch.Tensor, |
| 63 | num_tokens: int, |
| 64 | ) -> torch.Tensor: |
| 65 | if num_tokens == 0: |
| 66 | return seq_lens.new_empty((0,), dtype=torch.int64) |
| 67 | arange = torch.arange(num_tokens, device=query_start_loc.device) |
| 68 | requests = token_to_req[:num_tokens].long() |
| 69 | query_lens = torch.diff(query_start_loc) |
| 70 | within_query = arange - query_start_loc.index_select(0, requests) |
| 71 | return ( |
| 72 | seq_lens.index_select(0, requests).long() |
| 73 | - query_lens.index_select(0, requests).long() |
| 74 | + within_query.long() |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def _logical_to_physical_qsa_slots( |
| 79 | block_table: torch.Tensor, |
| 80 | request_indices: torch.Tensor, |
| 81 | logical_positions: torch.Tensor, |
| 82 | block_size: int, |
| 83 | ) -> torch.Tensor: |
| 84 | if block_size <= 0: |
| 85 | raise ValueError("QSA cache block size must be positive") |
| 86 | if block_table.ndim != 2: |
| 87 | raise ValueError("QSA block table must be two-dimensional") |
| 88 | if request_indices.shape != logical_positions.shape: |
| 89 | request_indices = torch.broadcast_to(request_indices, logical_positions.shape) |
| 90 | |
| 91 | requests = request_indices.to(device=block_table.device, dtype=torch.long) |
| 92 | positions = logical_positions.to(device=block_table.device, dtype=torch.long) |
| 93 | valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) |
| 94 | logical_blocks = torch.div( |
| 95 | positions.clamp_min(0), block_size, rounding_mode="floor" |
| 96 | ) |
| 97 | valid &= logical_blocks < block_table.shape[1] |
| 98 | safe_requests = requests.clamp(0, max(block_table.shape[0] - 1, 0)) |
| 99 | safe_blocks = logical_blocks.clamp(0, max(block_table.shape[1] - 1, 0)) |
| 100 | if not all(block_table.shape): |
| 101 | return torch.full_like(positions, PAD_SLOT_ID) |
| 102 | physical_blocks = block_table[safe_requests, safe_blocks].long() |
| 103 | valid &= physical_blocks >= 0 |
| 104 | slots = physical_blocks * block_size + positions.remainder(block_size) |
| 105 | return torch.where(valid, slots, PAD_SLOT_ID) |
| 106 | |
| 107 | |
| 108 | def circular_qsa_slot_mapping( |
| 109 | block_table: torch.Tensor, |
| 110 | token_to_req: torch.Tensor, |
| 111 | logical_positions: torch.Tensor, |
| 112 | compressor_state_size: int, |
| 113 | query_start_loc: torch.Tensor | None = None, |
| 114 | out: torch.Tensor | None = None, |
| 115 | ) -> torch.Tensor: |
| 116 | """Map each request to its fixed physical block as a circular token ring.""" |
| 117 | if compressor_state_size <= 0: |
| 118 | raise ValueError("QSA circular buffer size must be positive") |
| 119 | if block_table.ndim != 2: |
| 120 | raise ValueError("QSA block table must be two-dimensional") |
| 121 | |
| 122 | requests = token_to_req.to(device=block_table.device, dtype=torch.long) |
| 123 | positions = logical_positions.to(device=block_table.device, dtype=torch.long) |
| 124 | if not all(block_table.shape): |
| 125 | slots = torch.full_like(positions, PAD_SLOT_ID) |
| 126 | else: |
| 127 | valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) |
| 128 | safe_requests = requests.clamp(0, block_table.shape[0] - 1) |
| 129 | physical_blocks = block_table[safe_requests, 0].long() |
| 130 | valid &= physical_blocks >= 0 |
| 131 | slots = physical_blocks * compressor_state_size + positions.remainder( |
| 132 | compressor_state_size |
| 133 | ) |
| 134 | slots = torch.where(valid, slots, PAD_SLOT_ID) |
| 135 | |
| 136 | if query_start_loc is not None: |
| 137 | if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: |
| 138 | raise ValueError("QSA query starts must contain a terminal offset") |
| 139 | query_start_loc = query_start_loc.to(block_table.device) |
| 140 | num_requests = query_start_loc.shape[0] - 1 |
| 141 | safe_requests = requests.clamp(0, num_requests - 1) |
| 142 | request_ends = query_start_loc.index_select(0, safe_requests + 1) |
| 143 | rows = torch.arange(slots.numel(), device=slots.device) |
| 144 | keep = ( |
| 145 | (requests >= 0) |
| 146 | & (requests < num_requests) |
| 147 | & (rows + compressor_state_size >= request_ends) |
| 148 | ) |
| 149 | slots = torch.where(keep, slots, PAD_SLOT_ID) |
| 150 | |
| 151 | slots = slots.to(torch.int64) |
| 152 | if out is not None: |
| 153 | out.fill_(PAD_SLOT_ID) |
| 154 | out[: slots.numel()].copy_(slots) |
| 155 | return out[: slots.numel()] |
| 156 | return slots |
| 157 | |
| 158 | |
| 159 | def compressed_qsa_slot_mapping( |
| 160 | block_table: torch.Tensor, |
| 161 | token_to_req: torch.Tensor, |
| 162 | logical_positions: torch.Tensor, |
| 163 | storage_block_size: int, |
| 164 | compress_ratio: int, |
| 165 | out: torch.Tensor | None = None, |
| 166 | ) -> torch.Tensor: |
| 167 | """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache.""" |
| 168 | if storage_block_size <= 0 or compress_ratio <= 0: |
| 169 | raise ValueError("QSA block size and compression ratio must be positive") |
| 170 | compressed_positions = torch.div( |
| 171 | logical_positions.clamp_min(0), compress_ratio, rounding_mode="floor" |
| 172 | ) |
| 173 | slots = _logical_to_physical_qsa_slots( |
| 174 | block_table, |
| 175 | token_to_req, |
| 176 | compressed_positions, |
| 177 | storage_block_size, |
| 178 | ) |
| 179 | valid = (logical_positions >= 0) & ( |
| 180 | (logical_positions + 1).remainder(compress_ratio) == 0 |
| 181 | ) |
| 182 | slots = torch.where(valid, slots, PAD_SLOT_ID).to(torch.int64) |
| 183 | if out is not None: |
| 184 | out.fill_(PAD_SLOT_ID) |
| 185 | out[: slots.numel()].copy_(slots) |
| 186 | return out[: slots.numel()] |
| 187 | return slots |
| 188 | |
| 189 | |
| 190 | @cache |
| 191 | def _metadata_launch_pdl() -> bool: |
| 192 | return current_platform.is_arch_support_pdl() |
| 193 | |
| 194 | |
| 195 | @triton.jit( |
| 196 | do_not_specialize=[ |
| 197 | "num_reqs", |
| 198 | "num_mapped_tokens", |
| 199 | "num_tokens", |
| 200 | "max_num_work", |
| 201 | "num_search_steps", |
| 202 | "work_search_steps", |
| 203 | ] |
| 204 | ) |
| 205 | def _build_qsa_metadata_kernel( |
| 206 | query_start_loc_ptr, |
| 207 | seq_lens_ptr, |
| 208 | common_slot_mapping_ptr, |
| 209 | block_table_ptr, |
| 210 | token_to_req_ptr, |
| 211 | logical_positions_ptr, |
| 212 | visible_blocks_ptr, |
| 213 | slot_mapping_ptr, |
| 214 | k_work_metadata_ptr, |
| 215 | block_table_stride_0: tl.constexpr, |
| 216 | block_table_stride_1: tl.constexpr, |
| 217 | num_reqs, |
| 218 | num_mapped_tokens, |
| 219 | num_tokens, |
| 220 | max_num_work, |
| 221 | num_search_steps, |
| 222 | work_search_steps, |
| 223 | storage_block_size: tl.constexpr, |
| 224 | compress_ratio: tl.constexpr, |
| 225 | circular_buffer_size: tl.constexpr, |
| 226 | num_block_table_columns: tl.constexpr, |
| 227 | launch_pdl: tl.constexpr, |
| 228 | TOKEN_BLOCK_SIZE: tl.constexpr, |
| 229 | REQUEST_SCAN_SIZE: tl.constexpr, |
| 230 | WORK_BLOCK_SIZE: tl.constexpr, |
| 231 | ): |
| 232 | if launch_pdl: |
| 233 | tl.extra.cuda.gdc_wait() |
| 234 | |
| 235 | pid = tl.program_id(0) |
| 236 | token_idx = pid * TOKEN_BLOCK_SIZE + tl.arange(0, TOKEN_BLOCK_SIZE) |
| 237 | store_mask = token_idx < num_tokens |
| 238 | mapped = token_idx < num_mapped_tokens |
| 239 | search_token_idx = tl.minimum(token_idx, num_mapped_tokens - 1) |
| 240 | request_idx = tl.zeros((TOKEN_BLOCK_SIZE,), tl.int32) |
| 241 | # Find the last query start at or before each token. The dynamic loop avoids |
| 242 | # compiling a kernel variant for every ceil(log2(num_reqs)). |
| 243 | for step in tl.range(0, num_search_steps): |
| 244 | candidate = request_idx + (1 << (num_search_steps - step - 1)) |
| 245 | valid_candidate = candidate < num_reqs |
| 246 | candidate_start = tl.load( |
| 247 | query_start_loc_ptr + candidate, |
| 248 | mask=valid_candidate, |
| 249 | other=num_mapped_tokens + 1, |
| 250 | ) |
| 251 | advance = valid_candidate & (candidate_start <= search_token_idx) |
| 252 | request_idx = tl.where(advance, candidate, request_idx) |
| 253 | query_start = tl.load(query_start_loc_ptr + request_idx, mask=mapped, other=0) |
| 254 | query_end = tl.load(query_start_loc_ptr + request_idx + 1, mask=mapped, other=0) |
| 255 | seq_len = tl.load(seq_lens_ptr + request_idx, mask=mapped, other=0) |
| 256 | logical_position = seq_len - (query_end - query_start) + token_idx - query_start |
| 257 | logical_position = tl.where(mapped, logical_position, -1) |
| 258 | tl.store( |
| 259 | token_to_req_ptr + token_idx, |
| 260 | tl.where(mapped, request_idx, 0), |
| 261 | mask=store_mask, |
| 262 | ) |
| 263 | tl.store( |
| 264 | logical_positions_ptr + token_idx, |
| 265 | logical_position, |
| 266 | mask=store_mask, |
| 267 | ) |
| 268 | visible_blocks = tl.maximum( |
| 269 | 0, |
| 270 | tl.minimum( |
| 271 | (logical_position + 1) // compress_ratio, |
| 272 | seq_len // compress_ratio, |
| 273 | ), |
| 274 | ).to(tl.int32) |
| 275 | tl.store(visible_blocks_ptr + token_idx, visible_blocks, mask=store_mask) |
| 276 | |
| 277 | # circular_buffer_size is constexpr, so each builder instance compiles out |
| 278 | # the other QSA cache owner's slot-mapping rule. |
| 279 | if circular_buffer_size > 0: |
| 280 | valid = ( |
| 281 | mapped |
| 282 | & (logical_position >= 0) |
| 283 | & (token_idx + circular_buffer_size >= query_end) |
| 284 | & (num_block_table_columns > 0) |
| 285 | ) |
| 286 | physical_block = tl.load( |
| 287 | block_table_ptr + request_idx * block_table_stride_0, |
| 288 | mask=valid, |
| 289 | other=-1, |
| 290 | ) |
| 291 | valid &= physical_block >= 0 |
| 292 | slot = physical_block * circular_buffer_size + ( |
| 293 | logical_position % circular_buffer_size |
| 294 | ) |
| 295 | elif compress_ratio != 1: |
| 296 | compressed_position = tl.maximum(logical_position, 0) // compress_ratio |
| 297 | logical_block = compressed_position // storage_block_size |
| 298 | valid = ( |
| 299 | mapped |
| 300 | & (logical_position >= 0) |
| 301 | & ((logical_position + 1) % compress_ratio == 0) |
| 302 | & (logical_block < num_block_table_columns) |
| 303 | ) |
| 304 | physical_block = tl.load( |
| 305 | block_table_ptr |
| 306 | + request_idx * block_table_stride_0 |
| 307 | + logical_block * block_table_stride_1, |
| 308 | mask=valid, |
| 309 | other=-1, |
| 310 | ) |
| 311 | valid &= physical_block >= 0 |
| 312 | valid &= ( |
| 313 | tl.load(common_slot_mapping_ptr + token_idx, mask=mapped, other=-1) >= 0 |
| 314 | ) |
| 315 | slot = physical_block * storage_block_size + ( |
| 316 | compressed_position % storage_block_size |
| 317 | ) |
| 318 | if (circular_buffer_size > 0) or (compress_ratio != 1): |
| 319 | tl.store( |
| 320 | slot_mapping_ptr + token_idx, |
| 321 | tl.where(valid, slot, -1), |
| 322 | mask=store_mask, |
| 323 | ) |
| 324 | work_tile_start = pid * WORK_BLOCK_SIZE |
| 325 | has_work_tile = work_tile_start < max_num_work |
| 326 | if k_work_metadata_ptr is not None and has_work_tile: |
| 327 | # Every work CTA builds the request prefix in registers. Recomputing this |
| 328 | # small vector lets CTAs write disjoint work tiles without a grid barrier. |
| 329 | requests = tl.arange(0, REQUEST_SCAN_SIZE) |
| 330 | valid_request = requests < num_reqs |
| 331 | request_query_start = tl.load( |
| 332 | query_start_loc_ptr + requests, mask=valid_request, other=0 |
| 333 | ) |
| 334 | request_query_end = tl.load( |
| 335 | query_start_loc_ptr + requests + 1, mask=valid_request, other=0 |
| 336 | ) |
| 337 | request_seq_len = tl.load(seq_lens_ptr + requests, mask=valid_request, other=0) |
| 338 | request_query_len = request_query_end - request_query_start |
| 339 | chunk_start = request_seq_len - request_query_len |
| 340 | num_groups = request_seq_len // compress_ratio - chunk_start // compress_ratio |
| 341 | # Nonempty requests need one item even without a completed compression |
| 342 | # group because work item zero also commits the current raw-K suffix. |
| 343 | work_counts = tl.where(request_query_len > 0, tl.maximum(num_groups, 1), 0) |
| 344 | work_ends = tl.cumsum(tl.where(valid_request, work_counts, 0), axis=0) |
| 345 | total_work = tl.sum(tl.where(valid_request, work_counts, 0), axis=0) |
| 346 | work_offsets = tl.arange(0, WORK_BLOCK_SIZE) |
| 347 | work = work_tile_start + work_offsets |
| 348 | in_bounds = work < max_num_work |
| 349 | active = in_bounds & (work < total_work) |
| 350 | request = tl.zeros((WORK_BLOCK_SIZE,), dtype=tl.int32) |
| 351 | # Request zero starts at zero for every active work item. Descending |
| 352 | # steps find the last request starting at or before this item. |
| 353 | for step_idx in tl.range(0, work_search_steps): |
| 354 | step = 1 << (work_search_steps - step_idx - 1) |
| 355 | candidate = request + step |
| 356 | valid_candidate = candidate < num_reqs |
| 357 | candidate_start = tl.gather(work_ends, candidate - 1, 0) |
| 358 | advance = active & valid_candidate & (candidate_start <= work) |
| 359 | request = tl.where(advance, candidate, request) |
| 360 | |
| 361 | if launch_pdl: |
| 362 | # Let the dependent grid start launch setup while these CTAs finish |
| 363 | # stores; its gdc_wait still orders access to the completed metadata. |
| 364 | tl.extra.cuda.gdc_launch_dependents() |
| 365 | |
| 366 | owner_work_start = tl.gather(work_ends, tl.maximum(request - 1, 0), 0) |
| 367 | owner_work_start = tl.where(request == 0, 0, owner_work_start) |
| 368 | work_in_request = work - owner_work_start |
| 369 | tl.store( |
| 370 | k_work_metadata_ptr + work * 2, |
| 371 | tl.where(active, request, -1), |
| 372 | mask=in_bounds, |
| 373 | ) |
| 374 | tl.store( |
| 375 | k_work_metadata_ptr + work * 2 + 1, |
| 376 | tl.where(active, work_in_request, -1), |
| 377 | mask=in_bounds, |
| 378 | ) |
| 379 | |
| 380 | else: |
| 381 | if launch_pdl: |
| 382 | # A dependent grid launches only after every CTA has signaled. |
| 383 | tl.extra.cuda.gdc_launch_dependents() |
| 384 | |
| 385 | |
| 386 | def build_qsa_metadata_triton( |
| 387 | common_attn_metadata: CommonAttentionMetadata, |
| 388 | token_to_req_buffer: torch.Tensor, |
| 389 | logical_positions_buffer: torch.Tensor, |
| 390 | visible_blocks_buffer: torch.Tensor, |
| 391 | slot_mapping_buffer: torch.Tensor, |
| 392 | *, |
| 393 | storage_block_size: int, |
| 394 | compress_ratio: int, |
| 395 | circular_buffer_size: int = 0, |
| 396 | k_work_metadata_buffer: torch.Tensor | None = None, |
| 397 | request_capacity: int | None = None, |
| 398 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 399 | """Build QSA side-cache and optional pre-indexer work metadata.""" |
| 400 | num_tokens = common_attn_metadata.num_actual_tokens |
| 401 | num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) |
| 402 | token_to_req = token_to_req_buffer[:num_tokens] |
| 403 | logical_positions = logical_positions_buffer[:num_tokens] |
| 404 | visible_blocks = visible_blocks_buffer[:num_tokens] |
| 405 | slot_mapping = slot_mapping_buffer[:num_tokens] |
| 406 | num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1 |
| 407 | assert num_reqs > 0 |
| 408 | |
| 409 | if k_work_metadata_buffer is not None: |
| 410 | if request_capacity is None: |
| 411 | request_capacity = num_reqs |
| 412 | assert request_capacity >= num_reqs |
| 413 | # Pad for tl.arange while keeping the scan width stable across live batches. |
| 414 | request_scan_size = 1 << int(math.ceil(math.log2(request_capacity))) |
| 415 | max_num_work = k_work_metadata_buffer.shape[0] |
| 416 | else: |
| 417 | request_scan_size = 1 |
| 418 | max_num_work = 0 |
| 419 | |
| 420 | if num_tokens == 0 and k_work_metadata_buffer is None: |
| 421 | return token_to_req, logical_positions, visible_blocks, slot_mapping |
| 422 | |
| 423 | block_table = common_attn_metadata.block_table_tensor |
| 424 | num_search_steps = int(math.ceil(math.log2(num_reqs))) |
| 425 | work_search_steps = int(math.ceil(math.log2(num_reqs))) |
| 426 | # The same grid covers token tiles and, for the compressed cache, work tiles. |
| 427 | num_token_blocks = cdiv(num_tokens, 128) |
| 428 | num_work_blocks = ( |
| 429 | cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0 |
| 430 | ) |
| 431 | _build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)]( |
| 432 | common_attn_metadata.query_start_loc, |
| 433 | common_attn_metadata.seq_lens, |
| 434 | common_attn_metadata.slot_mapping, |
| 435 | block_table, |
| 436 | token_to_req, |
| 437 | logical_positions, |
| 438 | visible_blocks, |
| 439 | slot_mapping, |
| 440 | k_work_metadata_buffer, |
| 441 | block_table.stride(0), |
| 442 | block_table.stride(1), |
| 443 | num_reqs, |
| 444 | num_mapped_tokens, |
| 445 | num_tokens, |
| 446 | max_num_work, |
| 447 | num_search_steps, |
| 448 | work_search_steps, |
| 449 | storage_block_size, |
| 450 | compress_ratio, |
| 451 | circular_buffer_size, |
| 452 | block_table.shape[1], |
| 453 | launch_pdl=_metadata_launch_pdl(), |
| 454 | TOKEN_BLOCK_SIZE=128, |
| 455 | REQUEST_SCAN_SIZE=request_scan_size, |
| 456 | WORK_BLOCK_SIZE=256, |
| 457 | num_warps=4, |
| 458 | ) |
| 459 | if circular_buffer_size == 0 and compress_ratio == 1: |
| 460 | slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] |
| 461 | return token_to_req, logical_positions, visible_blocks, slot_mapping |
| 462 | |
| 463 | |
| 464 | def _build_qsa_metadata_torch( |
| 465 | common_attn_metadata: CommonAttentionMetadata, |
| 466 | token_to_req_buffer: torch.Tensor, |
| 467 | logical_positions_buffer: torch.Tensor, |
| 468 | visible_blocks_buffer: torch.Tensor, |
| 469 | slot_mapping_buffer: torch.Tensor, |
| 470 | *, |
| 471 | storage_block_size: int, |
| 472 | compress_ratio: int, |
| 473 | circular_buffer_size: int = 0, |
| 474 | k_work_metadata_buffer: torch.Tensor | None = None, |
| 475 | request_capacity: int | None = None, |
| 476 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 477 | del request_capacity |
| 478 | num_tokens = common_attn_metadata.num_actual_tokens |
| 479 | num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) |
| 480 | logical_positions = logical_positions_buffer[:num_tokens] |
| 481 | visible_blocks = visible_blocks_buffer[:num_tokens] |
| 482 | |
| 483 | token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[ |
| 484 | :num_tokens |
| 485 | ] |
| 486 | logical_positions[:num_mapped_tokens].copy_( |
| 487 | _logical_positions( |
| 488 | common_attn_metadata.query_start_loc, |
| 489 | common_attn_metadata.seq_lens, |
| 490 | token_to_req[:num_mapped_tokens], |
| 491 | num_mapped_tokens, |
| 492 | ) |
| 493 | ) |
| 494 | if num_mapped_tokens < num_tokens: |
| 495 | logical_positions[num_mapped_tokens:].fill_(-1) |
| 496 | row_seq_lens = common_attn_metadata.seq_lens.index_select(0, token_to_req.long()) |
| 497 | visible_blocks.copy_( |
| 498 | torch.minimum( |
| 499 | (logical_positions + 1) // compress_ratio, |
| 500 | row_seq_lens // compress_ratio, |
| 501 | ) |
| 502 | .clamp_min_(0) |
| 503 | .to(torch.int32) |
| 504 | ) |
| 505 | if circular_buffer_size > 0: |
| 506 | slot_mapping = circular_qsa_slot_mapping( |
| 507 | common_attn_metadata.block_table_tensor, |
| 508 | token_to_req, |
| 509 | logical_positions, |
| 510 | circular_buffer_size, |
| 511 | query_start_loc=common_attn_metadata.query_start_loc, |
| 512 | out=slot_mapping_buffer, |
| 513 | ) |
| 514 | elif compress_ratio == 1: |
| 515 | slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] |
| 516 | else: |
| 517 | slot_mapping = compressed_qsa_slot_mapping( |
| 518 | common_attn_metadata.block_table_tensor, |
| 519 | token_to_req, |
| 520 | logical_positions, |
| 521 | storage_block_size, |
| 522 | compress_ratio, |
| 523 | slot_mapping_buffer, |
| 524 | ) |
| 525 | slot_mapping.masked_fill_( |
| 526 | common_attn_metadata.slot_mapping[:num_tokens] < 0, -1 |
| 527 | ) |
| 528 | if k_work_metadata_buffer is not None: |
| 529 | query_lens = ( |
| 530 | common_attn_metadata.query_start_loc[1:] |
| 531 | - common_attn_metadata.query_start_loc[:-1] |
| 532 | ) |
| 533 | chunk_starts = common_attn_metadata.seq_lens - query_lens |
| 534 | num_work_per_request = ( |
| 535 | common_attn_metadata.seq_lens // compress_ratio |
| 536 | - chunk_starts // compress_ratio |
| 537 | ) |
| 538 | num_work_per_request = torch.where( |
| 539 | query_lens > 0, num_work_per_request.clamp_min(1), 0 |
| 540 | ) |
| 541 | k_start_loc = torch.empty( |
| 542 | query_lens.shape[0] + 1, |
| 543 | dtype=torch.int32, |
| 544 | device=query_lens.device, |
| 545 | ) |
| 546 | k_start_loc[0] = 0 |
| 547 | torch.cumsum(num_work_per_request, 0, out=k_start_loc[1:]) |
| 548 | work = torch.arange( |
| 549 | k_work_metadata_buffer.shape[0], |
| 550 | device=k_work_metadata_buffer.device, |
| 551 | ) |
| 552 | requests = torch.searchsorted(k_start_loc[1:], work, right=True) |
| 553 | active = work < k_start_loc[-1] |
| 554 | work_in_request = ( |
| 555 | work - k_start_loc[requests.clamp_max(query_lens.shape[0] - 1)] |
| 556 | ) |
| 557 | k_work_metadata_buffer[:, 0].copy_( |
| 558 | torch.where(active, requests, -1).to(torch.int32) |
| 559 | ) |
| 560 | k_work_metadata_buffer[:, 1].copy_( |
| 561 | torch.where(active, work_in_request, -1).to(torch.int32) |
| 562 | ) |
| 563 | return token_to_req, logical_positions, visible_blocks, slot_mapping |
| 564 | |
| 565 | |
| 566 | # Resolve the fallback outside the per-step metadata hot path. |
| 567 | build_qsa_metadata = ( |
| 568 | build_qsa_metadata_triton if HAS_TRITON else _build_qsa_metadata_torch |
| 569 | ) |
| 570 | |
| 571 | |
| 572 | @dataclass |
| 573 | class QSAForwardMetadata(AttentionMetadata): |
| 574 | """Common per-forward metadata for one QSA side cache.""" |
| 575 | |
| 576 | block_table: torch.Tensor |
| 577 | slot_mapping: torch.Tensor |
| 578 | seq_lens: torch.Tensor |
| 579 | query_start_loc: torch.Tensor |
| 580 | token_to_req: torch.Tensor |
| 581 | logical_positions: torch.Tensor |
| 582 | visible_blocks: torch.Tensor |
| 583 | k_work_metadata: torch.Tensor |
| 584 | num_actual_tokens: int |
| 585 | num_decodes: int |
| 586 | num_decode_tokens: int |
| 587 | num_prefills: int |
| 588 | num_prefill_tokens: int |
| 589 | max_query_len: int |
| 590 | decode_query_len: int |
| 591 | max_seq_len: int |
| 592 | storage_block_size: int |
| 593 | compress_ratio: int |
| 594 | |
| 595 | |
| 596 | class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]): |
| 597 | """Build QSA metadata from vLLM's cache-group-specific common metadata.""" |
| 598 | |
| 599 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 600 | |
| 601 | def __init__( |
| 602 | self, |
| 603 | kv_cache_spec: AttentionSpec, |
| 604 | layer_names: list[str], |
| 605 | vllm_config: VllmConfig, |
| 606 | device: torch.device, |
| 607 | ) -> None: |
| 608 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 609 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 610 | assert self.reorder_batch_threshold is not None |
| 611 | self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec) |
| 612 | if isinstance(kv_cache_spec, MLAAttentionSpec): |
| 613 | compress_ratio = kv_cache_spec.tokens_per_state |
| 614 | assert isinstance(compress_ratio, int), ( |
| 615 | "QSA compression requires an integer tokens_per_state" |
| 616 | ) |
| 617 | self.compress_ratio = compress_ratio |
| 618 | else: |
| 619 | self.compress_ratio = 1 |
| 620 | self.storage_block_size = kv_cache_spec.num_states |
| 621 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 622 | self.token_to_req_buffer = torch.empty( |
| 623 | max_tokens, dtype=torch.int32, device=device |
| 624 | ) |
| 625 | self.slot_mapping_buffer = torch.empty( |
| 626 | max_tokens, dtype=torch.int64, device=device |
| 627 | ) |
| 628 | self.logical_positions_buffer = torch.empty( |
| 629 | max_tokens, dtype=torch.int64, device=device |
| 630 | ) |
| 631 | self.visible_blocks_buffer = torch.empty( |
| 632 | max_tokens, dtype=torch.int32, device=device |
| 633 | ) |
| 634 | max_requests = vllm_config.scheduler_config.max_num_seqs |
| 635 | self.request_capacity = max_requests |
| 636 | if not self.is_circular_buffer and self.compress_ratio != 1: |
| 637 | max_k_work = ( |
| 638 | max_tokens + (self.compress_ratio - 1) * max_requests |
| 639 | ) // self.compress_ratio |
| 640 | self.k_work_metadata_buffer = torch.empty( |
| 641 | max_k_work, 2, dtype=torch.int32, device=device |
| 642 | ) |
| 643 | else: |
| 644 | self.k_work_metadata_buffer = torch.empty( |
| 645 | 0, 2, dtype=torch.int32, device=device |
| 646 | ) |
| 647 | |
| 648 | def build( |
| 649 | self, |
| 650 | common_prefix_len: int, |
| 651 | common_attn_metadata: CommonAttentionMetadata, |
| 652 | fast_build: bool = False, |
| 653 | ) -> QSAForwardMetadata: |
| 654 | del common_prefix_len, fast_build |
| 655 | num_tokens = common_attn_metadata.num_actual_tokens |
| 656 | decode_threshold = self.reorder_batch_threshold |
| 657 | assert decode_threshold is not None |
| 658 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 659 | split_decodes_and_prefills( |
| 660 | common_attn_metadata, |
| 661 | decode_threshold=decode_threshold, |
| 662 | require_uniform=True, |
| 663 | ) |
| 664 | ) |
| 665 | assert num_decodes + num_prefills == common_attn_metadata.num_reqs |
| 666 | assert num_decode_tokens + num_prefill_tokens == num_tokens |
| 667 | decode_query_len = 0 |
| 668 | if num_decodes > 0: |
| 669 | query_lens_cpu = torch.diff( |
| 670 | common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] |
| 671 | ) |
| 672 | nonzero_query_lens = query_lens_cpu[query_lens_cpu > 0] |
| 673 | if nonzero_query_lens.numel() > 0: |
| 674 | decode_query_len = int(nonzero_query_lens[0].item()) |
| 675 | assert torch.all(nonzero_query_lens == decode_query_len) |
| 676 | build_k_work = not self.is_circular_buffer and self.compress_ratio != 1 |
| 677 | k_work_metadata = self.k_work_metadata_buffer |
| 678 | request_capacity = None |
| 679 | if build_k_work: |
| 680 | num_requests = common_attn_metadata.query_start_loc.shape[0] - 1 |
| 681 | request_capacity = self.request_capacity |
| 682 | max_num_work = ( |
| 683 | num_tokens + (self.compress_ratio - 1) * num_requests |
| 684 | ) // self.compress_ratio |
| 685 | k_work_metadata = self.k_work_metadata_buffer[:max_num_work] |
| 686 | token_to_req, logical_positions, visible_blocks, slot_mapping = ( |
| 687 | build_qsa_metadata( |
| 688 | common_attn_metadata, |
| 689 | self.token_to_req_buffer, |
| 690 | self.logical_positions_buffer, |
| 691 | self.visible_blocks_buffer, |
| 692 | self.slot_mapping_buffer, |
| 693 | storage_block_size=self.storage_block_size, |
| 694 | compress_ratio=self.compress_ratio, |
| 695 | circular_buffer_size=( |
| 696 | self.kv_cache_spec.block_size if self.is_circular_buffer else 0 |
| 697 | ), |
| 698 | k_work_metadata_buffer=k_work_metadata if build_k_work else None, |
| 699 | request_capacity=request_capacity, |
| 700 | ) |
| 701 | ) |
| 702 | return QSAForwardMetadata( |
| 703 | block_table=common_attn_metadata.block_table_tensor, |
| 704 | slot_mapping=slot_mapping, |
| 705 | seq_lens=common_attn_metadata.seq_lens, |
| 706 | query_start_loc=common_attn_metadata.query_start_loc, |
| 707 | token_to_req=token_to_req, |
| 708 | logical_positions=logical_positions, |
| 709 | visible_blocks=visible_blocks, |
| 710 | k_work_metadata=k_work_metadata, |
| 711 | num_actual_tokens=num_tokens, |
| 712 | num_decodes=num_decodes, |
| 713 | num_decode_tokens=num_decode_tokens, |
| 714 | num_prefills=num_prefills, |
| 715 | num_prefill_tokens=num_prefill_tokens, |
| 716 | max_query_len=common_attn_metadata.max_query_len, |
| 717 | decode_query_len=decode_query_len, |
| 718 | max_seq_len=common_attn_metadata.max_seq_len, |
| 719 | storage_block_size=self.storage_block_size, |
| 720 | compress_ratio=self.compress_ratio, |
| 721 | ) |
| 722 | |
| 723 | |
| 724 | class QSAStateBackend(AttentionBackend): |
| 725 | """Key-only dummy backend for out-of-band QSA side-cache operations.""" |
| 726 | |
| 727 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] |
| 728 | # fp8 entries allow the optional e4m3 compressed indexer cache. |
| 729 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 730 | "auto", |
| 731 | "bfloat16", |
| 732 | "fp8", |
| 733 | "fp8_e4m3", |
| 734 | ] |
| 735 | |
| 736 | @staticmethod |
| 737 | def get_name() -> str: |
| 738 | return "QWEN4_EXP_EXP_QSA_STATE" |
| 739 | |
| 740 | @staticmethod |
| 741 | def get_impl_cls(): |
| 742 | raise NotImplementedError( |
| 743 | "QSA state caches run out-of-band and have no attention impl" |
| 744 | ) |
| 745 | |
| 746 | @staticmethod |
| 747 | def get_builder_cls() -> type[QSAMetadataBuilder]: |
| 748 | return QSAMetadataBuilder |
| 749 | |
| 750 | @classmethod |
| 751 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 752 | # QSA pages are packed beside the main KV pages within each block. |
| 753 | return (KVCacheLayout.BLNHC, KVCacheLayout.BLHNC) |
| 754 | |
| 755 | |
| 756 | class _QSAStateCache(nn.Module, AttentionLayerBase): |
| 757 | supports_dcp = False |
| 758 | |
| 759 | def __init__( |
| 760 | self, |
| 761 | *, |
| 762 | head_size: int, |
| 763 | dtype: torch.dtype, |
| 764 | cache_config: CacheConfig, |
| 765 | prefix: str, |
| 766 | vllm_config: VllmConfig, |
| 767 | compress_ratio: int = 1, |
| 768 | ) -> None: |
| 769 | super().__init__() |
| 770 | if head_size <= 0: |
| 771 | raise ValueError("QSA cache head size must be positive") |
| 772 | if compress_ratio <= 0: |
| 773 | raise ValueError("QSA compression ratio must be positive") |
| 774 | if cache_config.block_size % compress_ratio: |
| 775 | raise ValueError( |
| 776 | "QSA cache block size must be divisible by the compression ratio" |
| 777 | ) |
| 778 | self.head_size = head_size |
| 779 | self.dtype = dtype |
| 780 | self.cache_config = cache_config |
| 781 | self.prefix = prefix |
| 782 | self.compress_ratio = compress_ratio |
| 783 | self.kv_cache = torch.tensor([]) |
| 784 | |
| 785 | static_context = vllm_config.compilation_config.static_forward_context |
| 786 | if prefix in static_context: |
| 787 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 788 | static_context[prefix] = self |
| 789 | |
| 790 | def forward(self) -> None: ... |
| 791 | |
| 792 | def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: |
| 793 | """Adapt the unified [B, H, N, C] view to QSA's [B, N, H, C].""" |
| 794 | if kv_cache.ndim != 4 or kv_cache.shape[1] != 1: |
| 795 | raise ValueError("QSA state cache must be [blocks, 1, states, width]") |
| 796 | if kv_cache.dtype != self.dtype or kv_cache.shape[3] != self.head_size: |
| 797 | raise ValueError( |
| 798 | f"QSA state cache does not match its spec " |
| 799 | f"(dtype {kv_cache.dtype} != {self.dtype} or width " |
| 800 | f"{kv_cache.shape[3]} != {self.head_size})" |
| 801 | ) |
| 802 | super().bind_kv_cache(kv_cache.transpose(1, 2)) |
| 803 | |
| 804 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 805 | return QSAStateBackend |
| 806 | |
| 807 | |
| 808 | class QSAKeyStateCache(_QSAStateCache): |
| 809 | """Raw BF16 key, optionally followed by exact int64 MRoPE positions.""" |
| 810 | |
| 811 | _BF16_PER_INT64 = 4 |
| 812 | _NUM_ROPE_AXES = 3 |
| 813 | |
| 814 | def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None: |
| 815 | key_head_size = int(kwargs.pop("head_size")) |
| 816 | self.key_head_size = key_head_size |
| 817 | self.cache_rope_positions = bool(cache_rope_positions) |
| 818 | self.rope_position_offset = ( |
| 819 | (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 |
| 820 | ) * self._BF16_PER_INT64 |
| 821 | storage_head_size = key_head_size |
| 822 | if self.cache_rope_positions: |
| 823 | storage_head_size = self.rope_position_offset + ( |
| 824 | self._NUM_ROPE_AXES * self._BF16_PER_INT64 |
| 825 | ) |
| 826 | super().__init__(head_size=storage_head_size, **kwargs) |
| 827 | |
| 828 | # Derived on access so `kv_cache` stays the only reference to the bound |
| 829 | # storage: clearing it (e.g. after CUDA graph memory profiling) must free |
| 830 | # the cache, or the freed block stays pinned in the allocator. |
| 831 | @property |
| 832 | def key_cache(self) -> torch.Tensor: |
| 833 | return self.kv_cache[..., : self.key_head_size] |
| 834 | |
| 835 | @property |
| 836 | def rope_position_cache(self) -> torch.Tensor | None: |
| 837 | if not self.cache_rope_positions: |
| 838 | return None |
| 839 | return self.kv_cache[..., self.rope_position_offset :].view(torch.int64) |
| 840 | |
| 841 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 842 | # Hold the open group's committed keys plus every row a speculative |
| 843 | # step stores before acceptance is known, rounded up to whole groups so |
| 844 | # the ring divides the attention block size (it joins the LCM that sets |
| 845 | # the scheduler block size). Anything narrower lets a rejected draft row |
| 846 | # overwrite a committed key the next step needs to close the group. |
| 847 | span = self.compress_ratio + vllm_config.num_speculative_tokens |
| 848 | capacity = self.compress_ratio * cdiv(span, self.compress_ratio) |
| 849 | assert self.cache_config.block_size % capacity == 0, ( |
| 850 | f"QSA ring capacity {capacity} must divide the attention block " |
| 851 | f"size {self.cache_config.block_size}" |
| 852 | ) |
| 853 | return CircularBufferSpec( |
| 854 | block_size=capacity, |
| 855 | num_kv_heads=1, |
| 856 | head_size=self.head_size, |
| 857 | head_size_v=0, |
| 858 | dtype=self.dtype, |
| 859 | ) |
| 860 | |
| 861 | |
| 862 | class QSACompressedKeyCache(_QSAStateCache): |
| 863 | """Normed, group-first-RoPE key at one row per complete group.""" |
| 864 | |
| 865 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 866 | del vllm_config |
| 867 | return MLAAttentionSpec( |
| 868 | block_size=self.cache_config.block_size, |
| 869 | num_kv_heads=1, |
| 870 | head_size=self.head_size, |
| 871 | dtype=self.dtype, |
| 872 | tokens_per_state=self.compress_ratio, |
| 873 | ) |
| 874 | |
| 875 | |
| 876 | __all__ = [ |
| 877 | "QSACompressedKeyCache", |
| 878 | "QSAForwardMetadata", |
| 879 | "QSAKeyStateCache", |
| 880 | "QSAMetadataBuilder", |
| 881 | "QSAStateBackend", |
| 882 | "canonical_qsa_rope_positions", |
| 883 | "circular_qsa_slot_mapping", |
| 884 | "compressed_qsa_slot_mapping", |
| 885 | ] |