[Bugfix][HiSparse] Size the KV cache from the groups HiSparse allocates (#59450)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Matthew Bonanni committed
Oct 2, 2026 at 00:11 UTC
f7999d2e4489126f1c21868699f3890021f01b73
8 files changed
+187
-71
tests/v1/core/test_kv_cache_utils.py
+115
-23
@@ -59,7 +59,7 @@ from vllm.v1.core.kv_cache_utils import (
59
)
60
from vllm.v1.hisparse.layout import (
61
create_hisparse_layout,
62
- get_hisparse_gpu_memory_usage,
62
+ get_hisparse_kv_cache_groups,
63
)
64
from vllm.v1.kv_cache_interface import (
65
ChunkedLocalAttentionSpec,
@@ -118,9 +118,6 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(
118
cache_role=SparseCacheRole.INDEXER,
119
),
120
}
121
- group_spec = UniformTypeKVCacheSpecs.from_specs(specs)
122
- assert group_spec is not None
123
- group = KVCacheGroupSpec(list(specs), group_spec)
121
config = SimpleNamespace(
122
attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
123
model_config=SimpleNamespace(
@@ -142,14 +139,9 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(
139
get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC,
140
),
141
)
145
- indexer_spec = specs["model.layers.0.self_attn.indexer"]
146
- assert get_hisparse_gpu_memory_usage(config, [group]) == (
147
- indexer_spec.max_memory_usage_bytes(config)
148
- )
149
-
142
monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
143
cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
152
- config, [group], available_memory=2**30
144
+ config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30
145
)
146
assert cache_config.num_blocks == 7
147
assert cache_config.hisparse_host_num_blocks is not None
@@ -228,8 +220,6 @@ def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok
220
cache_role=SparseCacheRole.INDEXER,
221
),
222
}
231
- group_spec = UniformTypeKVCacheSpecs.from_specs(specs)
232
- assert group_spec is not None
223
config = SimpleNamespace(
224
attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
225
model_config=SimpleNamespace(
@@ -240,14 +230,123 @@ def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok
230
)
231
host_page = specs["model.layers.0.self_attn"].page_size_bytes
232
host_budget = (4 + extra_blocks) * host_page
243
- group = KVCacheGroupSpec(list(specs), group_spec)
233
+ groups = get_hisparse_kv_cache_groups(config, specs)
234
235
if ok:
246
- layout = create_hisparse_layout(config, [group], host_budget=host_budget)
236
+ layout = create_hisparse_layout(config, groups, host_budget=host_budget)
237
assert layout.host_num_blocks == 4 + extra_blocks
238
else:
239
with pytest.raises(ValueError, match="increase host_pool_gib"):
250
- create_hisparse_layout(config, [group], host_budget=host_budget)
240
+ create_hisparse_layout(config, groups, host_budget=host_budget)
241
+
242
+
243
+@pytest.mark.parametrize(
244
+ "max_model_len,num_gpu_blocks,ok",
245
+ [
246
+ (8192, 160, False),
247
+ (8192, 1000, True),
248
+ (32768, 656, False),
249
+ (32768, 657, True),
250
+ (-1, 657, True),
251
+ ],
252
+)
253
+def test_hisparse_pool_must_fit_max_model_len(
254
+ monkeypatch, max_model_len, num_gpu_blocks, ok
255
+):
256
+ """Each resident group takes its own blocks from HiSparse's shared GPU pool,
257
+ so a pool that fits only the indexer pages must be rejected at startup.
258
+ Resident pages are capped at the in-flight window: at 32768 tokens with 2048
259
+ in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks,
260
+ plus the null block. Auto-fit (-1) must pick a length that passes the same
261
+ check."""
262
+ monkeypatch.setattr(
263
+ hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
264
+ )
265
+ monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
266
+ specs: dict[str, KVCacheSpec] = {}
267
+ for i in range(4):
268
+ specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec(
269
+ block_size=64,
270
+ num_kv_heads=1,
271
+ head_size=576,
272
+ dtype=torch.bfloat16,
273
+ is_index_group_leader=True,
274
+ )
275
+ specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec(
276
+ block_size=64,
277
+ num_kv_heads=1,
278
+ head_size=128,
279
+ dtype=torch.bfloat16,
280
+ cache_role=SparseCacheRole.INDEXER,
281
+ )
282
+ config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
283
+ config.scheduler_config.max_num_batched_tokens = 2048
284
+ # Async scheduling doubles the in-flight window, and with it the bound.
285
+ config.scheduler_config.async_scheduling = False
286
+ config.attention_config.hisparse_config = HiSparseConfig()
287
+ config.model_config.hf_config.index_topk = 128
288
+ config.cache_config.num_gpu_blocks_override = num_gpu_blocks
289
+ config.cache_config.kv_cache_layout = "BLHNC"
290
+ # Workers report the resident/hot caches HiSparse derives.
291
+ specs |= {
292
+ name: group.kv_cache_spec
293
+ for group in get_hisparse_kv_cache_groups(config, specs) or []
294
+ if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
295
+ for name in group.layer_names
296
+ }
297
+
298
+ if ok:
299
+ kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34])
300
+ else:
301
+ with pytest.raises(ValueError, match="max seq len"):
302
+ kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34])
303
+
304
+
305
+def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch):
306
+ """Workers report HiSparse's resident/hot caches, which lay out their own raw
307
+ backing; KV cache layout resolution must skip them."""
308
+ from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
309
+
310
+ monkeypatch.setattr(
311
+ hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
312
+ )
313
+ monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
314
+ specs: dict[str, KVCacheSpec] = {
315
+ "model.layers.0.self_attn": MLAAttentionSpec(
316
+ block_size=64,
317
+ num_kv_heads=1,
318
+ head_size=576,
319
+ dtype=torch.bfloat16,
320
+ is_index_group_leader=True,
321
+ ),
322
+ "model.layers.0.self_attn.indexer": MLAAttentionSpec(
323
+ block_size=64,
324
+ num_kv_heads=1,
325
+ head_size=128,
326
+ dtype=torch.bfloat16,
327
+ cache_role=SparseCacheRole.INDEXER,
328
+ ),
329
+ }
330
+ config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
331
+ config.attention_config.hisparse_config = HiSparseConfig()
332
+ config.model_config.hf_config.index_topk = 128
333
+ reported = specs | {
334
+ name: group.kv_cache_spec
335
+ for group in get_hisparse_kv_cache_groups(config, specs) or []
336
+ if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
337
+ for name in group.layer_names
338
+ }
339
+ assert len(reported) > len(specs)
340
+
341
+ layouts = []
342
+ for worker_specs in (specs, reported):
343
+ config.cache_config.kv_cache_layout = None
344
+ layouts.append(
345
+ resolve_kv_cache_layout(
346
+ config, [["BLHNC", "LBNHC"]], list(worker_specs.values())
347
+ )
348
+ )
349
+ assert layouts[0] == layouts[1]
350
351
352
def test_hisparse_rejects_deepseek_v4():
@@ -260,9 +359,6 @@ def test_hisparse_rejects_deepseek_v4():
359
model_version="deepseek_v4",
360
)
361
}
263
- full_uniform = UniformTypeKVCacheSpecs.from_specs(full_specs)
264
- assert full_uniform is not None
265
- group = KVCacheGroupSpec(list(full_specs), full_uniform)
362
config = SimpleNamespace(
363
attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
364
model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)),
@@ -270,11 +366,7 @@ def test_hisparse_rejects_deepseek_v4():
366
)
367
368
with pytest.raises(ValueError, match="does not support DeepSeek V4"):
273
- create_hisparse_layout(
274
- config,
275
- [group],
276
- host_budget=2**30,
277
- )
369
+ get_hisparse_kv_cache_groups(config, full_specs)
370
371
372
@pytest.fixture(autouse=True)
tests/v1/worker/test_attn_utils.py
+3
@@ -86,6 +86,9 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size(
86
get_attn_backend=lambda backend=backend: backend,
87
)
88
monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers)
89
+ monkeypatch.setattr(
90
+ attn_utils_module, "get_hisparse_kv_cache_groups", lambda *_: []
91
+ )
92
config = SimpleNamespace(
93
attention_config=SimpleNamespace(hisparse_config=object() if enabled else None)
94
)
vllm/v1/attention/backends/utils.py
+4
-1
@@ -291,7 +291,10 @@ def resolve_kv_cache_layout(
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
- kv_cache_specs = tuple(kv_cache_specs or ())
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
vllm/v1/core/kv_cache_utils.py
+8
-10
@@ -21,7 +21,6 @@ from vllm.utils.math_utils import cdiv, round_up
21
from vllm.utils.mem_utils import format_gib
22
from vllm.utils.torch_utils import get_dtype_size
23
from vllm.v1.hisparse.layout import (
24
- get_hisparse_gpu_memory_usage,
24
get_hisparse_host_pool_bytes,
25
get_hisparse_kv_cache_config,
26
get_hisparse_kv_cache_groups,
@@ -1155,7 +1154,9 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
1154
`available_memory` into `num_blocks`. Used to compute the effective KV cache
1155
capacity once `num_gpu_blocks_override` is applied.
1156
"""
1158
- return _get_kv_cache_bytes_per_block(kv_cache_groups)
1157
+ return _get_kv_cache_bytes_per_block(
1158
+ [group for group in kv_cache_groups if not group.host_resident]
1159
+ )
1160
1161
1162
def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
@@ -2500,12 +2501,10 @@ def _max_memory_usage_bytes_from_groups(
2501
Each group independently claims blocks from the shared pool, so a request consumes
2502
the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``.
2503
"""
2504
+ kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident]
2505
if not kv_cache_groups:
2506
return 0
2507
2506
- if vllm_config.attention_config.hisparse_config is not None:
2507
- return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups)
2508
-
2508
if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
2509
(
2510
attn_group,
@@ -2563,12 +2562,11 @@ def _estimate_max_model_len_from_groups(
2562
vllm_config.model_config.max_model_len = model_len
2563
if hisparse_enabled:
2564
try:
2566
- config = get_kv_cache_config_from_groups(
2565
+ get_kv_cache_config_from_groups(
2566
vllm_config, kv_cache_groups, available_memory
2567
)
2568
except ValueError:
2569
return False
2571
- return get_max_concurrency_for_kv_cache_config(vllm_config, config) >= 1
2570
return (
2571
_max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
2572
<= available_memory
@@ -2678,6 +2676,9 @@ def _project_kv_cache_groups_to_worker(
2676
worker_layer_names = [
2677
layer_name for layer_name in group.layer_names if layer_name in worker_spec
2678
]
2679
+ if len(worker_layer_names) == len(group.layer_names):
2680
+ projected_groups.append(group)
2681
+ continue
2682
group_spec = group.kv_cache_spec
2683
if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs):
2684
group_spec = UniformTypeKVCacheSpecs(
@@ -2793,9 +2794,6 @@ def get_kv_cache_configs(
2794
adjusted_memory.append(override * bytes_per_block)
2795
available_memory = adjusted_memory
2796
2796
- if vllm_config.attention_config.hisparse_config is not None:
2797
- available_memory = [min(available_memory)] * len(available_memory)
2798
-
2797
# Reserve the null block BlockPool permanently holds back, so auto-fit and
2798
# the capacity check both plan against usable blocks. Allocation below
2799
# still uses the full memory.
vllm/v1/hisparse/binding.py
+17
-9
@@ -12,6 +12,7 @@ from vllm.config import VllmConfig
12
from vllm.v1.hisparse.layout import (
13
HISPARSE_HOT_SUFFIX,
14
HISPARSE_RESIDENT_SUFFIX,
15
+ get_hisparse_kv_cache_groups,
16
)
17
from vllm.v1.hisparse.runtime import (
18
HiSparseCacheHandle,
@@ -36,21 +37,20 @@ if TYPE_CHECKING:
37
from vllm.v1.worker.gpu.block_table import BlockTables
38
39
39
-def resolve_hisparse_block_size(
40
+def resolve_hisparse_specs(
41
vllm_config: VllmConfig,
42
kv_cache_spec: dict[str, KVCacheSpec],
43
attn_layers: Mapping[str, "AttentionLayerBase"],
43
-) -> None:
44
- """Resolve a common kernel block size in-place for HiSparse MLA specs."""
45
- if vllm_config.attention_config.hisparse_config is None:
46
- return
44
+) -> dict[str, KVCacheSpec]:
45
+ """Resolve a common kernel block size for HiSparse MLA specs, and add the
46
+ resident/hot specs HiSparse derives from them."""
47
mla_specs = {
48
name: spec
49
for name, spec in kv_cache_spec.items()
50
if isinstance(spec, MLAAttentionSpec)
51
}
52
if not mla_specs:
53
- return
53
+ return kv_cache_spec
54
block_sizes = {spec.block_size for spec in mla_specs.values()}
55
if len(block_sizes) != 1:
56
raise ValueError("HiSparse requires one scheduler block size.")
@@ -62,10 +62,18 @@ def resolve_hisparse_block_size(
62
"HiSparse requires a GPU block size supported by every sparse "
63
f"attention and indexer backend: {error}"
64
) from error
65
- kv_cache_spec.update(
66
- (name, spec.copy_with_new_block_size(block_size))
65
+ kv_cache_spec = kv_cache_spec | {
66
+ name: spec.copy_with_new_block_size(block_size)
67
for name, spec in mla_specs.items()
68
- )
68
+ }
69
+ groups = get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec)
70
+ assert groups is not None
71
+ return kv_cache_spec | {
72
+ layer_name: group.kv_cache_spec
73
+ for group in groups
74
+ if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
75
+ for layer_name in group.layer_names
76
+ }
77
78
79
def allocate_hisparse_kv_caches(
vllm/v1/hisparse/layout.py
+33
-24
@@ -47,6 +47,10 @@ class HiSparseLayout:
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
@@ -59,9 +63,11 @@ def get_hisparse_kv_cache_groups(
63
other_specs = {
64
name: spec
65
for name, spec in kv_cache_spec.items()
62
- if not isinstance(spec, MLAAttentionSpec)
66
+ if not isinstance(
67
+ spec, (MLAAttentionSpec, HiSparseResidentSpec, HiSparseHotSpec)
68
+ )
69
}
64
- if not mla_specs or not other_specs:
70
+ if not mla_specs:
71
return None
72
73
from vllm.v1.core.kv_cache_utils import get_kv_cache_groups
@@ -69,7 +75,10 @@ def get_hisparse_kv_cache_groups(
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)
72
- return [mla_group, *get_kv_cache_groups(vllm_config, other_specs)]
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:
@@ -112,24 +121,9 @@ def _partition_hisparse_specs(
121
return source_specs, indexer_specs
122
123
115
-def get_hisparse_gpu_memory_usage(
116
- vllm_config: VllmConfig,
117
- kv_cache_groups: list[KVCacheGroupSpec],
118
-) -> int:
119
- _, indexer_specs = _partition_hisparse_specs(kv_cache_groups)
120
- return sum(
121
- spec.max_memory_usage_bytes(vllm_config) for spec in indexer_specs.values()
122
- ) + sum(
123
- group.kv_cache_spec.max_memory_usage_bytes(vllm_config)
124
- for group in kv_cache_groups[1:]
125
- )
126
-
127
-
128
-def create_hisparse_layout(
129
- vllm_config: VllmConfig,
130
- groups: list[KVCacheGroupSpec],
131
- host_budget: int,
132
-) -> HiSparseLayout:
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())
@@ -225,11 +219,25 @@ def create_hisparse_layout(
219
enable_kv_transfer=True,
220
)
221
regular_groups = groups[1:]
228
- gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups]
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(
232
- sum(spec.page_size_bytes for spec in source_specs.values()),
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
@@ -237,6 +245,7 @@ def create_hisparse_layout(
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(
@@ -246,7 +255,7 @@ def create_hisparse_layout(
255
256
return HiSparseLayout(
257
source_group=source_group,
249
- device_groups=gpu_groups,
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,
vllm/v1/kv_cache_interface.py
+4
-2
@@ -472,9 +472,11 @@ class HiSparseResidentSpec(KVCacheSpec):
472
return cdiv(num_tokens, self.block_size)
473
474
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
475
- return cdiv(vllm_config.model_config.max_model_len, self.block_size) * (
476
- self.page_size
475
+ max_blocks = self.max_admission_blocks_per_request(
476
+ max_in_flight_tokens=vllm_config.max_in_flight_tokens,
477
+ max_model_len=vllm_config.model_config.max_model_len,
478
)
479
+ return max_blocks * self.page_size
480
481
@property
482
def has_layer_views(self) -> bool:
vllm/v1/worker/gpu/attn_utils.py
+3
-2
@@ -23,7 +23,7 @@ from vllm.v1.attention.backend import (
23
from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend
24
from vllm.v1.hisparse.binding import (
25
init_hisparse_kv_cache,
26
- resolve_hisparse_block_size,
26
+ resolve_hisparse_specs,
27
)
28
from vllm.v1.kv_cache_interface import (
29
AttentionSpec,
@@ -142,7 +142,8 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
142
if isinstance(spec, AttentionSpec):
143
spec = attn_module.get_attn_backend().customize_spec(spec)
144
kv_cache_spec[layer_name] = spec
145
- resolve_hisparse_block_size(vllm_config, kv_cache_spec, attn_layers)
145
+ if vllm_config.attention_config.hisparse_config is not None:
146
+ kv_cache_spec = resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers)
147
return kv_cache_spec
148
149