[Bugfix][HiSparse] Never allocate GPU pages without host backing (#59036)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Lucas Wilkinson committed
Sep 30, 2026 at 00:13 UTC
90e13fc757cffc11b55689d9cc33d844d9606516
8 files changed
+252
-85
tests/test_config.py
+13
@@ -420,6 +420,19 @@ def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch):
420
)
421
422
423
+def test_hisparse_rejects_disabled_full_isl_reservation(monkeypatch):
424
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
425
+ with pytest.raises(ValueError, match="requires --scheduler-reserve-full-isl"):
426
+ VllmConfig(
427
+ attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
428
+ scheduler_config=SchedulerConfig(
429
+ max_model_len=2048,
430
+ is_encoder_decoder=False,
431
+ scheduler_reserve_full_isl=False,
432
+ ),
433
+ )
434
+
435
+
436
def test_hisparse_rejects_non_cuda(monkeypatch):
437
monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
438
with pytest.raises(ValueError, match="requires NVIDIA CUDA"):
tests/v1/core/test_kv_cache_utils.py
+44
@@ -205,6 +205,50 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(
205
assert scheduler_block_size == hash_block_size == gpu_block_size
206
207
208
+@pytest.mark.parametrize("extra_blocks,ok", [(1, False), (2, True)])
209
+def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok):
210
+ """Every page needs a host block, so one max_model_len request must fit."""
211
+ monkeypatch.setattr(
212
+ hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
213
+ )
214
+ specs = {
215
+ "model.layers.0.self_attn": MLAAttentionSpec(
216
+ block_size=64,
217
+ num_kv_heads=1,
218
+ head_size=576,
219
+ dtype=torch.bfloat16,
220
+ is_index_group_leader=True,
221
+ ),
222
+ "model.layers.0.self_attn.indexer": MLAAttentionSpec(
223
+ block_size=64,
224
+ num_kv_heads=1,
225
+ head_size=128,
226
+ dtype=torch.bfloat16,
227
+ cache_role=SparseCacheRole.INDEXER,
228
+ ),
229
+ }
230
+ group_spec = UniformTypeKVCacheSpecs.from_specs(specs)
231
+ assert group_spec is not None
232
+ config = SimpleNamespace(
233
+ attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
234
+ model_config=SimpleNamespace(
235
+ hf_config=SimpleNamespace(index_topk=128), max_model_len=64 * 4
236
+ ),
237
+ parallel_config=SimpleNamespace(tensor_parallel_size=1, world_size=1),
238
+ cache_config=SimpleNamespace(num_gpu_blocks_override=7),
239
+ )
240
+ host_page = specs["model.layers.0.self_attn"].page_size_bytes
241
+ host_budget = (4 + extra_blocks) * host_page
242
+ group = KVCacheGroupSpec(list(specs), group_spec)
243
+
244
+ if ok:
245
+ layout = create_hisparse_layout(config, [group], host_budget=host_budget)
246
+ assert layout.host_num_blocks == 4 + extra_blocks
247
+ else:
248
+ with pytest.raises(ValueError, match="increase host_pool_gib"):
249
+ create_hisparse_layout(config, [group], host_budget=host_budget)
250
+
251
+
252
def test_hisparse_rejects_deepseek_v4():
253
full_specs = {
254
"model.layers.0.attn": MLAAttentionSpec(
tests/v1/core/test_prefix_caching.py
+80
-44
@@ -782,14 +782,15 @@ def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages():
782
assert required == 4
783
784
785
-@pytest.mark.parametrize(
786
- "full_sequence_must_fit,host_num_blocks,admitted",
787
- [(True, 2, False), (True, 3, True), (False, 2, False), (False, 3, True)],
788
-)
789
-def test_hisparse_async_admission_requires_only_import_destinations(
790
- full_sequence_must_fit, host_num_blocks, admitted, tmp_path
785
+@pytest.mark.parametrize("host_num_blocks,admitted", [(5, False), (6, True)])
786
+def test_hisparse_async_admission_requires_full_prompt_host_pages(
787
+ host_num_blocks, admitted, tmp_path
788
):
792
- """Imports need host destinations, but do not reserve future prefill pages."""
789
+ """An async load is admitted only if its whole prompt's host pages fit.
790
+
791
+ The in-flight prefill's host page is not free, so the waiting request's
792
+ four prompt pages need six host blocks including the null block.
793
+ """
794
from .utils import create_scheduler, mock_kv
795
796
(tmp_path / "config.json").write_text(
@@ -804,7 +805,7 @@ def test_hisparse_async_admission_requires_only_import_destinations(
805
manager = make_hisparse_kv_cache_manager(32, host_num_blocks)
806
scheduler.kv_cache_manager = manager
807
scheduler.kv_cache_config = manager.kv_cache_config
807
- scheduler.scheduler_reserve_full_isl = full_sequence_must_fit
808
+ assert scheduler.scheduler_reserve_full_isl
809
inflight = make_request(
810
"inflight", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
811
)
@@ -822,6 +823,52 @@ def test_hisparse_async_admission_requires_only_import_destinations(
823
assert bool(manager.get_blocks(request.request_id).blocks[0]) == admitted
824
825
826
+def test_hisparse_admitted_async_loads_can_finish_with_nothing_running(tmp_path):
827
+ """Admitted async loads must not strand each other on host capacity.
828
+
829
+ Waiting requests keep their blocks and are never preempted, so two imports
830
+ that each hold a host page while together needing more than the pool has
831
+ left would never be scheduled again. HiSparse relies on full-ISL admission:
832
+ a load is admitted only if its whole prompt's host pages fit, so the newest
833
+ admitted load can always finish and free host blocks for the others.
834
+ """
835
+ from .utils import create_scheduler, mock_kv
836
+
837
+ (tmp_path / "config.json").write_text(
838
+ '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
839
+ )
840
+ scheduler = create_scheduler(
841
+ model=str(tmp_path),
842
+ skip_tokenizer_init=True,
843
+ max_model_len=128,
844
+ use_kv_connector=mock_kv(matched_tokens=HISPARSE_BLOCK_SIZE, is_async=True),
845
+ )
846
+ # One null block plus three usable host blocks.
847
+ manager = make_hisparse_kv_cache_manager(32, 4)
848
+ scheduler.kv_cache_manager = manager
849
+ scheduler.kv_cache_config = manager.kv_cache_config
850
+ assert scheduler.scheduler_reserve_full_isl
851
+ first = make_request(
852
+ "first", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
853
+ )
854
+ second = make_request(
855
+ "second",
856
+ list(range(1000, 1000 + 3 * HISPARSE_BLOCK_SIZE)),
857
+ HISPARSE_BLOCK_SIZE,
858
+ sha256,
859
+ )
860
+ for request in (first, second):
861
+ scheduler.add_request(request)
862
+ scheduler.schedule()
863
+ admitted = {r.request_id for r in scheduler._inflight_prefills}
864
+
865
+ scheduler.finished_recving_kv_req_ids.update(admitted)
866
+ num_scheduled_tokens = 0
867
+ for _ in range(4):
868
+ num_scheduled_tokens += scheduler.schedule().total_num_scheduled_tokens
869
+ assert num_scheduled_tokens > 0
870
+
871
+
872
@pytest.mark.parametrize(
873
"evictable,local_tokens,free_blocks,admitted",
874
[
@@ -845,58 +892,47 @@ def test_hisparse_import_capacity_includes_hits_and_cow(
892
assert host_pool.get_num_free_blocks() == free_blocks
893
assert (
894
coordinator.host_manager.get_num_blocks_to_allocate(
848
- "waiting", 64, [hit], 32, local_tokens, 64
895
+ "waiting", 32, [hit], 32, local_tokens, 32
896
)
897
<= host_pool.get_num_free_blocks()
898
) == admitted
899
900
901
@pytest.mark.parametrize("enable_caching", [False, True])
855
-def test_hisparse_host_exhaustion_keeps_gpu_pages_readable(enable_caching):
856
- """Pages without a host destination must survive GPU cache reclamation."""
857
- manager = make_hisparse_kv_cache_manager(32, 2, enable_caching=enable_caching)
902
+def test_hisparse_host_exhaustion_defers_allocation(enable_caching):
903
+ """Pages are only allocated with host backing, so every one can be reclaimed."""
904
+ manager = make_hisparse_kv_cache_manager(32, 5, enable_caching=enable_caching)
905
+ donor = make_request("donor", [99] * 16, HISPARSE_BLOCK_SIZE, sha256)
906
+ assert manager.allocate_slots(donor, 16) is not None
907
request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
859
- assert manager.allocate_slots(request, 64) is not None
908
+ free_gpu = manager.block_pool.get_num_free_blocks()
909
+ assert manager.allocate_slots(request, 64) is None
910
+ assert manager.block_pool.get_num_free_blocks() == free_gpu
911
+
912
+ # The donor's in-flight write-back holds its host block until acked.
913
+ manager.free(donor)
914
coordinator = get_hisparse_coordinator(manager)
915
+ assert coordinator.has_pending_reclamation()
916
+ assert manager.allocate_slots(request, 64) is None
917
+ _publish_hisparse_pages(manager)
918
+ assert not coordinator.has_pending_reclamation()
919
+ assert manager.allocate_slots(request, 64) is not None
920
host = coordinator.host_manager
921
assert host is not None
863
- assert [b.is_null for b in host.req_to_blocks[request.request_id]] == [
864
- False,
865
- True,
866
- True,
867
- True,
868
- ]
922
+ assert not any(b.is_null for b in host.req_to_blocks[request.request_id])
923
_publish_hisparse_pages(manager)
924
for hot in coordinator.hot_managers:
925
hot.require_hot(request.request_id)
926
hot.allocate_new_blocks(request.request_id, 64, 64)
927
coordinator.update_residency(request.request_id)
928
+ assert coordinator.request_states[request.request_id].unpinned_pages == {0, 1}
929
for resident in coordinator.resident_managers:
875
- for block in resident.req_to_blocks[request.request_id][1:]:
876
- assert not block.is_null and block.ref_cnt > 0
877
- assert coordinator.request_states[request.request_id].valid_pages == {0}
878
- pool = coordinator.get_host_block_pool()
879
- assert pool is not None
880
- manager.free(request)
881
- assert pool.get_num_free_blocks() == 1
882
-
883
-
884
-def test_hisparse_host_cow_takes_priority_over_new_pages():
885
- """The last host block must hold the private tail, not a fresh page."""
886
- manager = make_hisparse_kv_cache_manager(32, 3)
887
- coordinator = get_hisparse_coordinator(manager)
888
- host = coordinator.host_manager
889
- assert host is not None
890
- source = host.block_pool.get_new_blocks(1)[0]
891
- host.req_to_blocks["cow"] = [source]
892
- host._partial_hit_reqs["cow"] = (0, source)
893
-
894
- host.allocate_new_blocks("cow", 32, 32)
895
-
896
- private, missing = host.req_to_blocks["cow"]
897
- assert not private.is_null and private is not source
898
- assert missing.is_null
899
- assert host.take_host_cow_copies() == [(source, private)]
930
+ assert [b.ref_cnt for b in resident.req_to_blocks[request.request_id]] == [
931
+ 0,
932
+ 0,
933
+ 1,
934
+ 1,
935
+ ]
936
937
938
@pytest.mark.parametrize("failed", [False, True])
tests/v1/e2e/general/test_hisparse.py
+78
@@ -223,3 +223,81 @@ def test_hisparse_spill_and_prefix_restore(
223
assert actual == expected
224
if with_offloading:
225
assert _offload_load_bytes() > load_bytes
226
+
227
+
228
+@pytest.mark.skipif(
229
+ not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA"
230
+)
231
+@fork_new_process_for_each_test
232
+def test_hisparse_host_exhaustion_defers_requests(
233
+ monkeypatch: pytest.MonkeyPatch,
234
+ vllm_runner: type[VllmRunner],
235
+):
236
+ """A full host pool defers requests instead of leaving pages GPU-only."""
237
+ capability = current_platform.get_device_capability()
238
+ if capability is None or capability.major < 9:
239
+ pytest.skip("Sparse MLA requires Hopper or newer")
240
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
241
+ monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
242
+ prompts = [
243
+ [1000 + request_idx * 64 + i % 64 for i in range(257)]
244
+ for request_idx in range(3)
245
+ ]
246
+ with vllm_runner(
247
+ MODEL,
248
+ load_format="dummy",
249
+ hf_overrides=_shrink_config,
250
+ attention_config=AttentionConfig(
251
+ hisparse_config=HiSparseConfig(device_buffer_size=512)
252
+ ),
253
+ kv_transfer_config=KVTransferConfig(
254
+ kv_connector="HiSparseConnector",
255
+ kv_role="kv_both",
256
+ kv_connector_extra_config={"host_pool_gib": 1},
257
+ ),
258
+ block_size=64,
259
+ max_model_len=320,
260
+ max_num_batched_tokens=512,
261
+ max_num_seqs=3,
262
+ num_gpu_blocks_override=128,
263
+ enable_prefix_caching=False,
264
+ enable_chunked_prefill=True,
265
+ enforce_eager=True,
266
+ ) as runner:
267
+ expected = [runner.generate_greedy([p], max_tokens=16)[0] for p in prompts]
268
+ engine = runner.llm.llm_engine.engine_core.engine_core
269
+ coordinator = get_hisparse_coordinator(engine.scheduler.kv_cache_manager)
270
+ host = coordinator.host_manager
271
+ assert host is not None
272
+ # Leave room for one full request plus part of another.
273
+ spare_blocks = 8
274
+ num_pressure = host.block_pool.get_num_free_blocks() - spare_blocks
275
+ host.allocate_new_blocks(
276
+ "host-pressure",
277
+ num_pressure * host.block_size,
278
+ num_pressure * host.block_size,
279
+ )
280
+ original_plan = coordinator.plan_prefix_materialization
281
+ original_count = host.get_num_blocks_to_allocate
282
+ checked_steps = 0
283
+ refusals = 0
284
+
285
+ def plan(request_id, num_computed_tokens):
286
+ nonlocal checked_steps
287
+ assert not any(b.is_null for b in host.req_to_blocks[request_id])
288
+ checked_steps += 1
289
+ original_plan(request_id, num_computed_tokens)
290
+
291
+ def count(*args, **kwargs):
292
+ nonlocal refusals
293
+ required = original_count(*args, **kwargs)
294
+ refusals += int(required > 0)
295
+ return required
296
+
297
+ monkeypatch.setattr(coordinator, "plan_prefix_materialization", plan)
298
+ monkeypatch.setattr(host, "get_num_blocks_to_allocate", count)
299
+ actual = runner.generate_greedy(prompts, max_tokens=16)
300
+
301
+ assert checked_steps > 0
302
+ assert refusals > 0
303
+ assert actual == expected
vllm/config/vllm.py
+8
@@ -1822,6 +1822,14 @@ class VllmConfig:
1822
raise ValueError(
1823
"HiSparse does not support decode context parallelism."
1824
)
1825
+ if not self.scheduler_config.scheduler_reserve_full_isl:
1826
+ # Without it, async loads admitted against free host blocks can
1827
+ # each wait on host pages the others hold, and waiting requests
1828
+ # are never preempted to free them.
1829
+ raise ValueError(
1830
+ "HiSparse requires --scheduler-reserve-full-isl; remove "
1831
+ "--no-scheduler-reserve-full-isl."
1832
+ )
1833
if self.model_config is not None and not hasattr(
1834
self.model_config.hf_config, "index_topk"
1835
):
vllm/v1/core/single_type_kv_cache_manager.py
+15
-38
@@ -2223,8 +2223,8 @@ class SinkFullAttentionManager(FullAttentionManager):
2223
class HiSparseSourceManager(FullAttentionManager):
2224
"""Host-tier manager with a private pool; publishes hashes once durable.
2225
2226
- Host capacity is best effort: a page that cannot get a host block keeps a
2227
- null host entry, so its GPU copy is never written back and stays pinned.
2226
+ Every page gets a host block: allocation fails when the host pool cannot
2227
+ back it, so each GPU-resident page can be written back and later reclaimed.
2228
"""
2229
2230
coordinator: "HiSparseCoordinator | None" = None
@@ -2273,44 +2273,21 @@ class HiSparseSourceManager(FullAttentionManager):
2273
num_tokens_main_model: int,
2274
apply_admission_cap: bool = False,
2275
) -> int:
2276
- if (
2277
- total_computed_tokens > num_local_computed_tokens
2278
- or self._has_partial_local_hit(
2279
- new_computed_blocks, num_local_computed_tokens
2280
- )
2281
- ):
2282
- # External loads need real destinations; future GPU-computed pages
2283
- # remain best effort. Use the same admission sentinel as Mamba.
2284
- required = super().get_num_blocks_to_allocate(
2285
- request_id,
2286
- total_computed_tokens,
2287
- new_computed_blocks,
2288
- total_computed_tokens,
2289
- num_local_computed_tokens,
2290
- total_computed_tokens,
2291
- )
2292
- if required > self.block_pool.get_num_free_blocks():
2293
- assert self.coordinator is not None
2294
- assert self.coordinator.gpu_pool is not None
2295
- return self.coordinator.gpu_pool.num_gpu_blocks + 1
2296
- return 0
2297
-
2298
- def allocate_new_blocks(
2299
- self, request_id: str, num_tokens: int, num_tokens_main_model: int
2300
- ) -> list[KVCacheBlock]:
2301
- req_blocks = self.req_to_blocks[request_id]
2302
- num_required = cdiv(num_tokens, self.block_size)
2303
- num_free = self.block_pool.get_num_free_blocks()
2304
- if request_id in self._partial_hit_reqs:
2305
- num_free -= 1
2306
- fit_blocks = min(num_required, len(req_blocks) + max(0, num_free))
2307
- new_blocks = super().allocate_new_blocks(
2276
+ required = super().get_num_blocks_to_allocate(
2277
request_id,
2309
- fit_blocks * self.block_size,
2310
- min(num_tokens_main_model, fit_blocks * self.block_size),
2278
+ num_tokens,
2279
+ new_computed_blocks,
2280
+ total_computed_tokens,
2281
+ num_local_computed_tokens,
2282
+ num_tokens_main_model,
2283
)
2312
- req_blocks.extend([self._null_block] * (num_required - len(req_blocks)))
2313
- return new_blocks
2284
+ # Host blocks come from a private pool, so they never count against
2285
+ # the GPU pool. Use the same admission sentinel as Mamba instead.
2286
+ if required > self.block_pool.get_num_free_blocks():
2287
+ assert self.coordinator is not None
2288
+ assert self.coordinator.gpu_pool is not None
2289
+ return self.coordinator.gpu_pool.num_gpu_blocks + 1
2290
+ return 0
2291
2292
def allocate_external_computed_blocks(
2293
self,
vllm/v1/hisparse/coordinator.py
+6
-3
@@ -660,10 +660,13 @@ class HiSparseCoordinator:
660
return command
661
662
def has_pending_reclamation(self) -> bool:
663
- """Whether an in-flight spill will free GPU blocks on completion."""
663
+ """Whether an in-flight spill will free GPU or host blocks on completion."""
664
return any(
665
- pending.page[0] in self.request_states
666
- and self._can_read_from_host(pending.page[0])
665
+ pending.host_block.ref_cnt == 1
666
+ or (
667
+ pending.page[0] in self.request_states
668
+ and self._can_read_from_host(pending.page[0])
669
+ )
670
for pending in self.pending_spills.values()
671
)
672
vllm/v1/hisparse/layout.py
+8
@@ -235,6 +235,14 @@ def create_hisparse_layout(
235
host_num_blocks = host_budget // host_block_stride
236
if host_num_blocks <= 0:
237
raise ValueError("HiSparse has no allocatable host blocks.")
238
+ # Every computed page needs a host block, so one request at max_model_len
239
+ # must fit alongside the pool's null block and a copy-on-write tail.
240
+ min_host_blocks = cdiv(vllm_config.model_config.max_model_len, gpu_block_size) + 2
241
+ if host_num_blocks < min_host_blocks:
242
+ raise ValueError(
243
+ f"HiSparse host pool has {host_num_blocks} blocks but max_model_len "
244
+ f"needs {min_host_blocks}; increase host_pool_gib."
245
+ )
246
247
return HiSparseLayout(
248
source_group=source_group,