| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Translate vLLM KV cache metadata for native offloading backends.""" |
| 4 | |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | from vllm.utils.math_utils import round_up |
| 8 | from vllm.v1.core.kv_cache_utils import ( |
| 9 | resolve_dcp_kv_block_size, |
| 10 | resolve_kv_cache_block_sizes, |
| 11 | ) |
| 12 | from vllm.v1.kv_cache_interface import ( |
| 13 | AttentionSpec, |
| 14 | FullAttentionSpec, |
| 15 | KVCacheGroupRole, |
| 16 | KVCacheSpec, |
| 17 | MLAAttentionSpec, |
| 18 | SlidingWindowMLASpec, |
| 19 | SlidingWindowSpec, |
| 20 | UniformTypeKVCacheSpecs, |
| 21 | iter_layer_specs, |
| 22 | ) |
| 23 | from vllm.v1.kv_offload.config import ( |
| 24 | OffloadingCacheConfig, |
| 25 | OffloadingConfig, |
| 26 | OffloadingGroupConfig, |
| 27 | OffloadingModelConfig, |
| 28 | OffloadingParallelConfig, |
| 29 | ) |
| 30 | |
| 31 | if TYPE_CHECKING: |
| 32 | from vllm.config import VllmConfig |
| 33 | from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec |
| 34 | |
| 35 | |
| 36 | def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]: |
| 37 | if kv_cache_config.hisparse_host_num_blocks is None: |
| 38 | return kv_cache_config.prefix_cacheable_group_ids |
| 39 | return tuple( |
| 40 | group_id |
| 41 | for group_id, group in enumerate(kv_cache_config.kv_cache_groups) |
| 42 | if group.role is KVCacheGroupRole.HISPARSE_INDEXER |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | def _group_kv_bytes_per_block(group: "KVCacheGroupSpec") -> int: |
| 47 | """Return the physical bytes occupied by one block of a cache group. |
| 48 | |
| 49 | Worker configs may retain ``UniformTypeKVCacheSpecs`` while scheduler |
| 50 | configs flatten that wrapper to one representative per-layer spec. Keep |
| 51 | the result invariant across those two representations. |
| 52 | """ |
| 53 | spec = group.kv_cache_spec |
| 54 | if isinstance(spec, UniformTypeKVCacheSpecs): |
| 55 | return spec.page_size_bytes |
| 56 | return spec.page_size_bytes * len(group.layer_names) |
| 57 | |
| 58 | |
| 59 | def build_offloading_config( |
| 60 | vllm_config: "VllmConfig", |
| 61 | kv_cache_config: "KVCacheConfig", |
| 62 | ) -> OffloadingConfig: |
| 63 | """Translate vLLM configuration into the native offloading boundary.""" |
| 64 | kv_transfer_config = vllm_config.kv_transfer_config |
| 65 | assert kv_transfer_config is not None |
| 66 | extra_config = kv_transfer_config.kv_connector_extra_config |
| 67 | assert kv_transfer_config.engine_id is not None |
| 68 | engine_id = kv_transfer_config.engine_id |
| 69 | |
| 70 | parallel_config = vllm_config.parallel_config |
| 71 | selected_groups = tuple( |
| 72 | (group_id, kv_cache_config.kv_cache_groups[group_id]) |
| 73 | for group_id in get_offloading_group_ids(kv_cache_config) |
| 74 | ) |
| 75 | if not selected_groups: |
| 76 | raise ValueError("KV offloading found no eligible cache groups.") |
| 77 | groups = tuple( |
| 78 | OffloadingGroupConfig( |
| 79 | group_id=group_id, |
| 80 | tokens_per_block=resolve_dcp_kv_block_size( |
| 81 | group.kv_cache_spec, |
| 82 | parallel_config.decode_context_parallel_size, |
| 83 | ), |
| 84 | layer_names=tuple(group.layer_names), |
| 85 | ) |
| 86 | for group_id, group in selected_groups |
| 87 | ) |
| 88 | |
| 89 | _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config) |
| 90 | for group in groups: |
| 91 | assert group.tokens_per_block % tokens_per_hash == 0, ( |
| 92 | f"tokens_per_block={group.tokens_per_block} not divisible by " |
| 93 | f"tokens_per_hash={tokens_per_hash}. " |
| 94 | f"Hybrid models (e.g. Mamba+Attention) need " |
| 95 | f"--enable-prefix-caching to align block sizes." |
| 96 | ) |
| 97 | |
| 98 | blocks_per_chunk = 1 |
| 99 | blocks_per_chunk_config = extra_config.get("blocks_per_chunk") |
| 100 | tokens_per_chunk = extra_config.get("block_size") |
| 101 | |
| 102 | if blocks_per_chunk_config is not None and tokens_per_chunk is not None: |
| 103 | raise ValueError( |
| 104 | "Specify only one of 'block_size' or 'blocks_per_chunk' " |
| 105 | "in kv_connector_extra_config." |
| 106 | ) |
| 107 | |
| 108 | if blocks_per_chunk_config is not None: |
| 109 | blocks_per_chunk = int(blocks_per_chunk_config) |
| 110 | |
| 111 | if blocks_per_chunk <= 0: |
| 112 | raise ValueError("'blocks_per_chunk' must be greater than 0.") |
| 113 | |
| 114 | elif tokens_per_chunk is not None: |
| 115 | tokens_per_chunk_int = int(tokens_per_chunk) |
| 116 | |
| 117 | unique_tokens_per_block = {group.tokens_per_block for group in groups} |
| 118 | |
| 119 | assert len(unique_tokens_per_block) == 1, ( |
| 120 | "If 'block_size' is specified in kv_connector_extra_config, " |
| 121 | "there must be at least one KV cache group, " |
| 122 | "and all groups must have the same block size." |
| 123 | ) |
| 124 | |
| 125 | tokens_per_block = unique_tokens_per_block.pop() |
| 126 | if tokens_per_chunk_int % tokens_per_block == 0: |
| 127 | blocks_per_chunk = tokens_per_chunk_int // tokens_per_block |
| 128 | else: |
| 129 | raise ValueError( |
| 130 | f"'block_size'={tokens_per_chunk_int} in kv_connector_extra_config " |
| 131 | f"must be a multiple of the GPU KV cache block size " |
| 132 | f"({tokens_per_block} tokens). Use " |
| 133 | f"{round_up(tokens_per_chunk_int, tokens_per_block)} instead, or set " |
| 134 | f"'blocks_per_chunk' to express the chunk size in blocks." |
| 135 | ) |
| 136 | |
| 137 | worker_kv_bytes_per_block = 0 |
| 138 | if ( |
| 139 | kv_cache_config.hisparse_host_num_blocks is None |
| 140 | and kv_cache_config.num_blocks > 0 |
| 141 | and kv_cache_config.kv_cache_tensors |
| 142 | ): |
| 143 | # Scratch filtering must preserve the scheduler/worker allocation stride. |
| 144 | # Every KVCacheTensor describes placement within the same backing allocation, |
| 145 | # so its size is the total, not a per-tensor share. |
| 146 | total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size |
| 147 | worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks |
| 148 | elif kv_cache_config.num_blocks > 0: |
| 149 | worker_kv_bytes_per_block = sum( |
| 150 | _group_kv_bytes_per_block(group) for _, group in selected_groups |
| 151 | ) |
| 152 | |
| 153 | single_group_spec = ( |
| 154 | kv_cache_config.kv_cache_groups[0].kv_cache_spec |
| 155 | if len(kv_cache_config.kv_cache_groups) == 1 |
| 156 | else None |
| 157 | ) |
| 158 | replicated_layout = ( |
| 159 | vllm_config.model_config.use_mla |
| 160 | and parallel_config.tensor_parallel_size > 1 |
| 161 | and kv_cache_config.kv_tp_replicas == parallel_config.tensor_parallel_size |
| 162 | and worker_kv_bytes_per_block > 0 |
| 163 | # Safe MVP boundary: TP-only, no other parallel axes. |
| 164 | and parallel_config.pipeline_parallel_size == 1 |
| 165 | and parallel_config.prefill_context_parallel_size == 1 |
| 166 | and parallel_config.decode_context_parallel_size == 1 |
| 167 | and parallel_config.world_size == parallel_config.tensor_parallel_size |
| 168 | # Shared /dev/shm mmap layout is single-node mp only. |
| 169 | and parallel_config.distributed_executor_backend == "mp" |
| 170 | and parallel_config.nnodes_within_dp == 1 |
| 171 | ) |
| 172 | |
| 173 | canonical_layout = bool(extra_config.get("canonical_layout", False)) |
| 174 | |
| 175 | # Only a single non-MLA full-attention group with genuinely head-sharded |
| 176 | # pages is parallelism-invariant: replicated latent or GQA heads, |
| 177 | # per-token-head scales, CP token sharding, and the V2 model runner's |
| 178 | # layout are all excluded. |
| 179 | is_parallelism_agnostic = ( |
| 180 | not vllm_config.use_v2_model_runner |
| 181 | and single_group_spec is not None |
| 182 | and isinstance(single_group_spec, FullAttentionSpec) |
| 183 | and not isinstance(single_group_spec, MLAAttentionSpec) |
| 184 | and single_group_spec.num_kv_heads * parallel_config.tensor_parallel_size |
| 185 | == vllm_config.model_config.get_total_num_kv_heads() |
| 186 | and not single_group_spec.kv_quant_mode.is_per_token_head |
| 187 | and parallel_config.decode_context_parallel_size == 1 |
| 188 | and parallel_config.prefill_context_parallel_size == 1 |
| 189 | ) |
| 190 | # Canonical pages are topology-free, so the gate widens to every config |
| 191 | # whose mappings derive portable, group by group; certification happens |
| 192 | # per layer at registration and create_worker fails closed on this flag. |
| 193 | if canonical_layout and not is_parallelism_agnostic: |
| 194 | tp_size = parallel_config.tensor_parallel_size |
| 195 | total_kv_heads = vllm_config.model_config.get_total_num_kv_heads() |
| 196 | |
| 197 | def spec_certifiable(spec: KVCacheSpec) -> bool: |
| 198 | """Conservative static mirror of _layer_mapping's per-layer checks.""" |
| 199 | if not isinstance(spec, AttentionSpec): |
| 200 | return False |
| 201 | if spec.kv_quant_mode.is_per_token_head: |
| 202 | return False |
| 203 | if type(spec) is MLAAttentionSpec: |
| 204 | return ( |
| 205 | spec.tokens_per_state == 1 |
| 206 | and spec.real_page_size_bytes % spec.block_size == 0 |
| 207 | ) |
| 208 | if isinstance(spec, (SlidingWindowMLASpec, MLAAttentionSpec)): |
| 209 | return False |
| 210 | if not isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): |
| 211 | return False |
| 212 | return ( |
| 213 | total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0 |
| 214 | ) and spec.num_kv_heads == max(1, total_kv_heads // tp_size) |
| 215 | |
| 216 | # UniformTypeKVCacheSpecs groups (e.g. MLA plus its DSA indexer) hold |
| 217 | # one spec per layer; certify per layer, as the mapping derivation does. |
| 218 | layer_specs = [ |
| 219 | spec |
| 220 | for group in kv_cache_config.kv_cache_groups |
| 221 | for spec in iter_layer_specs(group.kv_cache_spec) |
| 222 | ] |
| 223 | is_parallelism_agnostic = ( |
| 224 | len(layer_specs) > 0 |
| 225 | and all(spec_certifiable(spec) for spec in layer_specs) |
| 226 | and parallel_config.decode_context_parallel_size == 1 |
| 227 | and parallel_config.prefill_context_parallel_size == 1 |
| 228 | and parallel_config.world_size == tp_size |
| 229 | ) |
| 230 | |
| 231 | if canonical_layout: |
| 232 | replicated_layout = ( |
| 233 | is_parallelism_agnostic |
| 234 | and all( |
| 235 | type(spec) is MLAAttentionSpec |
| 236 | for _, group in selected_groups |
| 237 | for spec in iter_layer_specs(group.kv_cache_spec) |
| 238 | ) |
| 239 | and parallel_config.nnodes_within_dp == 1 |
| 240 | and ( |
| 241 | parallel_config.world_size == 1 |
| 242 | or parallel_config.distributed_executor_backend == "mp" |
| 243 | ) |
| 244 | ) |
| 245 | |
| 246 | kv_events_config = vllm_config.kv_events_config |
| 247 | cache_dtype = ( |
| 248 | vllm_config.model_config.dtype |
| 249 | if vllm_config.cache_config.cache_dtype == "auto" |
| 250 | else vllm_config.cache_config.cache_dtype |
| 251 | ) |
| 252 | |
| 253 | return OffloadingConfig( |
| 254 | groups=groups, |
| 255 | worker_kv_bytes_per_block=worker_kv_bytes_per_block, |
| 256 | enable_kv_cache_events=( |
| 257 | kv_events_config is not None and kv_events_config.enable_kv_cache_events |
| 258 | ), |
| 259 | extra_config=extra_config, |
| 260 | engine_id=engine_id, |
| 261 | model=OffloadingModelConfig( |
| 262 | name=vllm_config.model_config.model, |
| 263 | dtype=str(cache_dtype).removeprefix("torch."), |
| 264 | ), |
| 265 | cache=OffloadingCacheConfig( |
| 266 | tokens_per_hash=tokens_per_hash, |
| 267 | blocks_per_chunk=blocks_per_chunk, |
| 268 | ), |
| 269 | parallel=OffloadingParallelConfig( |
| 270 | rank=parallel_config.rank, |
| 271 | world_size=parallel_config.world_size, |
| 272 | tp_size=parallel_config.tensor_parallel_size, |
| 273 | pp_size=parallel_config.pipeline_parallel_size, |
| 274 | pcp_size=parallel_config.prefill_context_parallel_size, |
| 275 | dcp_size=parallel_config.decode_context_parallel_size, |
| 276 | data_parallel_index=parallel_config.data_parallel_index, |
| 277 | data_parallel_size=parallel_config.data_parallel_size, |
| 278 | data_parallel_rank_local=parallel_config.data_parallel_rank_local, |
| 279 | is_parallelism_agnostic=is_parallelism_agnostic, |
| 280 | ), |
| 281 | replicated_layout=replicated_layout, |
| 282 | canonical_layout=canonical_layout, |
| 283 | kv_cache_layout=vllm_config.cache_config.kv_cache_layout, |
| 284 | ) |