| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections import defaultdict |
| 5 | from dataclasses import dataclass |
| 6 | from typing import cast |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.config.kv_transfer import hisparse_host_pool_gib |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.utils.math_utils import cdiv |
| 12 | from vllm.v1.hisparse.runtime import ( |
| 13 | ResolvedHiSparseConfig, |
| 14 | get_hisparse_host_block_stride, |
| 15 | use_shared_hisparse_host_pool, |
| 16 | ) |
| 17 | from vllm.v1.kv_cache_interface import ( |
| 18 | HiSparseHotSpec, |
| 19 | HiSparseResidentSpec, |
| 20 | KVCacheConfig, |
| 21 | KVCacheGroupRole, |
| 22 | KVCacheGroupSpec, |
| 23 | KVCacheSpec, |
| 24 | KVCacheTensor, |
| 25 | MLAAttentionSpec, |
| 26 | SparseCacheRole, |
| 27 | UniformTypeKVCacheSpecs, |
| 28 | compute_layout_strides, |
| 29 | ) |
| 30 | from vllm.v1.kv_cache_layout import KVCacheLayout |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | HISPARSE_HOT_SUFFIX = ".hisparse_hot" |
| 35 | HISPARSE_RESIDENT_SUFFIX = ".hisparse_resident" |
| 36 | |
| 37 | |
| 38 | @dataclass(frozen=True) |
| 39 | class HiSparseLayout: |
| 40 | source_group: KVCacheGroupSpec |
| 41 | device_groups: list[KVCacheGroupSpec] |
| 42 | host_num_blocks: int |
| 43 | host_block_stride: int |
| 44 | shared_host_pool: bool |
| 45 | |
| 46 | |
| 47 | def get_hisparse_kv_cache_groups( |
| 48 | vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec] |
| 49 | ) -> list[KVCacheGroupSpec] | None: |
| 50 | """The groups HiSparse allocates: the host source group, then the GPU |
| 51 | groups sharing one device pool. Derived resident/hot specs in |
| 52 | `kv_cache_spec` (see `resolve_hisparse_specs`) are laid out again |
| 53 | from the attention specs they belong to.""" |
| 54 | attention_config = getattr(vllm_config, "attention_config", None) |
| 55 | if attention_config is None or attention_config.hisparse_config is None: |
| 56 | return None |
| 57 | |
| 58 | mla_specs: dict[str, KVCacheSpec] = { |
| 59 | name: spec |
| 60 | for name, spec in kv_cache_spec.items() |
| 61 | if isinstance(spec, MLAAttentionSpec) |
| 62 | } |
| 63 | other_specs = { |
| 64 | name: spec |
| 65 | for name, spec in kv_cache_spec.items() |
| 66 | if not isinstance( |
| 67 | spec, (MLAAttentionSpec, HiSparseResidentSpec, HiSparseHotSpec) |
| 68 | ) |
| 69 | } |
| 70 | if not mla_specs: |
| 71 | return None |
| 72 | |
| 73 | from vllm.v1.core.kv_cache_utils import get_kv_cache_groups |
| 74 | |
| 75 | mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) |
| 76 | assert mla_group_spec is not None |
| 77 | mla_group = KVCacheGroupSpec(list(mla_specs), mla_group_spec) |
| 78 | regular_groups = ( |
| 79 | get_kv_cache_groups(vllm_config, other_specs) if other_specs else [] |
| 80 | ) |
| 81 | return _lay_out_hisparse_groups(vllm_config, [mla_group, *regular_groups]) |
| 82 | |
| 83 | |
| 84 | def get_hisparse_host_pool_bytes(vllm_config: VllmConfig) -> int: |
| 85 | host_pool_gib = hisparse_host_pool_gib(vllm_config.kv_transfer_config) |
| 86 | if host_pool_gib is None: |
| 87 | raise ValueError("HiSparse requires HiSparseConnector with host_pool_gib") |
| 88 | return int(host_pool_gib * 2**30) |
| 89 | |
| 90 | |
| 91 | def _partition_hisparse_specs( |
| 92 | groups: list[KVCacheGroupSpec], |
| 93 | ) -> tuple[dict[str, MLAAttentionSpec], dict[str, MLAAttentionSpec]]: |
| 94 | group_spec = groups[0].kv_cache_spec |
| 95 | if not isinstance(group_spec, UniformTypeKVCacheSpecs): |
| 96 | raise ValueError("HiSparse requires uniform sparse-MLA cache specs.") |
| 97 | |
| 98 | specs = group_spec.kv_cache_specs |
| 99 | if any( |
| 100 | isinstance(spec, MLAAttentionSpec) and spec.model_version == "deepseek_v4" |
| 101 | for spec in specs.values() |
| 102 | ): |
| 103 | raise ValueError("HiSparse does not support DeepSeek V4.") |
| 104 | if not all(isinstance(spec, MLAAttentionSpec) for spec in specs.values()): |
| 105 | raise ValueError("HiSparse requires its first cache group to contain MLA only.") |
| 106 | |
| 107 | source_specs = { |
| 108 | name: spec |
| 109 | for name, spec in specs.items() |
| 110 | if isinstance(spec, MLAAttentionSpec) |
| 111 | and spec.cache_role is SparseCacheRole.SPARSE |
| 112 | } |
| 113 | indexer_specs = { |
| 114 | name: spec |
| 115 | for name, spec in specs.items() |
| 116 | if isinstance(spec, MLAAttentionSpec) |
| 117 | and spec.cache_role is SparseCacheRole.INDEXER |
| 118 | } |
| 119 | if not source_specs or not indexer_specs: |
| 120 | raise ValueError("HiSparse requires sparse-MLA and indexer cache specs.") |
| 121 | return source_specs, indexer_specs |
| 122 | |
| 123 | |
| 124 | def _lay_out_hisparse_groups( |
| 125 | vllm_config: VllmConfig, groups: list[KVCacheGroupSpec] |
| 126 | ) -> list[KVCacheGroupSpec]: |
| 127 | source_specs, indexer_specs = _partition_hisparse_specs(groups) |
| 128 | block_sizes = { |
| 129 | spec.block_size for spec in (*source_specs.values(), *indexer_specs.values()) |
| 130 | } |
| 131 | if len(block_sizes) != 1: |
| 132 | raise ValueError("HiSparse requires one resolved GPU block size.") |
| 133 | gpu_block_size = block_sizes.pop() |
| 134 | |
| 135 | config = ResolvedHiSparseConfig.from_vllm_config( |
| 136 | vllm_config, |
| 137 | vllm_config.model_config.hf_config.index_topk, |
| 138 | gpu_block_size, |
| 139 | ) |
| 140 | assert config is not None |
| 141 | indexer_group_spec = UniformTypeKVCacheSpecs.from_specs( |
| 142 | cast(dict[str, KVCacheSpec], indexer_specs) |
| 143 | ) |
| 144 | assert indexer_group_spec is not None |
| 145 | indexer_group = KVCacheGroupSpec( |
| 146 | list(indexer_specs), |
| 147 | indexer_group_spec, |
| 148 | enable_kv_transfer=True, |
| 149 | role=KVCacheGroupRole.HISPARSE_INDEXER, |
| 150 | ) |
| 151 | |
| 152 | indexer_page = sum(spec.page_size_bytes for spec in indexer_specs.values()) |
| 153 | hot_blocks_per_request = cdiv(config.device_buffer_size, gpu_block_size) |
| 154 | hot_units: list[list[tuple[str, MLAAttentionSpec]]] = [] |
| 155 | for layer_name, layer_spec in source_specs.items(): |
| 156 | if layer_spec.is_index_group_leader or not hot_units: |
| 157 | hot_units.append([]) |
| 158 | hot_units[-1].append((f"{layer_name}{HISPARSE_HOT_SUFFIX}", layer_spec)) |
| 159 | |
| 160 | resident_groups: list[KVCacheGroupSpec] = [] |
| 161 | hot_groups: list[KVCacheGroupSpec] = [] |
| 162 | |
| 163 | def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None: |
| 164 | page_sizes = {spec.page_size_bytes for _, spec in layers} |
| 165 | if len(page_sizes) != 1: |
| 166 | raise ValueError( |
| 167 | "HiSparse hot-cache groups require one page size, got " |
| 168 | f"{sorted(page_sizes)}." |
| 169 | ) |
| 170 | page_size = page_sizes.pop() |
| 171 | names = [name for name, _ in layers] |
| 172 | resident_groups.append( |
| 173 | KVCacheGroupSpec( |
| 174 | [ |
| 175 | name[: -len(HISPARSE_HOT_SUFFIX)] + HISPARSE_RESIDENT_SUFFIX |
| 176 | for name in names |
| 177 | ], |
| 178 | HiSparseResidentSpec( |
| 179 | block_size=gpu_block_size, |
| 180 | page_size=page_size, |
| 181 | ), |
| 182 | enable_kv_transfer=False, |
| 183 | ) |
| 184 | ) |
| 185 | hot_groups.append( |
| 186 | KVCacheGroupSpec( |
| 187 | names, |
| 188 | HiSparseHotSpec( |
| 189 | block_size=gpu_block_size, |
| 190 | page_size=page_size, |
| 191 | blocks_per_request=hot_blocks_per_request, |
| 192 | ), |
| 193 | enable_kv_transfer=False, |
| 194 | ) |
| 195 | ) |
| 196 | |
| 197 | current: list[tuple[str, MLAAttentionSpec]] = [] |
| 198 | current_page = 0 |
| 199 | for unit in hot_units: |
| 200 | unit_page = sum(spec.page_size_bytes for _, spec in unit) |
| 201 | if current and current_page + unit_page > indexer_page: |
| 202 | append_hot_group(current) |
| 203 | current = [] |
| 204 | current_page = 0 |
| 205 | current.extend(unit) |
| 206 | current_page += unit_page |
| 207 | if current: |
| 208 | append_hot_group(current) |
| 209 | |
| 210 | source_group_spec = UniformTypeKVCacheSpecs.from_specs( |
| 211 | cast(dict[str, KVCacheSpec], source_specs) |
| 212 | ) |
| 213 | assert source_group_spec is not None |
| 214 | source_group = KVCacheGroupSpec( |
| 215 | list(source_specs), |
| 216 | source_group_spec, |
| 217 | role=KVCacheGroupRole.HISPARSE_SOURCE, |
| 218 | host_resident=True, |
| 219 | enable_kv_transfer=True, |
| 220 | ) |
| 221 | regular_groups = groups[1:] |
| 222 | return [ |
| 223 | source_group, |
| 224 | indexer_group, |
| 225 | *resident_groups, |
| 226 | *hot_groups, |
| 227 | *regular_groups, |
| 228 | ] |
| 229 | |
| 230 | |
| 231 | def create_hisparse_layout( |
| 232 | vllm_config: VllmConfig, |
| 233 | groups: list[KVCacheGroupSpec], |
| 234 | host_budget: int, |
| 235 | ) -> HiSparseLayout: |
| 236 | """Size the host pool for groups laid out by `get_hisparse_kv_cache_groups`.""" |
| 237 | (source_group,) = [group for group in groups if group.host_resident] |
| 238 | shared_host_pool = use_shared_hisparse_host_pool(vllm_config) |
| 239 | host_block_stride = get_hisparse_host_block_stride( |
| 240 | source_group.kv_cache_spec.page_size_bytes, |
| 241 | use_shared_host_pool=shared_host_pool, |
| 242 | ) |
| 243 | host_num_blocks = host_budget // host_block_stride |
| 244 | if host_num_blocks <= 0: |
| 245 | raise ValueError("HiSparse has no allocatable host blocks.") |
| 246 | # Every computed page needs a host block, so one request at max_model_len |
| 247 | # must fit alongside the pool's null block and a copy-on-write tail. |
| 248 | gpu_block_size = source_group.kv_cache_spec.block_size |
| 249 | min_host_blocks = cdiv(vllm_config.model_config.max_model_len, gpu_block_size) + 2 |
| 250 | if host_num_blocks < min_host_blocks: |
| 251 | raise ValueError( |
| 252 | f"HiSparse host pool has {host_num_blocks} blocks but max_model_len " |
| 253 | f"needs {min_host_blocks}; increase host_pool_gib." |
| 254 | ) |
| 255 | |
| 256 | return HiSparseLayout( |
| 257 | source_group=source_group, |
| 258 | device_groups=[group for group in groups if not group.host_resident], |
| 259 | host_num_blocks=host_num_blocks, |
| 260 | host_block_stride=host_block_stride, |
| 261 | shared_host_pool=shared_host_pool, |
| 262 | ) |
| 263 | |
| 264 | |
| 265 | def _build_hisparse_kv_cache_tensors( |
| 266 | kv_cache_groups: list[KVCacheGroupSpec], |
| 267 | num_blocks: int, |
| 268 | size: int, |
| 269 | layout: KVCacheLayout, |
| 270 | bytes_per_block: int, |
| 271 | *, |
| 272 | host_resident: bool = False, |
| 273 | ) -> list[KVCacheTensor]: |
| 274 | interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None |
| 275 | tensors: list[KVCacheTensor] = [] |
| 276 | for group in kv_cache_groups: |
| 277 | group_spec = group.kv_cache_spec |
| 278 | layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) |
| 279 | if isinstance(group_spec, UniformTypeKVCacheSpecs): |
| 280 | for layer_name, spec in group_spec.kv_cache_specs.items(): |
| 281 | layers_by_spec[spec].append(layer_name) |
| 282 | elif group.layer_names: |
| 283 | layers_by_spec[group_spec].extend(group.layer_names) |
| 284 | |
| 285 | byte_offset = 0 |
| 286 | for spec, layer_names in layers_by_spec.items(): |
| 287 | if isinstance(spec, (HiSparseHotSpec, HiSparseResidentSpec)): |
| 288 | if not layout.is_block_outermost: |
| 289 | raise ValueError("HiSparse requires a block-outermost KV layout.") |
| 290 | layer_stride = spec.page_size_bytes |
| 291 | block_stride = bytes_per_block |
| 292 | else: |
| 293 | layer_stride, block_stride, _, _, _ = compute_layout_strides( |
| 294 | spec, |
| 295 | num_blocks, |
| 296 | len(layer_names), |
| 297 | layout, |
| 298 | fixed_strides=(None, interleaved_block_stride, None, None, None), |
| 299 | ) |
| 300 | offset = ( |
| 301 | byte_offset |
| 302 | * max(layer_stride, spec.page_size_bytes) |
| 303 | // spec.page_size_bytes |
| 304 | ) |
| 305 | tensors.append( |
| 306 | KVCacheTensor( |
| 307 | size=size, |
| 308 | layers=layer_names, |
| 309 | layer_stride=layer_stride, |
| 310 | block_stride=block_stride, |
| 311 | offset=offset, |
| 312 | host_resident=host_resident, |
| 313 | ) |
| 314 | ) |
| 315 | byte_offset += len(layer_names) * spec.page_size_bytes |
| 316 | return tensors |
| 317 | |
| 318 | |
| 319 | def get_hisparse_kv_cache_config( |
| 320 | vllm_config: VllmConfig, |
| 321 | kv_cache_groups: list[KVCacheGroupSpec], |
| 322 | available_memory: int, |
| 323 | host_budget: int, |
| 324 | ) -> KVCacheConfig: |
| 325 | from vllm.v1.core.kv_cache_utils import ( |
| 326 | _get_kv_cache_bytes_per_block, |
| 327 | may_override_num_blocks, |
| 328 | validate_kv_cache_layout, |
| 329 | ) |
| 330 | |
| 331 | hisparse_layout = create_hisparse_layout(vllm_config, kv_cache_groups, host_budget) |
| 332 | device_groups = hisparse_layout.device_groups |
| 333 | layout = vllm_config.cache_config.get_resolved_kv_cache_layout() |
| 334 | validate_kv_cache_layout(layout, device_groups) |
| 335 | bytes_per_block = _get_kv_cache_bytes_per_block(device_groups) |
| 336 | num_blocks = may_override_num_blocks( |
| 337 | vllm_config, available_memory // bytes_per_block |
| 338 | ) |
| 339 | size = bytes_per_block * num_blocks |
| 340 | kv_cache_tensors = _build_hisparse_kv_cache_tensors( |
| 341 | device_groups, num_blocks, size, layout, bytes_per_block |
| 342 | ) |
| 343 | |
| 344 | host_groups = [hisparse_layout.source_group] |
| 345 | host_layout = KVCacheLayout.LBNHC |
| 346 | validate_kv_cache_layout(host_layout, host_groups) |
| 347 | host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups) |
| 348 | host_size = host_bytes_per_block * hisparse_layout.host_num_blocks |
| 349 | kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors( |
| 350 | host_groups, |
| 351 | hisparse_layout.host_num_blocks, |
| 352 | host_size, |
| 353 | host_layout, |
| 354 | host_bytes_per_block, |
| 355 | host_resident=True, |
| 356 | ) |
| 357 | logger.info_once( |
| 358 | "HiSparse HMA: %.1f GiB host source (%d blocks), %.1f GiB shared " |
| 359 | "GPU indexer/resident/hot pool (%d blocks).", |
| 360 | host_size / 2**30, |
| 361 | hisparse_layout.host_num_blocks, |
| 362 | size / 2**30, |
| 363 | num_blocks, |
| 364 | ) |
| 365 | return KVCacheConfig( |
| 366 | num_blocks=num_blocks, |
| 367 | kv_cache_tensors=kv_cache_tensors, |
| 368 | kv_cache_groups=[*host_groups, *device_groups], |
| 369 | hisparse_host_num_blocks=hisparse_layout.host_num_blocks, |
| 370 | hisparse_host_block_stride=hisparse_layout.host_block_stride, |
| 371 | hisparse_shared_host_pool=hisparse_layout.shared_host_pool, |
| 372 | prefix_cache_retention_interval=( |
| 373 | vllm_config.cache_config.prefix_cache_retention_interval |
| 374 | ), |
| 375 | ) |