| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """KV cache helper for store.""" |
| 4 | |
| 5 | from collections.abc import Iterator |
| 6 | from dataclasses import dataclass |
| 7 | from typing import TYPE_CHECKING, Any, Literal, cast |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config import ( |
| 12 | VllmConfig, |
| 13 | get_current_vllm_config_or_none, |
| 14 | get_layers_from_vllm_config, |
| 15 | set_current_vllm_config, |
| 16 | ) |
| 17 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.utils.torch_utils import async_tensor_h2d |
| 22 | from vllm.v1.attention.backend import AttentionBackend |
| 23 | from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput |
| 24 | |
| 25 | if TYPE_CHECKING: |
| 26 | from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | EngineId = str |
| 31 | # block ids as returned by the hybrid KV cache manager. list[list[int]] are allow |
| 32 | # mutability and are for connector internal use only. |
| 33 | BlockIds = tuple[list[int], ...] | list[list[int]] |
| 34 | |
| 35 | |
| 36 | def clip_ssm_state_blocks(blocks: list[int], num_spec_blocks: int) -> list[int]: |
| 37 | """Drop speculative scratch slots and keep only the running-state slot.""" |
| 38 | if not blocks: |
| 39 | return blocks |
| 40 | # Keep at least the running-state slot, even before scratch allocation. |
| 41 | if num_scratch := min(num_spec_blocks, len(blocks) - 1): |
| 42 | blocks = blocks[:-num_scratch] |
| 43 | return blocks[-1:] |
| 44 | |
| 45 | |
| 46 | def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None): |
| 47 | # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is |
| 48 | # used for faster transfer. |
| 49 | if vllm_config is None: |
| 50 | vllm_config = get_current_vllm_config_or_none() |
| 51 | if vllm_config is None: |
| 52 | return None |
| 53 | kv_config = vllm_config.kv_transfer_config |
| 54 | if kv_config is not None: |
| 55 | connector_cls = KVConnectorFactory.get_connector_class(kv_config) |
| 56 | required_kvcache_layout = connector_cls.get_required_kvcache_layout(vllm_config) |
| 57 | if required_kvcache_layout is not None: |
| 58 | return required_kvcache_layout |
| 59 | return None |
| 60 | |
| 61 | |
| 62 | class KVOutputAggregator: |
| 63 | """Utility class to aggregate the output of all workers into a single |
| 64 | output corresponding to Rank 0 for scheduler.""" |
| 65 | |
| 66 | def __init__(self, expected_finished_count: int): |
| 67 | # Complete transfer tracker. Used to track finished requests |
| 68 | # [req_id -> n_remaining_workers] |
| 69 | self._recv_remaining_count = dict[str, int]() |
| 70 | self._send_remaining_count = dict[str, int]() |
| 71 | self._failed_recving_pending = set[str]() |
| 72 | self._expected_finished_count = expected_finished_count |
| 73 | |
| 74 | @classmethod |
| 75 | def from_connector(cls, connector: "KVConnectorBase", world_size: int): |
| 76 | return cls(connector.get_finished_count() or world_size) |
| 77 | |
| 78 | def aggregate( |
| 79 | self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0 |
| 80 | ) -> ModelRunnerOutput | None: |
| 81 | if not outputs[output_rank]: |
| 82 | return None |
| 83 | |
| 84 | # Aggregate kv_connector_output from all workers |
| 85 | |
| 86 | def update_finished_set( |
| 87 | req_ids: set[str] | None, |
| 88 | remaining_count_dict: dict[str, int], |
| 89 | finished_set: set[str], |
| 90 | ) -> None: |
| 91 | for req_id in req_ids or (): |
| 92 | remaining_count = remaining_count_dict.get( |
| 93 | req_id, self._expected_finished_count |
| 94 | ) |
| 95 | remaining_count_dict[req_id] = remaining_count - 1 |
| 96 | if remaining_count_dict[req_id] == 0: |
| 97 | finished_set.add(req_id) |
| 98 | del remaining_count_dict[req_id] |
| 99 | |
| 100 | finished_sending = set[str]() |
| 101 | finished_recving = set[str]() |
| 102 | aggregated_kv_connector_stats = None |
| 103 | aggregated_kv_connector_worker_meta = None |
| 104 | combined_kv_cache_events = None |
| 105 | invalid_block_ids = set[int]() |
| 106 | for model_runner_output in outputs: |
| 107 | assert model_runner_output is not None |
| 108 | kv_output = model_runner_output.kv_connector_output |
| 109 | if not kv_output: |
| 110 | continue |
| 111 | # Allow the worker to dynamically update the expected number of |
| 112 | # finished sending/recving for new requests. |
| 113 | if ( |
| 114 | kv_output.expected_finished_count > 0 |
| 115 | and kv_output.expected_finished_count != self._expected_finished_count |
| 116 | ): |
| 117 | logger.debug( |
| 118 | "Expected finished requests updated from %d to %d", |
| 119 | self._expected_finished_count, |
| 120 | kv_output.expected_finished_count, |
| 121 | ) |
| 122 | self._expected_finished_count = kv_output.expected_finished_count |
| 123 | |
| 124 | update_finished_set( |
| 125 | kv_output.finished_sending, self._send_remaining_count, finished_sending |
| 126 | ) |
| 127 | update_finished_set( |
| 128 | kv_output.finished_recving, self._recv_remaining_count, finished_recving |
| 129 | ) |
| 130 | |
| 131 | # Aggregate kv_connector_stats from all workers. |
| 132 | if aggregated_kv_connector_stats is None: |
| 133 | # Use the first worker's kv_connector_stats as accumulator. |
| 134 | aggregated_kv_connector_stats = kv_output.kv_connector_stats |
| 135 | elif kv_connector_stats := kv_output.kv_connector_stats: |
| 136 | assert isinstance( |
| 137 | aggregated_kv_connector_stats, type(kv_connector_stats) |
| 138 | ) |
| 139 | aggregated_kv_connector_stats = aggregated_kv_connector_stats.aggregate( |
| 140 | kv_connector_stats |
| 141 | ) |
| 142 | |
| 143 | # Aggregate kv_connector_worker_meta from all workers. |
| 144 | if aggregated_kv_connector_worker_meta is None: |
| 145 | # Use the first worker's kv_connector_worker_meta as accumulator. |
| 146 | aggregated_kv_connector_worker_meta = kv_output.kv_connector_worker_meta |
| 147 | elif kv_connector_worker_meta := kv_output.kv_connector_worker_meta: |
| 148 | aggregated_kv_connector_worker_meta = ( |
| 149 | aggregated_kv_connector_worker_meta.aggregate( |
| 150 | kv_connector_worker_meta |
| 151 | ) |
| 152 | ) |
| 153 | |
| 154 | # Combine kv_cache_events from all workers. |
| 155 | if combined_kv_cache_events is None: |
| 156 | # Use the first worker's kv_cache events as start event list. |
| 157 | combined_kv_cache_events = kv_output.kv_cache_events |
| 158 | elif kv_cache_events := kv_output.kv_cache_events: |
| 159 | assert isinstance( |
| 160 | combined_kv_cache_events, |
| 161 | type(kv_cache_events), |
| 162 | ) |
| 163 | worker_kv_cache_events = kv_cache_events.get_all_events() |
| 164 | combined_kv_cache_events.add_events(worker_kv_cache_events) |
| 165 | combined_kv_cache_events.increment_workers(1) |
| 166 | |
| 167 | invalid_block_ids |= kv_output.invalid_block_ids |
| 168 | self._failed_recving_pending |= kv_output.failed_recving |
| 169 | |
| 170 | failed_recving = self._failed_recving_pending & finished_recving |
| 171 | self._failed_recving_pending -= failed_recving |
| 172 | |
| 173 | # select output of the worker specified by output_rank |
| 174 | output = outputs[output_rank] |
| 175 | |
| 176 | assert output is not None |
| 177 | output.kv_connector_output = KVConnectorOutput( |
| 178 | finished_sending=finished_sending or None, |
| 179 | finished_recving=finished_recving or None, |
| 180 | kv_connector_stats=aggregated_kv_connector_stats or None, |
| 181 | kv_cache_events=combined_kv_cache_events or None, |
| 182 | kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None, |
| 183 | invalid_block_ids=invalid_block_ids, |
| 184 | failed_recving=failed_recving, |
| 185 | expected_finished_count=self._expected_finished_count, |
| 186 | ) |
| 187 | |
| 188 | return output |
| 189 | |
| 190 | |
| 191 | def _make_src_and_dst_indices( |
| 192 | src_block_ids: list[int], |
| 193 | dst_block_ids: list[int], |
| 194 | src_device: torch.device | str, |
| 195 | dst_device: torch.device | str, |
| 196 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 197 | def _to(block_ids: list[int], device: torch.device | str) -> torch.Tensor: |
| 198 | device = torch.device(device) if isinstance(device, str) else device |
| 199 | if device.type == "cpu": |
| 200 | return torch.tensor(block_ids, dtype=torch.int64, device=device) |
| 201 | return async_tensor_h2d(block_ids, dtype=torch.int64, device=device) |
| 202 | |
| 203 | return _to(src_block_ids, src_device), _to(dst_block_ids, dst_device) |
| 204 | |
| 205 | |
| 206 | def copy_kv_blocks( |
| 207 | src_kv_caches: dict[str, torch.Tensor], |
| 208 | dst_kv_caches: dict[str, torch.Tensor], |
| 209 | src_block_ids: list[int], |
| 210 | dst_block_ids: list[int], |
| 211 | direction: Literal["h2d", "d2h"], |
| 212 | ) -> None: |
| 213 | """Copy kv blocks between different buffers.""" |
| 214 | if ( |
| 215 | not src_kv_caches |
| 216 | or not dst_kv_caches |
| 217 | or not src_block_ids |
| 218 | or not dst_block_ids |
| 219 | or len(src_block_ids) != len(dst_block_ids) |
| 220 | ): |
| 221 | return |
| 222 | |
| 223 | src_device = next(iter(src_kv_caches.values())).device |
| 224 | dst_device = next(iter(dst_kv_caches.values())).device |
| 225 | |
| 226 | src_indices, dst_indices = _make_src_and_dst_indices( |
| 227 | src_block_ids=src_block_ids, |
| 228 | dst_block_ids=dst_block_ids, |
| 229 | src_device=src_device, |
| 230 | dst_device=dst_device, |
| 231 | ) |
| 232 | |
| 233 | if direction == "h2d": |
| 234 | copy_fn = current_platform.insert_blocks_to_device |
| 235 | else: |
| 236 | copy_fn = current_platform.swap_out_blocks_to_host |
| 237 | for layer_name in src_kv_caches: |
| 238 | src_tensor = src_kv_caches[layer_name] |
| 239 | dst_tensor = dst_kv_caches[layer_name] |
| 240 | copy_fn(src_tensor, dst_tensor, src_indices, dst_indices) |
| 241 | |
| 242 | |
| 243 | def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): |
| 244 | """Transforms the layout of received KV cache blocks to the local block_size. |
| 245 | (Only works for local blocksize > remote blocksize) |
| 246 | |
| 247 | Example: |
| 248 | local blocksize = 16 tokens, remote blocksize = 4 tokens |
| 249 | local block[0] = remote block[0, 1, 2, 3] |
| 250 | remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|... |
| 251 | local is |h0-b0..................|h1-b0..................|... |
| 252 | permute is to: |
| 253 | 1. view => view remote as n_blocks * remote_shape(H,remoteN,D) |
| 254 | 2. permute => (H, nblocks, remoteN, D) |
| 255 | 3. flatten => (H, localN, D) |
| 256 | |
| 257 | """ |
| 258 | blocks_to_update = cache.index_select(0, indices) |
| 259 | # use physical order |
| 260 | blocks_to_update = blocks_to_update.permute(0, 2, 1, 3) |
| 261 | n_kv_heads, block_size, head_size = blocks_to_update.shape[1:] |
| 262 | remote_block_size = block_size // block_size_ratio |
| 263 | n_blocks = block_size_ratio |
| 264 | |
| 265 | permuted_blocks = ( |
| 266 | blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size) |
| 267 | .permute(0, 2, 1, 3, 4) |
| 268 | .flatten(2, 3) |
| 269 | ) |
| 270 | permuted_blocks = permuted_blocks.permute(0, 2, 1, 3) |
| 271 | cache.index_copy_(0, indices, permuted_blocks) |
| 272 | |
| 273 | |
| 274 | def kv_postprocess_layout_on_receive(cache, indices): |
| 275 | """Transforms the layout of received KV cache blocks to the local format. |
| 276 | |
| 277 | This method corrects layout mismatches from direct memory copies by |
| 278 | permuting the tensor dimensions. |
| 279 | |
| 280 | 4D cache: |
| 281 | - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]` |
| 282 | - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]` |
| 283 | 5D cache: |
| 284 | - **Source Layout:** `[num_blocks, kv_dim, n_kv_head, block_size, head_dim]` |
| 285 | - **Target Layout:** `[num_blocks, kv_dim, block_size, n_kv_head, head_dim]` |
| 286 | |
| 287 | Implementation: |
| 288 | - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout |
| 289 | - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size |
| 290 | - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back |
| 291 | |
| 292 | """ |
| 293 | blocks_to_update = cache.index_select(0, indices) |
| 294 | target_shape = list(blocks_to_update.shape) |
| 295 | target_shape[0] = -1 |
| 296 | inv_order = [0, 1, 3, 2, 4] if blocks_to_update.ndim == 5 else [0, 2, 1, 3] |
| 297 | src_shape = tuple(target_shape[i] for i in inv_order) |
| 298 | blocks_to_update = cache.index_select(0, indices) |
| 299 | permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order) |
| 300 | cache.index_copy_(0, indices, permuted_blocks) |
| 301 | |
| 302 | |
| 303 | def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio): |
| 304 | """Transforms the layout of received KV cache to the local block_size and LBHNC. |
| 305 | (Only works for local blocksize > remote blocksize) |
| 306 | |
| 307 | prefill is LBHNC, smaller block_size |
| 308 | decode(local) is LBNHC, larger block_size |
| 309 | """ |
| 310 | blocks_to_update = cache.index_select(0, indices) |
| 311 | |
| 312 | block_size, n_kv_heads, head_size = blocks_to_update.shape[1:] |
| 313 | remote_block_size = block_size // block_size_ratio |
| 314 | n_blocks = block_size_ratio |
| 315 | |
| 316 | permuted_blocks = ( |
| 317 | blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size) |
| 318 | .permute(0, 1, 3, 2, 4) |
| 319 | .flatten(1, 2) |
| 320 | ) |
| 321 | cache.index_copy_(0, indices, permuted_blocks) |
| 322 | |
| 323 | |
| 324 | def yield_req_data( |
| 325 | scheduler_output, |
| 326 | ) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]: |
| 327 | """Yields: |
| 328 | (req_id, new_block_id_groups, preempted) |
| 329 | |
| 330 | """ |
| 331 | # new requests |
| 332 | for req_data in scheduler_output.scheduled_new_reqs: |
| 333 | yield req_data.req_id, req_data.block_ids, False |
| 334 | |
| 335 | # cached requests |
| 336 | cached_reqs = scheduler_output.scheduled_cached_reqs |
| 337 | yield from zip( |
| 338 | cached_reqs.req_ids, |
| 339 | cached_reqs.new_block_ids, |
| 340 | (req_id in cached_reqs.resumed_req_ids for req_id in cached_reqs.req_ids), |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | def get_current_attn_backends( |
| 345 | vllm_config: VllmConfig, layer_names: list[str] | None = None |
| 346 | ) -> list[type[AttentionBackend]]: |
| 347 | """Get all distinct attention backends for the given layers. |
| 348 | |
| 349 | Args: |
| 350 | vllm_config: The current vLLM configuration. |
| 351 | layer_names: Optional list of layer names to scope the lookup. |
| 352 | When None, all attention layers are considered. |
| 353 | |
| 354 | Returns: |
| 355 | Deduplicated list of attention backend classes. |
| 356 | |
| 357 | """ |
| 358 | layer_type = cast(type[Any], AttentionLayerBase) |
| 359 | layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names) |
| 360 | if layers: |
| 361 | seen: dict[str, type[AttentionBackend]] = {} |
| 362 | for layer in layers.values(): |
| 363 | backend = layer.get_attn_backend() |
| 364 | seen[backend.full_cls_name()] = backend |
| 365 | return list(seen.values()) |
| 366 | |
| 367 | # Fallback for tests, when static_forward_context is empty. |
| 368 | logger.debug( |
| 369 | "No layers found in the vLLM config. Falling back to default attention backend." |
| 370 | ) |
| 371 | from vllm.v1.attention.selector import get_attn_backend |
| 372 | |
| 373 | with set_current_vllm_config(vllm_config): |
| 374 | return [ |
| 375 | get_attn_backend( |
| 376 | head_size=vllm_config.model_config.get_head_size(), |
| 377 | dtype=vllm_config.model_config.dtype, |
| 378 | kv_cache_dtype=vllm_config.cache_config.cache_dtype, |
| 379 | use_mla=vllm_config.model_config.use_mla, |
| 380 | ) |
| 381 | ] |
| 382 | |
| 383 | |
| 384 | def get_current_attn_backend( |
| 385 | vllm_config: VllmConfig, layer_names: list[str] | None = None |
| 386 | ) -> type[AttentionBackend]: |
| 387 | """Get the first attention backend for the given layers.""" |
| 388 | return get_current_attn_backends(vllm_config, layer_names)[0] |
| 389 | |
| 390 | |
| 391 | # ---- Per-engine transfer info ---- |
| 392 | |
| 393 | |
| 394 | @dataclass(frozen=True) |
| 395 | class EngineTransferInfo: |
| 396 | """Common per-remote-engine transfer state, computed at handshake. |
| 397 | |
| 398 | Stored per ``(engine_id, pp_rank)`` inside ``TransferTopology._engines``. |
| 399 | """ |
| 400 | |
| 401 | remote_tp_size: int |
| 402 | |
| 403 | remote_block_len: int |
| 404 | """Block length (bytes)""" |
| 405 | |
| 406 | remote_block_size: int |
| 407 | """Tokens per block.""" |
| 408 | |
| 409 | remote_physical_blocks_per_logical: int |
| 410 | """Physical blocks per logical block.""" |
| 411 | |
| 412 | remote_pp_rank: int = 0 |
| 413 | """Remote producer PP rank for this engine.""" |
| 414 | |
| 415 | start_layer: int = 0 |
| 416 | """Global index of the first layer owned by this PP rank.""" |
| 417 | |
| 418 | end_layer: int = 0 |
| 419 | """Exclusive global index after the last layer owned by this PP rank.""" |
| 420 | |
| 421 | remote_dcp_size: int = 1 |
| 422 | """Remote decode context parallel size.""" |
| 423 | |
| 424 | |
| 425 | # ---- Transfer topology ---- |
| 426 | |
| 427 | |
| 428 | @dataclass |
| 429 | class TransferTopology: |
| 430 | """Single source of truth for local TP identity and per-engine remote info.""" |
| 431 | |
| 432 | tp_rank: int |
| 433 | tp_size: int |
| 434 | block_size: int |
| 435 | engine_id: EngineId |
| 436 | is_mla: bool |
| 437 | is_mamba: bool |
| 438 | total_num_kv_heads: int |
| 439 | attn_backends: list[type[AttentionBackend]] |
| 440 | dcp_size: int = 1 |
| 441 | tensor_shape: torch.Size | None = None |
| 442 | |
| 443 | def __post_init__(self): |
| 444 | self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size) |
| 445 | |
| 446 | self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {} |
| 447 | |
| 448 | # ============================================================ |
| 449 | # Engine registration |
| 450 | # ============================================================ |
| 451 | |
| 452 | def register_remote_engine( |
| 453 | self, |
| 454 | remote_engine_id: EngineId, |
| 455 | info: EngineTransferInfo, |
| 456 | ) -> EngineTransferInfo: |
| 457 | """Register a remote engine, unifying worker dicts state. |
| 458 | |
| 459 | The caller (worker) is responsible for computing the info via |
| 460 | the transfer policy. This method only stores and deduplicates. |
| 461 | """ |
| 462 | assert remote_engine_id != self.engine_id, ( |
| 463 | f"Cannot register local engine {self.engine_id} as remote. " |
| 464 | f"Local identity is set via __init__ params." |
| 465 | ) |
| 466 | engine_key = (remote_engine_id, info.remote_pp_rank) |
| 467 | if engine_key in self._engines: |
| 468 | return self._engines[engine_key] |
| 469 | self._engines[engine_key] = info |
| 470 | return info |
| 471 | |
| 472 | def get_engine_info( |
| 473 | self, remote_engine_id: EngineId, remote_pp_rank: int = 0 |
| 474 | ) -> EngineTransferInfo: |
| 475 | return self._engines[(remote_engine_id, remote_pp_rank)] |
| 476 | |
| 477 | def unregister_remote_engine(self, remote_engine_id: EngineId) -> None: |
| 478 | # Remove all pp_rank entries for the remote engine. |
| 479 | for key in [k for k in self._engines if k[0] == remote_engine_id]: |
| 480 | del self._engines[key] |
| 481 | |
| 482 | @property |
| 483 | def dcp_rank(self) -> int: |
| 484 | """This rank's DCP coverage rank. |
| 485 | |
| 486 | with ``dcp_size in (1, tp_size)`` enforced at the connector boundary, a |
| 487 | rank's DCP identity is always exactly ``tp_rank % dcp_size``. |
| 488 | """ |
| 489 | return self.tp_rank % self.dcp_size |
| 490 | |
| 491 | # ============================================================ |
| 492 | # Common methods |
| 493 | # ============================================================ |
| 494 | |
| 495 | def tp_ratio(self, remote_tp_size: int) -> int: |
| 496 | """Calculate the tensor parallel ratio between local and remote TP. |
| 497 | |
| 498 | Positive when local_tp >= remote_tp (local workers read from the |
| 499 | same remote worker in groups of size ``tp_ratio``). Negative when |
| 500 | remote_tp > local_tp (ratio is flipped). |
| 501 | """ |
| 502 | if self.tp_size >= remote_tp_size: |
| 503 | assert self.tp_size % remote_tp_size == 0, ( |
| 504 | f"Local tensor parallel size {self.tp_size} is not divisible " |
| 505 | f"by remote tensor parallel size {remote_tp_size}." |
| 506 | ) |
| 507 | return self.tp_size // remote_tp_size |
| 508 | assert remote_tp_size % self.tp_size == 0, ( |
| 509 | f"Remote tensor parallel size {remote_tp_size} is not divisible " |
| 510 | f"by local tensor parallel size {self.tp_size}." |
| 511 | ) |
| 512 | return -(remote_tp_size // self.tp_size) |
| 513 | |
| 514 | def block_size_ratio(self, remote_block_size: int) -> int: |
| 515 | """Calculate the block size ratio between local and remote.""" |
| 516 | assert self.block_size % remote_block_size == 0, ( |
| 517 | f"Local block size {self.block_size} is not divisible " |
| 518 | f"by remote block size {remote_block_size} or vice versa." |
| 519 | ) |
| 520 | return self.block_size // remote_block_size |
| 521 | |
| 522 | def is_kv_replicated( |
| 523 | self, remote_engine_id: EngineId, remote_pp_rank: int = 0 |
| 524 | ) -> bool: |
| 525 | """Whether the KV cache is replicated across TP workers due to the |
| 526 | number of TP workers being greater than the number of KV heads. |
| 527 | """ |
| 528 | return ( |
| 529 | self._engines[(remote_engine_id, remote_pp_rank)].remote_tp_size |
| 530 | > self.total_num_kv_heads |
| 531 | ) |
| 532 | |
| 533 | def replicates_kv_cache( |
| 534 | self, remote_engine_id: EngineId, remote_pp_rank: int = 0 |
| 535 | ) -> bool: |
| 536 | # MLA is always replicated as the hidden dim can't be split. |
| 537 | return self.is_mla or self.is_kv_replicated(remote_engine_id, remote_pp_rank) |
| 538 | |
| 539 | @property |
| 540 | def local_replicates_kv_cache(self) -> bool: |
| 541 | """Whether the local engine's KV cache is replicated.""" |
| 542 | return self.is_mla or self.tp_size > self.total_num_kv_heads |
| 543 | |
| 544 | def dcp_source_ranks(self, remote_tp_size: int, remote_dcp_size: int) -> list[int]: |
| 545 | """Remote ranks whose DCP slice overlaps mine (MLA, ``remote_dcp_size > 1``). |
| 546 | |
| 547 | Shared by ``handshake_target_ranks`` (who to query metadata from) and |
| 548 | ``compute_tp_mapping`` (who to actually read from) — for MLA the two |
| 549 | questions have the identical answer, since DCP sharding is the only |
| 550 | thing keeping a remote rank from being interchangeable with any other. |
| 551 | """ |
| 552 | local_dcp_size = self.dcp_size |
| 553 | local_dcp_rank = self.dcp_rank |
| 554 | if local_dcp_size <= remote_dcp_size: |
| 555 | # Keep every remote rank whose slice sits inside mine. When |
| 556 | # local_dcp_size == 1 (replicated locally), local_dcp_rank == 0 reduces to |
| 557 | # every remote rank, since no single one holds the whole sequence. |
| 558 | return [ |
| 559 | r for r in range(remote_tp_size) if r % local_dcp_size == local_dcp_rank |
| 560 | ] |
| 561 | # Local finer-grained: exactly one remote rank covers my whole slice |
| 562 | return [local_dcp_rank % remote_dcp_size] |
| 563 | |
| 564 | def handshake_target_ranks( |
| 565 | self, remote_tp_size: int, remote_dcp_size: int = 1 |
| 566 | ) -> list[int]: |
| 567 | """Pre-registration: compute which remote TP ranks to handshake with. |
| 568 | |
| 569 | Pure math based on local/remote TP (and DCP, when the remote shards |
| 570 | its KV cache) sizes — does not require the remote engine to be |
| 571 | registered yet. |
| 572 | |
| 573 | DCP support is scoped to ``dcp_size in (1, tp_size)`` on each side |
| 574 | and DCP sizes that divide one another: neither side ever has a |
| 575 | partially-duplicated, partially-sharded KV cache. When the |
| 576 | remote is not sharded (``remote_dcp_size == 1``) this reduces |
| 577 | exactly to the DTP>=PTP case, since a sharded local side |
| 578 | already has ``tp_size == dcp_size``. |
| 579 | """ |
| 580 | if remote_dcp_size > 1: |
| 581 | return self.dcp_source_ranks(remote_tp_size, remote_dcp_size) |
| 582 | |
| 583 | tp_ratio = self.tp_ratio(remote_tp_size) |
| 584 | if tp_ratio > 0: |
| 585 | return [self.tp_rank // tp_ratio] |
| 586 | abs_ratio = -tp_ratio |
| 587 | return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)] |
| 588 | |
| 589 | def dcp_consumer_count(self, remote_tp_size: int, remote_dcp_size: int) -> int: |
| 590 | """How many local ranks (in aggregate) read from a given remote rank. |
| 591 | |
| 592 | Used by the producer side to know how many reader notifications to |
| 593 | wait for before freeing a request's blocks. Reuses ``tp_ratio`` |
| 594 | whenever the remote isn't sharded — a sharded local side already |
| 595 | has ``tp_size == dcp_size``, so the existing TP-ratio formula is |
| 596 | already correct there unmodified. |
| 597 | """ |
| 598 | if remote_dcp_size > 1: |
| 599 | if self.dcp_size == 1: |
| 600 | # Replicated locally: every local rank reads every shard. |
| 601 | return self.tp_size |
| 602 | # Both sharded, different degrees. |
| 603 | return max(1, self.dcp_size // remote_dcp_size) |
| 604 | # Remote replicated: `tp_ratio` local ranks share each remote rank |
| 605 | # when local_tp >= remote_tp, else each remote rank has one reader. |
| 606 | return max(1, self.tp_ratio(remote_tp_size)) |
| 607 | |
| 608 | def target_remote_ranks( |
| 609 | self, remote_engine_id: EngineId, remote_pp_rank: int = 0 |
| 610 | ) -> list[int]: |
| 611 | """Get the remote TP rank(s) that the current local TP rank will |
| 612 | read from. When remote tp_size > local tp_size, reads from |
| 613 | multiple remote ranks. |
| 614 | """ |
| 615 | info = self._engines[(remote_engine_id, remote_pp_rank)] |
| 616 | tp_ratio = self.tp_ratio(info.remote_tp_size) |
| 617 | if tp_ratio > 0: |
| 618 | return [self.tp_rank // tp_ratio] |
| 619 | # remote TP > local TP: read from |tp_ratio| remote workers |
| 620 | abs_ratio = -tp_ratio |
| 621 | return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)] |
| 622 | |
| 623 | def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str: |
| 624 | """One-line summary of transfer config for logging.""" |
| 625 | info = self._engines[(remote_engine_id, remote_pp_rank)] |
| 626 | return ( |
| 627 | f"TransferTopology(" |
| 628 | f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, " |
| 629 | f"num_kv_heads={self.total_num_kv_heads if not self.is_mla else 1}, " |
| 630 | f"local_tp={self.tp_size}, " |
| 631 | f"remote_tp={info.remote_tp_size}, " |
| 632 | f"remote_pp={remote_pp_rank}, " |
| 633 | f"local_dcp={self.dcp_size}, " |
| 634 | f"remote_dcp={info.remote_dcp_size}, " |
| 635 | f"local_rank={self.tp_rank}, " |
| 636 | f"remote_block_len={info.remote_block_len})" |
| 637 | ) |