| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from collections import defaultdict |
| 5 | from collections.abc import Callable, Iterable, Sequence |
| 6 | from dataclasses import dataclass, field, fields |
| 7 | from typing import ( |
| 8 | TYPE_CHECKING, |
| 9 | Any, |
| 10 | Protocol, |
| 11 | TypeVar, |
| 12 | cast, |
| 13 | ) |
| 14 | |
| 15 | import numpy as np |
| 16 | import torch |
| 17 | from typing_extensions import runtime_checkable |
| 18 | |
| 19 | from vllm.config import CacheConfig, VllmConfig, get_layers_from_vllm_config |
| 20 | from vllm.config.cache import _layout_from_name |
| 21 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 22 | from vllm.utils.math_utils import cdiv |
| 23 | from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor |
| 24 | from vllm.v1.kv_cache_interface import ( |
| 25 | AttentionSpec, |
| 26 | KVCacheLayout, |
| 27 | KVCacheSpec, |
| 28 | MambaSpec, |
| 29 | ) |
| 30 | |
| 31 | if TYPE_CHECKING: |
| 32 | from vllm.v1.core.sched.output import SchedulerOutput |
| 33 | from vllm.v1.worker.gpu_input_batch import InputBatch |
| 34 | |
| 35 | import vllm.envs as envs |
| 36 | from vllm.distributed.kv_transfer.kv_connector.utils import ( |
| 37 | get_kv_connector_cache_layout, |
| 38 | ) |
| 39 | from vllm.logger import init_logger |
| 40 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 41 | from vllm.v1.attention.backend import ( |
| 42 | AttentionBackend, |
| 43 | AttentionImpl, |
| 44 | AttentionMetadata, |
| 45 | CommonAttentionMetadata, |
| 46 | subclass_attention_backend, |
| 47 | ) |
| 48 | |
| 49 | logger = init_logger(__name__) |
| 50 | |
| 51 | PAD_SLOT_ID = -1 |
| 52 | NULL_BLOCK_ID = 0 |
| 53 | |
| 54 | |
| 55 | ImplT = TypeVar("ImplT", bound=AttentionImpl) |
| 56 | |
| 57 | |
| 58 | def find_attention_impl_variant( |
| 59 | impl: AttentionImpl, impl_cls: type[ImplT] |
| 60 | ) -> ImplT | None: |
| 61 | for variant in impl.get_impl_variants(): |
| 62 | if isinstance(variant, impl_cls): |
| 63 | return cast(ImplT, variant) |
| 64 | return None |
| 65 | |
| 66 | |
| 67 | _LN_2 = math.log(2.0) |
| 68 | |
| 69 | |
| 70 | def log2_lse_to_ln(lse: torch.Tensor) -> torch.Tensor: |
| 71 | """Convert a base-2 log-sum-exp tensor to natural-log units.""" |
| 72 | return lse * _LN_2 |
| 73 | |
| 74 | |
| 75 | def compute_mm_prefix_range_tensor( |
| 76 | mm_prefix_range: dict[int, list[tuple[int, int]]] | None, |
| 77 | num_seqs: int, |
| 78 | device: torch.device, |
| 79 | ) -> torch.Tensor | None: |
| 80 | """Convert mm_prefix_range dict to padded tensor for Triton kernel. |
| 81 | |
| 82 | Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges. |
| 83 | Empty ranges have start==end==0, which kernel skips via is_valid check. |
| 84 | """ |
| 85 | if mm_prefix_range is None: |
| 86 | return None |
| 87 | |
| 88 | range_lists = [ |
| 89 | mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs) |
| 90 | ] |
| 91 | |
| 92 | if all(r == [(0, 0)] for r in range_lists): |
| 93 | return None |
| 94 | |
| 95 | max_ranges = max(len(r) for r in range_lists) |
| 96 | padded = [] |
| 97 | for r in range_lists: |
| 98 | padded_r = list(r) + [(0, 0)] * (max_ranges - len(r)) |
| 99 | padded.append(padded_r) |
| 100 | padded = async_tensor_h2d(padded, dtype=torch.int32, device=device) |
| 101 | return padded.view(num_seqs, max_ranges, 2) |
| 102 | |
| 103 | |
| 104 | def fill_mm_prefix_query_ranges( |
| 105 | out: np.ndarray, |
| 106 | mm_prefix_range: dict[int, list[tuple[int, int]]] | None, |
| 107 | query_start_loc_cpu: torch.Tensor, |
| 108 | seq_lens_cpu: torch.Tensor, |
| 109 | ) -> int: |
| 110 | """Map each scheduled query token to the mm_prefix range containing it. |
| 111 | |
| 112 | Writes into ``out``, a caller-owned ``(max_num_batched_tokens, 2)`` int32 |
| 113 | staging buffer, and returns the number of rows written (0 if no range |
| 114 | covers any scheduled query token, in which case ``out`` is untouched and |
| 115 | the caller should skip the mask_mod entirely). Row ``i`` holds the absolute |
| 116 | ``[start, end]`` bounds of the bidirectional range that query token ``i`` |
| 117 | belongs to, or ``(-1, -1)`` when it is outside every range. |
| 118 | |
| 119 | mm_prefix ranges never overlap, so "query and key share a range" is |
| 120 | equivalent to "the key lies inside the query's own range". The kernel |
| 121 | therefore needs no key-side lookup, and this metadata is sized by scheduled |
| 122 | query tokens rather than by context length -- bounded by |
| 123 | ``max_num_batched_tokens`` instead of ``num_seqs * max_seq_len``. |
| 124 | |
| 125 | Ranges are absolute prompt positions and may extend past the tokens |
| 126 | scheduled so far under chunked prefill; the portion outside the current |
| 127 | chunk is simply not recorded. Degenerate ranges (``start >= end``) are |
| 128 | skipped to match the Triton path's ``start < end`` validity check. |
| 129 | |
| 130 | ``seq_lens_cpu`` only needs to be exact for prefill rows, since mm_prefix |
| 131 | ranges cover prompt tokens: an over-estimate on a decode row shifts that |
| 132 | row's query position further past every range, which still matches nothing. |
| 133 | """ |
| 134 | if mm_prefix_range is None: |
| 135 | return 0 |
| 136 | |
| 137 | query_start_loc = query_start_loc_cpu.numpy() |
| 138 | num_actual_tokens = int(query_start_loc[-1]) |
| 139 | if num_actual_tokens <= 0: |
| 140 | return 0 |
| 141 | assert num_actual_tokens <= out.shape[0], ( |
| 142 | f"mm_prefix staging buffer holds {out.shape[0]} tokens, got {num_actual_tokens}" |
| 143 | ) |
| 144 | |
| 145 | # Resolve every span before touching `out`, so batches whose ranges all |
| 146 | # fall outside the scheduled tokens skip the fill entirely. |
| 147 | spans: list[tuple[int, int, int, int]] = [] |
| 148 | for req_idx, req_ranges in mm_prefix_range.items(): |
| 149 | if not req_ranges: |
| 150 | continue |
| 151 | token_start = int(query_start_loc[req_idx]) |
| 152 | query_len = int(query_start_loc[req_idx + 1]) - token_start |
| 153 | if query_len <= 0: |
| 154 | continue |
| 155 | # Absolute position of this request's first scheduled query token. |
| 156 | context_len = int(seq_lens_cpu[req_idx]) - query_len |
| 157 | for start, end in req_ranges: |
| 158 | if start >= end: |
| 159 | continue |
| 160 | first = max(start - context_len, 0) |
| 161 | last = min(end - context_len, query_len - 1) |
| 162 | if first > last: |
| 163 | continue |
| 164 | spans.append((token_start + first, token_start + last + 1, start, end)) |
| 165 | |
| 166 | if not spans: |
| 167 | return 0 |
| 168 | |
| 169 | out[:num_actual_tokens] = -1 |
| 170 | for row_start, row_end, start, end in spans: |
| 171 | out[row_start:row_end] = (start, end) |
| 172 | return num_actual_tokens |
| 173 | |
| 174 | |
| 175 | _FLASHINFER_LAYOUT_NAMES = { |
| 176 | "LBNHC": "NHD", |
| 177 | "LBHNC": "HND", |
| 178 | "BLHNC": "HND", |
| 179 | "BLNHC": "NHD", |
| 180 | "BHLNC": "HND", |
| 181 | } |
| 182 | |
| 183 | |
| 184 | def get_flashinfer_layout_string(layout: KVCacheLayout) -> str: |
| 185 | """Return the layout name in FlashInfer's convention (NHD/HND).""" |
| 186 | assert layout.name in _FLASHINFER_LAYOUT_NAMES, ( |
| 187 | f"KV cache layout {layout.name} has no FlashInfer equivalent; FlashInfer " |
| 188 | "rejects it in supported_kv_cache_layouts" |
| 189 | ) |
| 190 | return _FLASHINFER_LAYOUT_NAMES[layout.name] |
| 191 | |
| 192 | |
| 193 | # Preference order when no backend declares a supported set; LBNHC (NHD) |
| 194 | # first to match main's default. |
| 195 | _DEFAULT_LAYOUT_PREFERENCE = ( |
| 196 | KVCacheLayout.LBNHC, |
| 197 | KVCacheLayout.LBHNC, |
| 198 | KVCacheLayout.BLNHC, |
| 199 | KVCacheLayout.BLHNC, |
| 200 | KVCacheLayout.BHLNC, |
| 201 | KVCacheLayout.LHBNC, |
| 202 | ) |
| 203 | |
| 204 | |
| 205 | def _layout_names(layouts: Iterable[KVCacheLayout]) -> list[str]: |
| 206 | return [layout.name for layout in layouts] |
| 207 | |
| 208 | |
| 209 | def get_supported_kv_cache_layouts( |
| 210 | backends: Iterable[type[AttentionBackend]], |
| 211 | ) -> list[KVCacheLayout]: |
| 212 | """Layouts every one of the worker's backends supports, most preferred first. |
| 213 | |
| 214 | Every backend declares the layouts its kernels support, most preferred first |
| 215 | (``supported_kv_cache_layouts``), or None when any layout works; workers where |
| 216 | nothing declares follow the default preference. Identical declarations keep |
| 217 | their order; otherwise the layout the most backends put first wins, ties |
| 218 | keeping the enum order. An empty intersection is a hard error. |
| 219 | """ |
| 220 | supported_layouts_lists: list[Sequence[KVCacheLayout]] = [ |
| 221 | layouts |
| 222 | for backend in backends |
| 223 | if (layouts := backend.supported_kv_cache_layouts()) is not None |
| 224 | ] or [_DEFAULT_LAYOUT_PREFERENCE] |
| 225 | |
| 226 | first = supported_layouts_lists[0] |
| 227 | if all(layouts == first for layouts in supported_layouts_lists[1:]): |
| 228 | return list(first) |
| 229 | |
| 230 | priorities: dict[KVCacheLayout, int] = defaultdict(int) |
| 231 | for preferred_layout, *_ in supported_layouts_lists: |
| 232 | priorities[preferred_layout] += 1 |
| 233 | supported_layouts = set.intersection(*map(set, supported_layouts_lists)) |
| 234 | candidates = sorted( |
| 235 | (layout for layout in KVCacheLayout if layout in supported_layouts), |
| 236 | key=lambda layout: priorities[layout], |
| 237 | reverse=True, |
| 238 | ) |
| 239 | if not candidates: |
| 240 | raise ValueError( |
| 241 | "No KV cache layout satisfies every supported set: " |
| 242 | f"{list(map(_layout_names, supported_layouts_lists))}." |
| 243 | ) |
| 244 | return candidates |
| 245 | |
| 246 | |
| 247 | def record_kv_cache_layout(cache_config: CacheConfig, layout_name: str) -> None: |
| 248 | """Adopt a layout resolved elsewhere (the engine core) in this process.""" |
| 249 | layout = _layout_from_name(layout_name) |
| 250 | existing = cache_config.kv_cache_layout |
| 251 | if existing is not None and existing != layout.name: |
| 252 | raise ValueError( |
| 253 | f"KV cache layout is already resolved to {existing}; " |
| 254 | f"cannot change it to {layout.name}." |
| 255 | ) |
| 256 | cache_config.kv_cache_layout = layout.name |
| 257 | |
| 258 | |
| 259 | def resolve_kv_cache_layout( |
| 260 | vllm_config: VllmConfig, |
| 261 | supported_layouts: list[list[str]], |
| 262 | kv_cache_specs: Iterable[KVCacheSpec] | None = None, |
| 263 | ) -> KVCacheLayout: |
| 264 | """Resolve one KV cache layout for the whole model. |
| 265 | |
| 266 | Runs once in the engine core. Every worker reports the layouts its backends |
| 267 | support, most preferred first (``get_supported_kv_cache_layouts``); all |
| 268 | ranks run the same backends, so their lists must agree. Specs mixing HNC |
| 269 | shapes narrow the candidates to block-compact layouts. An explicit |
| 270 | ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails, |
| 271 | with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the |
| 272 | connector's preference is used when compatible and dropped with a warning |
| 273 | otherwise. A layout already present on ``cache_config`` wins outright, and |
| 274 | the result is recorded there (see ``CacheConfig.kv_cache_layout``); it |
| 275 | reaches workers through the ``set_kv_cache_layout`` RPC and |
| 276 | ``KVCacheConfig.kv_cache_layout``. |
| 277 | """ |
| 278 | cache_config = vllm_config.cache_config |
| 279 | if cache_config.kv_cache_layout is not None: |
| 280 | return cache_config.get_resolved_kv_cache_layout() |
| 281 | |
| 282 | assert supported_layouts and all(supported_layouts), ( |
| 283 | "No worker reported supported KV cache layouts." |
| 284 | ) |
| 285 | assert all(names == supported_layouts[0] for names in supported_layouts[1:]), ( |
| 286 | f"Workers disagree on supported KV cache layouts: {supported_layouts}." |
| 287 | ) |
| 288 | candidates = [_layout_from_name(name) for name in supported_layouts[0]] |
| 289 | |
| 290 | # A block-compact layout means the block is densely packed in memory, so any mix of |
| 291 | # specs can re-interpret HNC with different sizes as long as the total number of |
| 292 | # bytes is the same. If not block-compact, each spec must agree on HNC to alias |
| 293 | # the same page (this aliasing is done by the Hybrid Memory Allocator, HMA). |
| 294 | # Specs without per-layer views lay out their own raw backing tensor. |
| 295 | kv_cache_specs = tuple( |
| 296 | spec for spec in kv_cache_specs or () if spec.has_layer_views |
| 297 | ) |
| 298 | hnc_shapes = { |
| 299 | (spec.num_heads, spec.num_states, spec.page_size_bytes) |
| 300 | for spec in kv_cache_specs |
| 301 | } |
| 302 | if len(hnc_shapes) > 1: |
| 303 | candidates = [m for m in candidates if m.is_block_compact] |
| 304 | if not candidates: |
| 305 | raise ValueError( |
| 306 | "Specs with mixed HNC shapes need a block-compact layout, but " |
| 307 | f"none is in every supported set: {supported_layouts}." |
| 308 | ) |
| 309 | |
| 310 | dcp_sharding = { |
| 311 | spec.dcp_sharded for spec in kv_cache_specs if isinstance(spec, AttentionSpec) |
| 312 | } |
| 313 | page_sizes = {spec.page_size_bytes for spec in kv_cache_specs} |
| 314 | if len(dcp_sharding) > 1 and len(page_sizes) > 1: |
| 315 | # Sharded target and replicated draft caches need independent groups. |
| 316 | # Block-outer layouts can pack those groups without equalizing pages. |
| 317 | candidates = [layout for layout in candidates if layout.is_block_outermost] |
| 318 | if not candidates: |
| 319 | raise ValueError( |
| 320 | "DCP with a replicated draft and mixed KV page sizes requires " |
| 321 | f"a block-outer KV cache layout; supported sets: {supported_layouts}." |
| 322 | ) |
| 323 | |
| 324 | if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None: |
| 325 | layout = _layout_from_name(requested) |
| 326 | if layout not in candidates: |
| 327 | raise ValueError( |
| 328 | f"VLLM_KV_CACHE_LAYOUT={requested} does not satisfy every " |
| 329 | f"supported set; valid layouts: {_layout_names(candidates)}." |
| 330 | ) |
| 331 | elif (connector := get_kv_connector_cache_layout(vllm_config)) is not None: |
| 332 | layout = _layout_from_name(connector) |
| 333 | if layout not in candidates: |
| 334 | logger.warning_once( |
| 335 | f"KV connector cache layout {connector} does not satisfy every " |
| 336 | f"supported set; valid layouts: {_layout_names(candidates)}. " |
| 337 | f"Using {candidates[0].name} instead." |
| 338 | ) |
| 339 | layout = candidates[0] |
| 340 | else: |
| 341 | layout = candidates[0] |
| 342 | |
| 343 | logger.info_once("Using %s KV cache layout.", layout.name) |
| 344 | cache_config.kv_cache_layout = layout.name |
| 345 | return layout |
| 346 | |
| 347 | |
| 348 | @dataclass |
| 349 | class PerLayerParameters: |
| 350 | """Currently, FlashInfer backend only support models in which all layers share |
| 351 | the same values for the following hyperparameters. Should not be used for |
| 352 | trtllm-gen backend since it supports different values for the following |
| 353 | hyperparameters. |
| 354 | """ |
| 355 | |
| 356 | window_left: int |
| 357 | logits_soft_cap: float | None |
| 358 | sm_scale: float |
| 359 | has_sinks: bool = False |
| 360 | # has same params for all layers |
| 361 | has_same_window_lefts: bool | None = field(default=None, compare=False) |
| 362 | has_same_all_params: bool | None = field(default=None, compare=False) |
| 363 | |
| 364 | |
| 365 | def get_per_layer_parameters( |
| 366 | vllm_config: VllmConfig, layer_names: list[str], cls_: type["AttentionImpl"] |
| 367 | ) -> dict[str, PerLayerParameters]: |
| 368 | """Scan layers in `layer_names` and determine some hyperparameters |
| 369 | to use during `plan`. |
| 370 | """ |
| 371 | layers = get_layers_from_vllm_config( |
| 372 | vllm_config, |
| 373 | AttentionLayerBase, # type: ignore[type-abstract] |
| 374 | layer_names, |
| 375 | ) |
| 376 | per_layer_params: dict[str, PerLayerParameters] = {} |
| 377 | |
| 378 | for key, layer in layers.items(): |
| 379 | impl = find_attention_impl_variant(layer.impl, cls_) |
| 380 | assert impl is not None |
| 381 | |
| 382 | # Infer hyperparameters from the attention layer |
| 383 | window_size = getattr(impl, "sliding_window", None) |
| 384 | window_left = window_size[0] if window_size is not None else -1 |
| 385 | logits_soft_cap = getattr(impl, "logits_soft_cap", None) |
| 386 | sm_scale = impl.scale |
| 387 | has_sinks = getattr(impl, "sinks", None) is not None |
| 388 | |
| 389 | per_layer_params[key] = PerLayerParameters( |
| 390 | window_left, logits_soft_cap, sm_scale, has_sinks |
| 391 | ) |
| 392 | |
| 393 | return per_layer_params |
| 394 | |
| 395 | |
| 396 | def get_num_attention_heads_from_layers( |
| 397 | vllm_config: VllmConfig, layer_names: list[str] |
| 398 | ) -> int | None: |
| 399 | """Per-TP-rank ``num_heads`` shared by the named Attention layers. |
| 400 | |
| 401 | Use in metadata builders whose plan-time allocations depend on the |
| 402 | head count: the model-wide ``get_num_attention_heads()`` is wrong |
| 403 | for models with non-uniform per-layer head counts. All layers in |
| 404 | one attention group must agree on ``num_heads``; this is asserted. |
| 405 | Returns ``None`` when no matching Attention layer is found. |
| 406 | """ |
| 407 | attn_layers = get_layers_from_vllm_config( |
| 408 | vllm_config, |
| 409 | AttentionLayerBase, # type: ignore[type-abstract] |
| 410 | layer_names, |
| 411 | ) |
| 412 | if not attn_layers: |
| 413 | return None |
| 414 | heads = { |
| 415 | cast(Any, getattr(layer, "impl", layer)).num_heads |
| 416 | for layer in attn_layers.values() |
| 417 | } |
| 418 | assert len(heads) == 1, ( |
| 419 | f"All layers in one attention group must share num_heads; " |
| 420 | f"got {heads} for {layer_names}." |
| 421 | ) |
| 422 | return heads.pop() |
| 423 | |
| 424 | |
| 425 | def infer_global_hyperparameters( |
| 426 | per_layer_params: dict[str, PerLayerParameters], |
| 427 | ) -> PerLayerParameters: |
| 428 | """Currently, FlashInfer backend other than trtllm-gen |
| 429 | only support models in which all layers share |
| 430 | the same values for the following hyperparameters: |
| 431 | - `window_left` |
| 432 | - `logits_soft_cap` |
| 433 | - `sm_scale` |
| 434 | |
| 435 | So this function asserts that all layers share the same values for these |
| 436 | hyperparameters and returns the global values. |
| 437 | """ |
| 438 | assert len(per_layer_params) > 0, "No attention layers found in the model." |
| 439 | |
| 440 | param_sets = list(per_layer_params.values()) |
| 441 | global_params = param_sets[0] |
| 442 | |
| 443 | global_params.has_same_window_lefts = all( |
| 444 | params.window_left == global_params.window_left for params in param_sets |
| 445 | ) |
| 446 | global_params.has_same_all_params = all( |
| 447 | params == global_params for params in param_sets |
| 448 | ) |
| 449 | |
| 450 | return global_params |
| 451 | |
| 452 | |
| 453 | # |
| 454 | # Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into |
| 455 | # local attention blocks, where each block is passed to the attention kernel |
| 456 | # as an independent local ("virtual") batch item. |
| 457 | # |
| 458 | # For example, if are performing a chunked prefill a batch of 3 sequences: |
| 459 | # q_seqlens = [4, 10, 5] |
| 460 | # kv_seqlens = [6, 17, 9] |
| 461 | # Then normally for regular attention we would compute with an attention mask |
| 462 | # for batch idx 0 (q_seqlens = 4, kv_seqlens = 6) like: |
| 463 | # batch idx: 0 (q_seqlens = 4, kv_seqlens = 6) |
| 464 | # k_toks > 0 1 2 3 4 5 |
| 465 | # q_toks v _____________ |
| 466 | # 0 | 1 1 1 |
| 467 | # 1 | 1 1 1 1 |
| 468 | # 2 | 1 1 1 1 1 |
| 469 | # 3 | 1 1 1 1 1 1 |
| 470 | # |
| 471 | # for local attention (with attn_chunk_size = 4) we would compute with an |
| 472 | # attention mask like: |
| 473 | # batch idx: 0 (q_seqlens = 4, kv_seqlens = 6, attn_chunk_size = 4) |
| 474 | # k_toks > 0 1 2 3 4 5 |
| 475 | # q_toks v _____________ |
| 476 | # 0 | 1 1 1 |
| 477 | # 1 | 1 1 1 1 |
| 478 | # 2 | 1 |
| 479 | # 3 | 1 1 |
| 480 | # |
| 481 | # We can simulate this mask using standard flash-attention by breaking the |
| 482 | # sequences into local ("virtual") batches, where each local batch item is a |
| 483 | # local attention block, so in this case batch idx 0 would be broken up into: |
| 484 | # |
| 485 | # local-batch idx: 0 (q_seqlens = 2, kv_seqlens = 4) (batch 0) |
| 486 | # k_toks > 0 1 2 3 |
| 487 | # q_toks v _____________ |
| 488 | # 0 | 1 1 1 |
| 489 | # 1 | 1 1 1 1 |
| 490 | # local-batch idx: 1 (q_seqlens = 2, kv_seqlens = 2) (batch 0) |
| 491 | # k_toks > 4 5 |
| 492 | # q_toks v _____________ |
| 493 | # 2 | 1 |
| 494 | # 3 | 1 1 |
| 495 | # |
| 496 | # e.g. if we have: |
| 497 | # attn_chunk_size = 4 |
| 498 | # query_start_loc_np = [0, 4, 14, 19] (q_seqlens = [4, 10, 5]) |
| 499 | # Then this function would return: |
| 500 | # __b0__ ______b1______ __b2__ < orig batch indices |
| 501 | # q_seqlens_local = [ 2, 2, 1, 4, 4, 1, 4, 1] |
| 502 | # cu_seqlens_q_local = [0, 4, 6, 10, 14, 18, 19, 23, 24] |
| 503 | # seqlens_k_local = [ 4, 2, 4, 4, 4, 1, 4, 1] |
| 504 | # block_table_local : shape[local_virtual_batches, pages_per_local_batch] |
| 505 | def make_local_attention_virtual_batches( |
| 506 | attn_chunk_size: int, |
| 507 | common_attn_metadata: CommonAttentionMetadata, |
| 508 | block_size: int = 0, |
| 509 | ) -> tuple[CommonAttentionMetadata, Callable[[torch.Tensor], torch.Tensor]]: |
| 510 | query_start_loc_np = common_attn_metadata.query_start_loc_cpu.numpy() |
| 511 | with gpu_sync_allowed(): |
| 512 | # TODO see https://github.com/vllm-project/vllm/pull/31852 |
| 513 | seq_lens_np = common_attn_metadata.seq_lens.cpu().numpy() |
| 514 | block_table = common_attn_metadata.block_table_tensor |
| 515 | device = common_attn_metadata.query_start_loc.device |
| 516 | |
| 517 | q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1] |
| 518 | actual_batch_size = seq_lens_np.shape[0] |
| 519 | |
| 520 | # Handle if we are starting in the middle of a local attention block, |
| 521 | # we assume q_seqlens > 0 (for all elements), for each batch idx we compute |
| 522 | # the number of tokens that are not in the first local attention block and |
| 523 | # then we can simply use a cdiv for the rest. |
| 524 | # For example if we have: |
| 525 | # attn_chunk_size = 4 |
| 526 | # q_seqlens = [4, 10, 5] |
| 527 | # k_seqlens = [6, 17, 9] |
| 528 | # Then we would get: |
| 529 | # new_tokens_in_first_block = [2, 1, 4] |
| 530 | # local_blocks = [2, 4, 2] |
| 531 | q_tokens_in_first_block = np.minimum( |
| 532 | attn_chunk_size - ((seq_lens_np - q_seqlens) % attn_chunk_size), q_seqlens |
| 533 | ).astype(np.int32) |
| 534 | tokens_in_last_block = attn_chunk_size + (seq_lens_np % -attn_chunk_size) |
| 535 | local_blocks = 1 + cdiv(q_seqlens - q_tokens_in_first_block, attn_chunk_size) |
| 536 | |
| 537 | # Once we know the number of local blocks we can compute the request spans |
| 538 | # for each batch idx, we can figure out the number of "virtual" requests we |
| 539 | # have to make, |
| 540 | # For the above example we would get: |
| 541 | # seqlens_q_local = [2, 2, 1, 4, 4, 1, 4, 1] |
| 542 | # |
| 543 | # First Get batched arange. (E.g., [2, 4, 2] -> [0, 1, 0, 1, 2, 3, 0, 1]) |
| 544 | # (TODO: make a utility to share this code with _prepare_inputs) |
| 545 | # arange step 1. [2, 4, 2] -> [2, 6, 8] |
| 546 | cu_num_blocks = np.cumsum(local_blocks) |
| 547 | virtual_batches = cu_num_blocks[-1] |
| 548 | # arange step 2. [2, 6, 8] -> [0, 0, 2, 2, 2, 2, 6, 6] |
| 549 | block_offsets = np.repeat(cu_num_blocks - local_blocks, local_blocks) |
| 550 | # arange step 3. [0, 1, 0, 1, 2, 3, 0, 1] |
| 551 | arange = np.arange(virtual_batches, dtype=np.int32) - block_offsets |
| 552 | # also compute reverse arange (i.e. [1, 0, 3, 2, 1, 0, 1, 0]) |
| 553 | rarange = np.repeat(local_blocks, local_blocks) - arange - 1 |
| 554 | # Then we can compute the seqlens_q_local, handling the fact that the |
| 555 | # first and last blocks could be partial |
| 556 | seqlens_q_local = np.repeat(q_seqlens - q_tokens_in_first_block, local_blocks) |
| 557 | # set the first block since this may be a partial block |
| 558 | seqlens_q_local[arange == 0] = q_tokens_in_first_block |
| 559 | # set the remaining blocks |
| 560 | seqlens_q_local[arange > 0] = np.minimum( |
| 561 | seqlens_q_local - attn_chunk_size * (arange - 1), attn_chunk_size |
| 562 | )[arange > 0] |
| 563 | |
| 564 | # convert from q_seqlens to cu_seqlens_q |
| 565 | cu_seqlens_q_local = np.empty(virtual_batches + 1, dtype=np.int32) |
| 566 | np.cumsum(seqlens_q_local, out=cu_seqlens_q_local[1:]) |
| 567 | cu_seqlens_q_local[0] = 0 |
| 568 | |
| 569 | # compute the seqlens_k_local, |
| 570 | # basically a full local attention block for all but the last block in each |
| 571 | # batch |
| 572 | # For our example this will be: |
| 573 | # seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1] |
| 574 | seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32) |
| 575 | seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block |
| 576 | k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - ( |
| 577 | rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks) |
| 578 | ) |
| 579 | # For the example the local attention blocks start at: |
| 580 | # _b0_ _____b1_____ _b2_ |
| 581 | # k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8] |
| 582 | block_starts = k_seqstarts_absolute // block_size |
| 583 | assert attn_chunk_size % block_size == 0, ( |
| 584 | f"attn_chunk_size {attn_chunk_size} is not divisible by block_size {block_size}" |
| 585 | ) |
| 586 | pages_per_local_batch = attn_chunk_size // block_size |
| 587 | |
| 588 | # Create a block_table for the local attention blocks |
| 589 | # For out example if we have a block-table like (assuming block_size=2): |
| 590 | # block_table = [ |
| 591 | # [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], < batch 0 |
| 592 | # [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], < batch 1 |
| 593 | # [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], < batch 2 |
| 594 | # ] |
| 595 | # Then for the local batches we would want a block-table like |
| 596 | # block_table_local = [ |
| 597 | # [ 0, 1 ], < local-batch 0, (batch 0, starting from k[0]) |
| 598 | # [ 2, 3 ], < local-batch 1, (batch 0, starting from k[4]) |
| 599 | # [ 12, 13 ], < local-batch 2, (batch 1, starting from k[4]) |
| 600 | # [ 14, 15 ], < local-batch 3, (batch 1, starting from k[8]) |
| 601 | # [ 16, 17 ], < local-batch 4, (batch 1, starting from k[12]) |
| 602 | # [ 18, 19 ], < local-batch 5, (batch 1, starting from k[16]) |
| 603 | # [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4]) |
| 604 | # [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8]) |
| 605 | # ] |
| 606 | block_indices = block_starts[:, None] + np.arange( |
| 607 | pages_per_local_batch, dtype=np.int32 |
| 608 | ) |
| 609 | block_indices = block_indices.reshape(-1).clip(max=block_table.shape[1] - 1) |
| 610 | batch_indices = np.repeat( |
| 611 | np.arange(actual_batch_size, dtype=np.int32), |
| 612 | local_blocks * pages_per_local_batch, |
| 613 | ) |
| 614 | |
| 615 | # NOTE: https://github.com/pytorch/pytorch/pull/160256 causes performance |
| 616 | # regression when using numpy arrays (batch and block indices) to index into |
| 617 | # torch tensor (block_table). As a workaround, convert numpy arrays to torch |
| 618 | # tensor first, which recovers perf. |
| 619 | # Upload the index tensors to the block_table's device up-front so that the |
| 620 | # fancy indexing below doesn't implicitly force a synchronous H2D copy. |
| 621 | batch_indices_torch = async_tensor_h2d(batch_indices, device=device) |
| 622 | block_indices_torch = async_tensor_h2d(block_indices, device=device) |
| 623 | |
| 624 | # Save as a lambda so we can return this for update_block_table |
| 625 | make_block_table = lambda block_table: block_table[ |
| 626 | batch_indices_torch, block_indices_torch |
| 627 | ].view(virtual_batches, -1) |
| 628 | block_table_local = make_block_table(block_table) |
| 629 | |
| 630 | query_start_loc_cpu = torch.from_numpy(cu_seqlens_q_local) |
| 631 | seq_lens_cpu = torch.from_numpy(seqlens_k_local) |
| 632 | max_seq_len = int(seq_lens_cpu.max()) |
| 633 | |
| 634 | return CommonAttentionMetadata( |
| 635 | query_start_loc_cpu=query_start_loc_cpu, |
| 636 | query_start_loc=async_tensor_h2d(query_start_loc_cpu, device=device), |
| 637 | seq_lens=async_tensor_h2d(seq_lens_cpu, device=device), |
| 638 | num_reqs=len(seq_lens_cpu), |
| 639 | num_actual_tokens=common_attn_metadata.num_actual_tokens, |
| 640 | max_query_len=seqlens_q_local.max(), |
| 641 | max_seq_len=max_seq_len, |
| 642 | block_table_tensor=block_table_local, |
| 643 | slot_mapping=common_attn_metadata.slot_mapping, |
| 644 | causal=True, |
| 645 | seq_lens_cpu_upper_bound=seq_lens_cpu, |
| 646 | ), make_block_table |
| 647 | |
| 648 | |
| 649 | def make_kv_sharing_fast_prefill_common_attn_metadata( |
| 650 | common_attn_metadata: CommonAttentionMetadata, |
| 651 | ) -> CommonAttentionMetadata: |
| 652 | if common_attn_metadata.max_query_len == 1: |
| 653 | # All requests are decode (assume 1 token for now) |
| 654 | # Skip computing fast prefill path |
| 655 | return common_attn_metadata |
| 656 | |
| 657 | if ( |
| 658 | common_attn_metadata.logits_indices_padded is None |
| 659 | or common_attn_metadata.num_logits_indices is None |
| 660 | ): |
| 661 | # Fast prefill not armed for this step (e.g. cudagraph capture, or a |
| 662 | # pure-decode step): run the KV-sharing layers on the full batch. |
| 663 | return common_attn_metadata |
| 664 | |
| 665 | logits_indices_padded = common_attn_metadata.logits_indices_padded |
| 666 | num_logits_indices = common_attn_metadata.num_logits_indices |
| 667 | # Get rid of CUDAGraph padding, if any |
| 668 | logits_indices = logits_indices_padded[:num_logits_indices] |
| 669 | num_reqs = common_attn_metadata.num_reqs |
| 670 | query_start_loc = common_attn_metadata.query_start_loc |
| 671 | # Example inputs |
| 672 | # num_reqs: 3 |
| 673 | # generation_indices: [14, 18, 19, 27] |
| 674 | # query_start_loc: [0, 15, 20, 28] |
| 675 | # seq_lens: [41, 31, 40] |
| 676 | |
| 677 | # Find how many decode indices belong to each request |
| 678 | # request_ids: [0, 1, 1, 2] |
| 679 | request_ids = torch.bucketize(logits_indices, query_start_loc[1:], right=True) |
| 680 | |
| 681 | # Figure out how many tokens are in each request |
| 682 | # num_decode_tokens: [1, 2, 1] |
| 683 | # Avoid `torch.bincount` here — on CUDA it forces a sync to determine |
| 684 | # the output size (even with `minlength`, the kernel must confirm no |
| 685 | # value exceeds the bound). `scatter_add_` into a preallocated buffer |
| 686 | # is equivalent and stays async. |
| 687 | num_decode_tokens = torch.zeros( |
| 688 | num_reqs, dtype=request_ids.dtype, device=request_ids.device |
| 689 | ) |
| 690 | num_decode_tokens.scatter_add_( |
| 691 | 0, request_ids.to(num_decode_tokens.dtype), torch.ones_like(request_ids) |
| 692 | ) |
| 693 | |
| 694 | # Calculate new query_start_loc with tokens in generation_indices |
| 695 | # decode_query_start_loc: [0, 1, 3, 4] |
| 696 | decode_query_start_loc = torch.empty( |
| 697 | num_reqs + 1, device=query_start_loc.device, dtype=query_start_loc.dtype |
| 698 | ) |
| 699 | |
| 700 | decode_query_start_loc[:1].fill_(0) # Avoid sync from scalar assignment. |
| 701 | decode_query_start_loc[1:] = torch.cumsum(num_decode_tokens, dim=0) |
| 702 | |
| 703 | # `num_decode_tokens` is a histogram over `logits_indices`, so its total is |
| 704 | # just how many there were -- already known as a Python int. |
| 705 | total_num_decode_tokens = num_logits_indices |
| 706 | |
| 707 | # Largest per-request logits count. |
| 708 | decode_max_query_len = common_attn_metadata.max_logits_per_req |
| 709 | assert decode_max_query_len is not None |
| 710 | |
| 711 | common_attn_metadata = CommonAttentionMetadata( |
| 712 | query_start_loc=decode_query_start_loc, |
| 713 | query_start_loc_cpu=decode_query_start_loc.to("cpu", non_blocking=True), |
| 714 | seq_lens=common_attn_metadata.seq_lens, |
| 715 | num_reqs=num_reqs, |
| 716 | num_actual_tokens=total_num_decode_tokens, |
| 717 | max_query_len=decode_max_query_len, |
| 718 | max_seq_len=common_attn_metadata.max_seq_len, |
| 719 | block_table_tensor=common_attn_metadata.block_table_tensor, |
| 720 | slot_mapping=common_attn_metadata.slot_mapping, |
| 721 | causal=True, |
| 722 | seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, |
| 723 | ) |
| 724 | return common_attn_metadata |
| 725 | |
| 726 | |
| 727 | def split_decodes_prefills_and_extends( |
| 728 | common_attn_metadata: CommonAttentionMetadata, |
| 729 | decode_threshold: int = 1, |
| 730 | ) -> tuple[int, int, int, int, int, int]: |
| 731 | """Assuming a reordered batch, finds the boundary between prefill and decode |
| 732 | requests. |
| 733 | |
| 734 | Args: |
| 735 | common_attn_metadata: CommonAttentionMetadata object containing the |
| 736 | batch metadata. |
| 737 | decode_threshold: The maximum query length to be considered a decode. |
| 738 | |
| 739 | Returns: |
| 740 | num_decodes: The number of decode requests. |
| 741 | num_extends: The number of extend requests. |
| 742 | num_prefills: The number of prefill requests. |
| 743 | num_decode_tokens: The number of tokens in the decode requests. |
| 744 | num_extend_tokens: The number of tokens in the extend requests. |
| 745 | num_prefill_tokens: The number of tokens in the prefill requests. |
| 746 | |
| 747 | """ |
| 748 | max_query_len = common_attn_metadata.max_query_len |
| 749 | num_reqs = common_attn_metadata.num_reqs |
| 750 | num_tokens = common_attn_metadata.num_actual_tokens |
| 751 | query_start_loc = common_attn_metadata.query_start_loc_cpu |
| 752 | |
| 753 | if max_query_len <= decode_threshold: |
| 754 | return num_reqs, 0, 0, num_tokens, 0, 0 |
| 755 | |
| 756 | # Upper bound is exact for prefill rows; decode rows still satisfy |
| 757 | # seq_len > query_len under the optimistic bound, so `seq_lens == |
| 758 | # query_lens` identifies prefills correctly either way. |
| 759 | assert common_attn_metadata.seq_lens_cpu_upper_bound is not None |
| 760 | seq_lens = common_attn_metadata.seq_lens_cpu_upper_bound |
| 761 | |
| 762 | query_lens = query_start_loc[1:] - query_start_loc[:-1] |
| 763 | is_prefill_or_extend = query_lens > decode_threshold |
| 764 | is_prefill = (seq_lens == query_lens) & is_prefill_or_extend |
| 765 | first_extend = is_prefill_or_extend.int().argmax(dim=-1).item() |
| 766 | first_prefill = is_prefill.int().argmax(dim=-1).item() |
| 767 | num_decodes = first_extend |
| 768 | num_decode_tokens = query_start_loc[first_extend].item() |
| 769 | if not torch.any(is_prefill_or_extend): |
| 770 | return (num_decodes, 0, 0, num_decode_tokens, 0, 0) |
| 771 | |
| 772 | num_prefills_or_extends = num_reqs - num_decodes |
| 773 | num_prefill_or_extend_tokens = num_tokens - num_decode_tokens |
| 774 | if not torch.any(is_prefill): |
| 775 | return ( |
| 776 | num_decodes, |
| 777 | num_prefills_or_extends, |
| 778 | 0, |
| 779 | num_decode_tokens, |
| 780 | num_prefill_or_extend_tokens, |
| 781 | 0, |
| 782 | ) |
| 783 | |
| 784 | num_extends = first_prefill - num_decodes |
| 785 | num_prefills = num_reqs - first_prefill |
| 786 | |
| 787 | num_prefill_tokens = num_tokens - query_start_loc[first_prefill] |
| 788 | num_extend_tokens = num_prefill_or_extend_tokens - num_prefill_tokens |
| 789 | return ( |
| 790 | num_decodes, |
| 791 | num_extends, |
| 792 | num_prefills, |
| 793 | num_decode_tokens, |
| 794 | num_extend_tokens, |
| 795 | num_prefill_tokens, |
| 796 | ) |
| 797 | |
| 798 | |
| 799 | def split_decodes_and_prefills( |
| 800 | common_attn_metadata: CommonAttentionMetadata, |
| 801 | decode_threshold: int = 1, |
| 802 | require_uniform: bool = False, |
| 803 | treat_short_extends_as_decodes: bool = True, |
| 804 | ) -> tuple[int, int, int, int]: |
| 805 | """Assuming a reordered batch, finds the boundary between prefill and decode |
| 806 | requests. |
| 807 | |
| 808 | The batch is expected to be ordered as: |
| 809 | decode → short_extend → long_extend → prefill |
| 810 | |
| 811 | Args: |
| 812 | common_attn_metadata: CommonAttentionMetadata object containing the |
| 813 | batch metadata. |
| 814 | decode_threshold: The maximum query length to be considered a decode. |
| 815 | require_uniform: If True, requires that all decode requests have the |
| 816 | same query length. When set, some queries may be considered prefills |
| 817 | even if they are <= decode_threshold, in order to ensure uniformity. |
| 818 | treat_short_extends_as_decodes: If True (default), short extends |
| 819 | (query_len <= threshold but still prefilling) are counted as |
| 820 | decodes. If False, they are counted as prefills. |
| 821 | |
| 822 | Returns: |
| 823 | num_decodes: The number of decode requests. |
| 824 | num_prefills: The number of prefill requests. |
| 825 | num_decode_tokens: The number of tokens in the decode requests. |
| 826 | num_prefill_tokens: The number of tokens in the prefill requests. |
| 827 | |
| 828 | """ |
| 829 | max_query_len = common_attn_metadata.max_query_len |
| 830 | num_reqs = common_attn_metadata.num_reqs |
| 831 | num_tokens = common_attn_metadata.num_actual_tokens |
| 832 | query_start_loc = common_attn_metadata.query_start_loc_cpu |
| 833 | |
| 834 | if ( |
| 835 | max_query_len <= decode_threshold |
| 836 | and (not require_uniform or decode_threshold <= 1) |
| 837 | and treat_short_extends_as_decodes |
| 838 | ): |
| 839 | return num_reqs, 0, num_tokens, 0 |
| 840 | |
| 841 | query_lens = query_start_loc[1:] - query_start_loc[:-1] |
| 842 | if query_lens[0].item() > decode_threshold: |
| 843 | # first request is not decode, so no decode requests |
| 844 | return 0, num_reqs, 0, num_tokens |
| 845 | |
| 846 | if require_uniform: |
| 847 | # check if we are in a padded uniform batch; this is used for full-CGs, some |
| 848 | # requests may have a query length of 0 but since they are padding its fine |
| 849 | # to treat them as decodes (ensures num_decodes matches the captured size) |
| 850 | if treat_short_extends_as_decodes and torch.all( |
| 851 | (query_lens == query_lens[0]) | (query_lens == 0) |
| 852 | ): |
| 853 | return num_reqs, 0, num_tokens, 0 # all decodes |
| 854 | is_prefill = (query_lens != query_lens[0]) & (query_lens != 0) |
| 855 | else: |
| 856 | is_prefill = query_lens > decode_threshold |
| 857 | |
| 858 | if not treat_short_extends_as_decodes: |
| 859 | assert common_attn_metadata.is_prefilling is not None |
| 860 | is_prefill |= common_attn_metadata.is_prefilling |
| 861 | |
| 862 | if not torch.any(is_prefill): |
| 863 | return num_reqs, 0, num_tokens, 0 |
| 864 | |
| 865 | first_prefill = is_prefill.int().argmax(dim=-1).item() |
| 866 | num_decodes = first_prefill |
| 867 | num_prefills = num_reqs - num_decodes |
| 868 | num_decode_tokens = query_start_loc[first_prefill].item() |
| 869 | num_prefill_tokens = num_tokens - num_decode_tokens |
| 870 | return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) |
| 871 | |
| 872 | |
| 873 | def split_prefill_chunks( |
| 874 | seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0 |
| 875 | ) -> list[tuple[int, int]]: |
| 876 | """Split the prefill requests into chunks such that the total sequence length |
| 877 | of each chunk is less than or equal to the workspace size. |
| 878 | |
| 879 | Args: |
| 880 | seq_lens_cpu: The sequence lengths of the prefill requests on CPU. |
| 881 | workspace_size: The maximum workspace size (in tokens) per chunk. |
| 882 | request_offset: The offset to add to the request indices. |
| 883 | |
| 884 | Returns: |
| 885 | A list of tuples of (reqs_start, reqs_end) representing chunk boundaries. |
| 886 | |
| 887 | """ |
| 888 | chunk_bounds = [] |
| 889 | i, n = 0, len(seq_lens_cpu) |
| 890 | assert torch.all(seq_lens_cpu <= workspace_size).item() |
| 891 | |
| 892 | while i < n: |
| 893 | start, chunk_total = i, 0 |
| 894 | while i < n and (chunk_total + (s := seq_lens_cpu[i].item())) <= workspace_size: |
| 895 | chunk_total += s |
| 896 | i += 1 |
| 897 | chunk_bounds.append((start + request_offset, i + request_offset)) |
| 898 | return chunk_bounds |
| 899 | |
| 900 | |
| 901 | def reorder_batch_to_split_decodes_and_prefills( |
| 902 | input_batch: "InputBatch", |
| 903 | scheduler_output: "SchedulerOutput", |
| 904 | decode_threshold: int = 1, |
| 905 | ) -> bool: |
| 906 | """Reorders the batch to split into prefill and decode requests; places all |
| 907 | requests with <= decode_threshold tokens at the front of the batch. |
| 908 | |
| 909 | The batch is reordered into 4 regions: |
| 910 | decode: (num_scheduled <= threshold AND is not prefilling) |
| 911 | short_extend: (num_scheduled <= threshold AND is chunked prefilling) |
| 912 | long_extend: (num_scheduled > threshold AND is chunked prefilling) |
| 913 | prefill: (num_computed == 0) # First chunks |
| 914 | |
| 915 | Returns: |
| 916 | True if the batch was modified, False otherwise. |
| 917 | |
| 918 | """ |
| 919 | num_reqs = len(input_batch.req_ids) |
| 920 | num_scheduled_tokens = [ |
| 921 | scheduler_output.num_scheduled_tokens[id] for id in input_batch.req_ids |
| 922 | ] |
| 923 | num_scheduled_tokens_np = np.array(num_scheduled_tokens) |
| 924 | num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs] |
| 925 | num_prompt_tokens_np = input_batch.num_prompt_tokens[:num_reqs] |
| 926 | |
| 927 | has_context = num_computed_tokens_np > 0 |
| 928 | is_below_threshold = num_scheduled_tokens_np <= decode_threshold |
| 929 | done_prefilling = num_computed_tokens_np >= num_prompt_tokens_np |
| 930 | |
| 931 | # Mutually exclusive categories (exactly one True per request): |
| 932 | # 1. No context yet -> prefill |
| 933 | # 2. Has context, above threshold -> long_extend |
| 934 | # 3. Has context, below threshold, still prefilling -> short_extend |
| 935 | # 4. Has context, below threshold, done prefilling -> decode |
| 936 | is_pure_prefill = ~has_context |
| 937 | is_long_extend = has_context & ~is_below_threshold |
| 938 | is_short_extend = has_context & is_below_threshold & ~done_prefilling |
| 939 | is_decode = has_context & is_below_threshold & done_prefilling |
| 940 | |
| 941 | # Desired order: decode → short_extend → long_extend → prefill |
| 942 | req_regions = np.zeros(num_reqs, dtype=np.int32) # 0 = decode by default |
| 943 | req_regions[is_short_extend] = 1 |
| 944 | req_regions[is_long_extend] = 2 |
| 945 | req_regions[is_pure_prefill] = 3 |
| 946 | |
| 947 | num_decodes = int(is_decode.sum()) |
| 948 | num_short_extends = int(is_short_extend.sum()) |
| 949 | num_long_extends = int(is_long_extend.sum()) |
| 950 | num_prefills = int(is_pure_prefill.sum()) |
| 951 | |
| 952 | target_regions = np.repeat( |
| 953 | [0, 1, 2, 3], |
| 954 | [num_decodes, num_short_extends, num_long_extends, num_prefills], |
| 955 | ).astype(np.int32) |
| 956 | |
| 957 | needs_swap = req_regions != target_regions |
| 958 | |
| 959 | if not needs_swap.any(): |
| 960 | return False |
| 961 | |
| 962 | # Extract indices that need swapping and sort by target region |
| 963 | orig_indices = np.where(needs_swap)[0] |
| 964 | sorted_order = np.argsort(req_regions[needs_swap], kind="stable") |
| 965 | src_indices = orig_indices[sorted_order] |
| 966 | |
| 967 | src_dest_map = {int(src): int(dst) for src, dst in zip(src_indices, orig_indices)} |
| 968 | |
| 969 | for src in src_dest_map: |
| 970 | dst = src_dest_map[src] |
| 971 | while src != dst: |
| 972 | input_batch.swap_states(src, dst) |
| 973 | # Mark dst as done by updating its destination to itself |
| 974 | next_dst = src_dest_map.get(dst, dst) |
| 975 | src_dest_map[dst] = dst |
| 976 | dst = next_dst |
| 977 | |
| 978 | return True |
| 979 | |
| 980 | |
| 981 | def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch.Tensor: |
| 982 | """Reshapes the query tensor for the specified batch size, so that |
| 983 | it has shape (batch_size, seq_len, num_heads, head_dim). |
| 984 | """ |
| 985 | assert query.dim() == 3, f"query must be 3D, got {query.dim()}D" |
| 986 | total_tokens = query.shape[0] |
| 987 | num_heads = query.shape[1] |
| 988 | head_dim = query.shape[2] |
| 989 | assert total_tokens % batch_size == 0, ( |
| 990 | f"{total_tokens=} is not divisible by {batch_size=}" |
| 991 | ) |
| 992 | seq_len = total_tokens // batch_size |
| 993 | return query.view(batch_size, seq_len, num_heads, head_dim) |
| 994 | |
| 995 | |
| 996 | def reshape_attn_output_for_spec_decode(attn_output: torch.Tensor) -> torch.Tensor: |
| 997 | """Reshapes the attention output tensor, so that |
| 998 | the batch_size and seq_len dimensions are combined. |
| 999 | """ |
| 1000 | if attn_output.dim() == 3: |
| 1001 | # Already in the correct shape |
| 1002 | return attn_output |
| 1003 | assert attn_output.dim() == 4, f"attn_output must be 4D, got {attn_output.dim()}D" |
| 1004 | total_tokens = attn_output.shape[0] * attn_output.shape[1] |
| 1005 | return attn_output.view(total_tokens, attn_output.shape[2], attn_output.shape[3]) |
| 1006 | |
| 1007 | |
| 1008 | @runtime_checkable |
| 1009 | class KVSharingFastPrefillMetadata(Protocol): |
| 1010 | logits_indices_padded: torch.Tensor | None = None |
| 1011 | num_logits_indices: int | None = None |
| 1012 | |
| 1013 | |
| 1014 | def create_fast_prefill_custom_backend( |
| 1015 | prefix: str, |
| 1016 | underlying_attn_backend: type[AttentionBackend], |
| 1017 | ) -> type[AttentionBackend]: |
| 1018 | underlying_builder = underlying_attn_backend.get_builder_cls() |
| 1019 | |
| 1020 | class FastPrefillAttentionBuilder(underlying_builder): # type: ignore |
| 1021 | def build( |
| 1022 | self, |
| 1023 | common_prefix_len: int, |
| 1024 | common_attn_metadata: CommonAttentionMetadata, |
| 1025 | fast_build: bool = False, |
| 1026 | ) -> AttentionMetadata: |
| 1027 | new_common_attn_metadata = ( |
| 1028 | make_kv_sharing_fast_prefill_common_attn_metadata(common_attn_metadata) |
| 1029 | ) |
| 1030 | metadata = super().build( |
| 1031 | common_prefix_len, new_common_attn_metadata, fast_build |
| 1032 | ) |
| 1033 | |
| 1034 | class KVSharingFastPrefillAttentionMetadata( |
| 1035 | metadata.__class__, # type: ignore |
| 1036 | KVSharingFastPrefillMetadata, |
| 1037 | ): |
| 1038 | def __init__(self, metadata, common_attn_metadata): |
| 1039 | # Shallow copy all fields in metadata cls |
| 1040 | for _field in fields(metadata.__class__): |
| 1041 | setattr(self, _field.name, getattr(metadata, _field.name)) |
| 1042 | |
| 1043 | self.logits_indices_padded = ( |
| 1044 | common_attn_metadata.logits_indices_padded |
| 1045 | ) |
| 1046 | self.num_logits_indices = common_attn_metadata.num_logits_indices |
| 1047 | self._attention_backend_variant = getattr( |
| 1048 | metadata, "_attention_backend_variant", 0 |
| 1049 | ) |
| 1050 | |
| 1051 | return KVSharingFastPrefillAttentionMetadata(metadata, common_attn_metadata) |
| 1052 | |
| 1053 | attn_backend = subclass_attention_backend( |
| 1054 | name_prefix=prefix, |
| 1055 | attention_backend_cls=underlying_attn_backend, |
| 1056 | builder_cls=FastPrefillAttentionBuilder, |
| 1057 | ) |
| 1058 | |
| 1059 | return attn_backend |
| 1060 | |
| 1061 | |
| 1062 | def compute_causal_conv1d_metadata( |
| 1063 | query_start_loc_p_cpu: torch.Tensor, *, device: torch.device |
| 1064 | ) -> tuple[dict[int, dict[str, Any]], torch.Tensor, torch.Tensor]: |
| 1065 | # Needed for causal_conv1d. Use the CPU query_start_loc to avoid DtoH sync. |
| 1066 | assert query_start_loc_p_cpu.device.type == "cpu" |
| 1067 | seqlens = query_start_loc_p_cpu.diff() |
| 1068 | nums_dict: dict[int, dict[str, Any]] = {} |
| 1069 | batch_ptr = None |
| 1070 | token_chunk_offset_ptr = None |
| 1071 | for BLOCK_M in [8]: # cover all BLOCK_M values |
| 1072 | nums = -(-seqlens // BLOCK_M) |
| 1073 | nums_dict[BLOCK_M] = {} |
| 1074 | nums_dict[BLOCK_M]["nums"] = nums |
| 1075 | nums_dict[BLOCK_M]["tot"] = nums.sum().item() |
| 1076 | mlist = np_to_pinned_tensor(np.repeat(np.arange(len(nums)), nums)) |
| 1077 | nums_dict[BLOCK_M]["mlist"] = mlist |
| 1078 | mlist_len = len(nums_dict[BLOCK_M]["mlist"]) |
| 1079 | nums_dict[BLOCK_M]["mlist_len"] = mlist_len |
| 1080 | MAX_NUM_PROGRAMS = max(1024, mlist_len) * 2 |
| 1081 | offsetlist = [] # type: ignore |
| 1082 | for idx, num in enumerate(nums): |
| 1083 | offsetlist.extend(range(num)) |
| 1084 | offsetlist = torch.tensor(offsetlist, dtype=torch.int32, pin_memory=PIN_MEMORY) |
| 1085 | nums_dict[BLOCK_M]["offsetlist"] = offsetlist |
| 1086 | |
| 1087 | if batch_ptr is None: |
| 1088 | # Update default value after class definition |
| 1089 | batch_ptr = torch.full( |
| 1090 | (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=device |
| 1091 | ) |
| 1092 | token_chunk_offset_ptr = torch.full( |
| 1093 | (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=device |
| 1094 | ) |
| 1095 | else: |
| 1096 | if batch_ptr.nelement() < MAX_NUM_PROGRAMS: |
| 1097 | batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID) |
| 1098 | assert token_chunk_offset_ptr is not None |
| 1099 | token_chunk_offset_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID) |
| 1100 | |
| 1101 | assert batch_ptr is not None |
| 1102 | batch_ptr[0:mlist_len].copy_(mlist, non_blocking=True) |
| 1103 | assert token_chunk_offset_ptr is not None |
| 1104 | token_chunk_offset_ptr[0:mlist_len].copy_(offsetlist, non_blocking=True) |
| 1105 | nums_dict[BLOCK_M]["batch_ptr"] = batch_ptr |
| 1106 | nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr |
| 1107 | |
| 1108 | return nums_dict, batch_ptr, token_chunk_offset_ptr |
| 1109 | |
| 1110 | |
| 1111 | def get_dcp_local_seq_lens( |
| 1112 | seq_lens: torch.Tensor, |
| 1113 | dcp_size: int = 1, |
| 1114 | dcp_rank: int | torch.Tensor | None = None, |
| 1115 | cp_kv_cache_interleave_size: int = 1, |
| 1116 | ) -> torch.Tensor: |
| 1117 | """While using dcp, kv_cache size stored on each rank may be different, |
| 1118 | use this function to calculate split decode seq_lens of each dcp rank. |
| 1119 | Only consider dcp now, we can extend the case of cp based on this. |
| 1120 | """ |
| 1121 | seq_lens_i32 = seq_lens.to(torch.int32) |
| 1122 | rank_offsets: int | torch.Tensor |
| 1123 | if dcp_rank is None: |
| 1124 | rank_offsets = torch.arange( |
| 1125 | dcp_size, |
| 1126 | dtype=torch.int32, |
| 1127 | device=seq_lens.device, |
| 1128 | ).view( |
| 1129 | *((1,) * seq_lens_i32.dim()), |
| 1130 | dcp_size, |
| 1131 | ) |
| 1132 | seq_lens_tiled = seq_lens_i32.unsqueeze(-1) |
| 1133 | elif isinstance(dcp_rank, torch.Tensor): |
| 1134 | assert dcp_rank.dtype == torch.int32 |
| 1135 | assert dcp_rank.device == seq_lens.device |
| 1136 | assert dcp_rank.numel() == 1 |
| 1137 | rank_offsets = dcp_rank |
| 1138 | seq_lens_tiled = seq_lens_i32 |
| 1139 | else: |
| 1140 | # Use the Python scalar directly to avoid a synchronizing H2D copy. |
| 1141 | rank_offsets = dcp_rank |
| 1142 | seq_lens_tiled = seq_lens_i32 |
| 1143 | base = ( |
| 1144 | seq_lens_tiled |
| 1145 | // cp_kv_cache_interleave_size |
| 1146 | // dcp_size |
| 1147 | * cp_kv_cache_interleave_size |
| 1148 | ) |
| 1149 | remainder = seq_lens_tiled - base * dcp_size |
| 1150 | remainder = torch.clip( |
| 1151 | remainder - rank_offsets * cp_kv_cache_interleave_size, |
| 1152 | 0, |
| 1153 | cp_kv_cache_interleave_size, |
| 1154 | ) |
| 1155 | dcp_local_seq_lens = base + remainder |
| 1156 | return dcp_local_seq_lens |
| 1157 | |
| 1158 | |
| 1159 | def mamba_get_block_table_tensor( |
| 1160 | block_table: torch.Tensor, |
| 1161 | seq_lens: torch.Tensor, |
| 1162 | kv_cache_spec: KVCacheSpec, |
| 1163 | mamba_cache_mode: str, |
| 1164 | ) -> torch.Tensor: |
| 1165 | """Get the block table tensor for mamba kernels from the input |
| 1166 | common_attn_metadata.block_table_tensor given different mamba cache modes. |
| 1167 | |
| 1168 | - "none": input (#requests, 1 + num_speculative_blocks); |
| 1169 | output (#requests, 1 + num_speculative_blocks). |
| 1170 | |
| 1171 | - "align": input (#requests, cdiv(max_model_len, block_size)); |
| 1172 | output (#requests, 1 + num_speculative_blocks), which are the last |
| 1173 | 1 + num_speculative_blocks of each request. |
| 1174 | """ |
| 1175 | if mamba_cache_mode == "none": |
| 1176 | return block_table |
| 1177 | else: |
| 1178 | assert isinstance(kv_cache_spec, MambaSpec) |
| 1179 | # NOTE: For 0-length requests in CUDA graph, use a start_index of 0 |
| 1180 | # to handle the invalid block table. |
| 1181 | start_indices = (seq_lens - 1) // kv_cache_spec.block_size |
| 1182 | start_indices.clamp_(min=0) |
| 1183 | # Use int32 for arithmetic to avoid dtype promotion overhead, |
| 1184 | # then convert to int64 for gather (which requires Long indices) |
| 1185 | offsets = torch.arange( |
| 1186 | 1 + kv_cache_spec.num_speculative_blocks, |
| 1187 | device=block_table.device, |
| 1188 | dtype=torch.int32, |
| 1189 | ) |
| 1190 | indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64) |
| 1191 | return torch.gather(block_table, 1, indices_to_gather) |