A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 06:00 UTC
[Core][Refactor]: thread `scheduler_block_size` into KVCacheManager and KVCacheCoordinator (#44165)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Yifan Qiao committed
Jun 2, 2026 at 01:14 UTC
7c37096620fa4e161c1d8c1db5c43c8514545d84
9 files changed
+99
-50
tests/v1/core/test_kv_cache_utils.py
+12
-3
@@ -1447,7 +1447,10 @@ def test_allocate_with_lookahead():
1447
1448
# Test case 1: Requires additional lookahead tokens
1449
kv_cache_manager = KVCacheManager(
1450
- kv_cache_config=config, max_model_len=100, hash_block_size=block_size
1450
+ kv_cache_config=config,
1451
+ max_model_len=100,
1452
+ scheduler_block_size=block_size,
1453
+ hash_block_size=block_size,
1454
)
1455
blocks = kv_cache_manager.allocate_slots(
1456
request,
@@ -1458,7 +1461,10 @@ def test_allocate_with_lookahead():
1461
1462
# Test case 2: With precomputed blocks
1463
kv_cache_manager = KVCacheManager(
1461
- kv_cache_config=config, max_model_len=100, hash_block_size=block_size
1464
+ kv_cache_config=config,
1465
+ max_model_len=100,
1466
+ scheduler_block_size=block_size,
1467
+ hash_block_size=block_size,
1468
)
1469
# required_blocks = ceil((3 + 2) /4) = 2
1470
blocks = kv_cache_manager.allocate_slots(
@@ -1471,7 +1477,10 @@ def test_allocate_with_lookahead():
1477
# Test case 3: With precomputed blocks
1478
# required_blocks = ceil((3 + 4) / 4) = 2
1479
kv_cache_manager = KVCacheManager(
1474
- kv_cache_config=config, max_model_len=100, hash_block_size=block_size
1480
+ kv_cache_config=config,
1481
+ max_model_len=100,
1482
+ scheduler_block_size=block_size,
1483
+ hash_block_size=block_size,
1484
)
1485
blocks = kv_cache_manager.allocate_slots(
1486
request,
tests/v1/core/test_prefix_caching.py
+43
-30
@@ -4,6 +4,7 @@
4
5
import copy
6
from collections.abc import Callable
7
+from math import lcm
8
9
import pytest
10
import torch
@@ -92,6 +93,18 @@ def make_request(
93
)
94
95
96
+def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheManager:
97
+ """Build a ``KVCacheManager``, deriving ``scheduler_block_size`` from the
98
+ config (LCM of group block sizes) unless explicitly provided. This mirrors
99
+ ``resolve_kv_cache_block_sizes`` for the non-context-parallel case used by
100
+ these tests, so callers don't have to pass it at every site."""
101
+ kwargs.setdefault(
102
+ "scheduler_block_size",
103
+ lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)),
104
+ )
105
+ return KVCacheManager(kv_cache_config, **kwargs)
106
+
107
+
108
def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
109
return KVCacheConfig(
110
num_blocks=num_blocks,
@@ -208,7 +221,7 @@ def make_kv_cache_config_three_types(
221
@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
222
def test_prefill(hash_fn):
223
block_size = 16
211
- manager = KVCacheManager(
224
+ manager = make_kv_cache_manager(
225
make_kv_cache_config(block_size, 11),
226
max_model_len=8192,
227
enable_caching=True,
@@ -331,7 +344,7 @@ def test_prefill(hash_fn):
344
345
def test_prefill_hybrid_model():
346
block_size = 16
334
- manager = KVCacheManager(
347
+ manager = make_kv_cache_manager(
348
make_kv_cache_config_hybrid_model(block_size, 21, 2),
349
max_model_len=8192,
350
enable_caching=True,
@@ -500,7 +513,7 @@ def test_prefill_hybrid_model():
513
def test_prefill_hybrid_model_eagle():
514
block_size = 16
515
kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 3)
503
- manager = KVCacheManager(
516
+ manager = make_kv_cache_manager(
517
kv_cache_config,
518
max_model_len=8192,
519
enable_caching=True,
@@ -837,7 +850,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
850
num_blocks = 10 * num_groups
851
852
kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)
840
- manager = KVCacheManager(
853
+ manager = make_kv_cache_manager(
854
kv_cache_config,
855
max_model_len=8192,
856
enable_caching=True,
@@ -912,7 +925,7 @@ def test_prefill_hybrid_model_combinations_eagle(
925
num_blocks = 10 * num_groups
926
927
kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)
915
- manager = KVCacheManager(
928
+ manager = make_kv_cache_manager(
929
kv_cache_config,
930
max_model_len=8192,
931
enable_caching=True,
@@ -984,7 +997,7 @@ def test_prefill_hybrid_model_mamba_align():
997
kv_cache_config = _make_hybrid_kv_cache_config(
998
block_size, num_blocks, ["full", "mamba_align"]
999
)
987
- manager = KVCacheManager(
1000
+ manager = make_kv_cache_manager(
1001
kv_cache_config,
1002
max_model_len=8192,
1003
enable_caching=True,
@@ -1017,7 +1030,7 @@ def test_prefill_plp():
1030
3. Schedule plp request; no hit should occur; validate blocks
1031
"""
1032
block_size = 16
1020
- manager = KVCacheManager(
1033
+ manager = make_kv_cache_manager(
1034
make_kv_cache_config(block_size, 11),
1035
max_model_len=8192,
1036
enable_caching=True,
@@ -1125,7 +1138,7 @@ def test_prefill_plp():
1138
1139
def test_decode():
1140
block_size = 16
1128
- manager = KVCacheManager(
1141
+ manager = make_kv_cache_manager(
1142
make_kv_cache_config(block_size, 11),
1143
max_model_len=8192,
1144
enable_caching=True,
@@ -1188,7 +1201,7 @@ def test_decode():
1201
1202
def test_evict():
1203
block_size = 16
1191
- manager = KVCacheManager(
1204
+ manager = make_kv_cache_manager(
1205
make_kv_cache_config(block_size, 11),
1206
max_model_len=8192,
1207
enable_caching=True,
@@ -1247,7 +1260,7 @@ def test_hash_block_correct_reuse():
1260
its hash metadata should be correctly reset.
1261
"""
1262
block_size = 16
1250
- manager = KVCacheManager(
1263
+ manager = make_kv_cache_manager(
1264
make_kv_cache_config(16, 2),
1265
max_model_len=8192,
1266
enable_caching=True,
@@ -1288,7 +1301,7 @@ def test_computed_blocks_not_evicted():
1301
for a request if there are any other free blocks.
1302
"""
1303
block_size = 16
1291
- manager = KVCacheManager(
1304
+ manager = make_kv_cache_manager(
1305
make_kv_cache_config(block_size, 3),
1306
max_model_len=8192,
1307
enable_caching=True,
@@ -1347,7 +1360,7 @@ def test_basic_prefix_caching_disabled():
1360
This tests that the prefix caching is disabled.
1361
"""
1362
block_size = 4
1350
- manager = KVCacheManager(
1363
+ manager = make_kv_cache_manager(
1364
make_kv_cache_config(block_size, 5),
1365
max_model_len=8192,
1366
enable_caching=False,
@@ -1531,7 +1544,7 @@ def test_mm_prefix_caching():
1544
"""
1545
1546
block_size = 16
1534
- manager = KVCacheManager(
1547
+ manager = make_kv_cache_manager(
1548
make_kv_cache_config(block_size, 11),
1549
max_model_len=8192,
1550
enable_caching=True,
@@ -1639,7 +1652,7 @@ def test_cache_key_salting():
1652
is separated cache as expected.
1653
"""
1654
block_size = 16
1642
- manager = KVCacheManager(
1655
+ manager = make_kv_cache_manager(
1656
make_kv_cache_config(block_size, 11),
1657
max_model_len=8192,
1658
enable_caching=True,
@@ -1721,7 +1734,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1734
the computed blocks should not be touched.
1735
"""
1736
block_size = 16
1724
- manager = KVCacheManager(
1737
+ manager = make_kv_cache_manager(
1738
make_kv_cache_config(block_size, 11),
1739
max_model_len=8192,
1740
enable_caching=True,
@@ -1794,7 +1807,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1807
1808
def test_reset_prefix_cache():
1809
block_size = 16
1797
- manager = KVCacheManager(
1810
+ manager = make_kv_cache_manager(
1811
make_kv_cache_config(block_size, 11),
1812
max_model_len=8192,
1813
enable_caching=True,
@@ -1835,7 +1848,7 @@ def test_reset_prefix_cache():
1848
def test_prefix_cache_stats_disabled():
1849
"""Test that prefix_cache_stats is None when log_stats is False."""
1850
block_size = 16
1838
- manager = KVCacheManager(
1851
+ manager = make_kv_cache_manager(
1852
make_kv_cache_config(block_size, 11),
1853
max_model_len=8192,
1854
enable_caching=True,
@@ -1915,7 +1928,7 @@ def test_kv_cache_events(blocks_to_cache: int):
1928
# Should see a single block stored event with a blocks_to_cache number of
1929
# block hashes
1930
# take_events should reset the kv_event_queue
1918
- manager = KVCacheManager(
1931
+ manager = make_kv_cache_manager(
1932
make_kv_cache_config(block_size, num_blocks),
1933
max_model_len=8192,
1934
enable_caching=True,
@@ -2043,7 +2056,7 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int):
2056
num_blocks = blocks_to_cache + 1
2057
2058
# Create KVCacheManager with events enabled
2046
- manager = KVCacheManager(
2059
+ manager = make_kv_cache_manager(
2060
make_kv_cache_config(block_size, num_blocks),
2061
max_model_len=8192,
2062
enable_caching=True,
@@ -2101,7 +2114,7 @@ def test_block_stored_event_group_idx(group_id: int):
2114
block_size = 4
2115
num_tokens = block_size * 2
2116
2104
- manager = KVCacheManager(
2117
+ manager = make_kv_cache_manager(
2118
make_kv_cache_config_three_types(block_size, num_blocks=5),
2119
max_model_len=8192,
2120
enable_caching=True,
@@ -2161,7 +2174,7 @@ def test_block_stored_event_group_idx_multiple_groups():
2174
block_size = 4
2175
num_tokens = block_size * 2
2176
2164
- manager = KVCacheManager(
2177
+ manager = make_kv_cache_manager(
2178
KVCacheConfig(
2179
num_blocks=5,
2180
kv_cache_tensors=[],
@@ -2238,7 +2251,7 @@ def test_block_stored_event_group_idx_multiple_groups():
2251
def test_block_stored_event_group_idx_out_of_bounds(monkeypatch):
2252
"""Out-of-range group_idx events are returned without metadata annotation."""
2253
block_size = 4
2241
- manager = KVCacheManager(
2254
+ manager = make_kv_cache_manager(
2255
make_kv_cache_config(block_size, num_blocks=5),
2256
max_model_len=8192,
2257
enable_caching=True,
@@ -2328,7 +2341,7 @@ def test_eagle_enabled_removes_last_block():
2341
"""Verify Eagle does NOT remove blocks when request
2342
length is divisible by block size."""
2343
block_size = 16
2331
- manager = KVCacheManager(
2344
+ manager = make_kv_cache_manager(
2345
make_kv_cache_config(block_size, num_blocks=10),
2346
max_model_len=8192,
2347
enable_caching=True,
@@ -2361,7 +2374,7 @@ def test_eagle_enabled_removes_last_block():
2374
def test_eagle_with_partial_blocks():
2375
"""Test Eagle behavior with requests containing partial blocks."""
2376
block_size = 16
2364
- manager = KVCacheManager(
2377
+ manager = make_kv_cache_manager(
2378
make_kv_cache_config(block_size, num_blocks=10),
2379
max_model_len=8192,
2380
enable_caching=True,
@@ -2397,7 +2410,7 @@ def test_eagle_with_sliding_window():
2410
dtype=torch.float32,
2411
sliding_window=block_size,
2412
)
2400
- manager = KVCacheManager(
2413
+ manager = make_kv_cache_manager(
2414
KVCacheConfig(
2415
num_blocks=10,
2416
kv_cache_tensors=[],
@@ -2482,7 +2495,7 @@ def test_different_block_size():
2495
),
2496
],
2497
)
2485
- manager = KVCacheManager(
2498
+ manager = make_kv_cache_manager(
2499
kv_cache_config=kv_cache_config,
2500
max_model_len=8192,
2501
enable_caching=True,
@@ -2565,7 +2578,7 @@ def test_hybrid_cache_blocks_swa_tail_window_only():
2578
),
2579
],
2580
)
2568
- manager = KVCacheManager(
2581
+ manager = make_kv_cache_manager(
2582
kv_cache_config=kv_cache_config,
2583
max_model_len=8192,
2584
enable_caching=True,
@@ -2633,7 +2646,7 @@ def test_hybrid_cache_blocks_clamped_to_lcm():
2646
),
2647
],
2648
)
2636
- manager = KVCacheManager(
2649
+ manager = make_kv_cache_manager(
2650
kv_cache_config=kv_cache_config,
2651
max_model_len=8192,
2652
enable_caching=True,
@@ -2781,7 +2794,7 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
2794
],
2795
)
2796
2784
- manager = KVCacheManager(
2797
+ manager = make_kv_cache_manager(
2798
config,
2799
max_model_len=max_model_len,
2800
max_num_batched_tokens=max_num_batched_tokens,
@@ -2837,7 +2850,7 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():
2850
],
2851
)
2852
2840
- manager = KVCacheManager(
2853
+ manager = make_kv_cache_manager(
2854
config,
2855
max_model_len=max_model_len,
2856
max_num_batched_tokens=max_num_batched_tokens,
tests/v1/core/test_single_type_kv_cache_manager.py
+3
@@ -28,6 +28,7 @@ def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=T
28
block_pool=block_pool,
29
enable_caching=enable_caching,
30
kv_cache_group_id=0,
31
+ scheduler_block_size=sliding_window_spec.block_size,
32
max_admission_blocks_per_request=10**9,
33
)
34
@@ -40,6 +41,7 @@ def get_chunked_local_attention_manager(
41
block_pool=block_pool,
42
enable_caching=enable_caching,
43
kv_cache_group_id=0,
44
+ scheduler_block_size=chunked_local_attention_spec.block_size,
45
max_admission_blocks_per_request=10**9,
46
)
47
@@ -458,6 +460,7 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
460
block_pool=block_pool,
461
enable_caching=False,
462
kv_cache_group_id=0,
463
+ scheduler_block_size=spec.block_size,
464
max_admission_blocks_per_request=cap,
465
)
466
tests/v1/kv_connector/unit/offloading_connector/utils.py
+7
-1
@@ -28,6 +28,7 @@ from vllm.utils.hashing import sha256
28
from vllm.v1.core.kv_cache_utils import (
29
get_request_block_hasher,
30
init_none_hash,
31
+ resolve_kv_cache_block_sizes,
32
)
33
from vllm.v1.core.sched.async_scheduler import AsyncScheduler
34
from vllm.v1.core.sched.scheduler import Scheduler
@@ -230,13 +231,18 @@ class RequestRunner:
231
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
232
self.num_kv_groups = len(kv_cache_config.kv_cache_groups)
233
234
+ scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
235
+ kv_cache_config, vllm_config
236
+ )
237
+
238
scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
239
self.scheduler = scheduler_cls(
240
vllm_config=vllm_config,
241
kv_cache_config=kv_cache_config,
242
log_stats=True,
243
structured_output_manager=StructuredOutputManager(vllm_config),
239
- block_size=block_size,
244
+ block_size=scheduler_block_size,
245
+ hash_block_size=hash_block_size,
246
)
247
248
self.worker_connector = OffloadingConnector(
vllm/v1/core/kv_cache_coordinator.py
+26
-16
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from abc import ABC, abstractmethod
4
from collections.abc import Sequence
5
-from math import lcm
5
6
from vllm.v1.core.block_pool import BlockPool
7
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
@@ -40,12 +39,20 @@ class KVCacheCoordinator(ABC):
39
enable_kv_cache_events: bool,
40
dcp_world_size: int,
41
pcp_world_size: int,
42
+ scheduler_block_size: int,
43
hash_block_size: int,
44
metrics_collector: KVCacheMetricsCollector | None = None,
45
):
46
self.kv_cache_config = kv_cache_config
47
self.max_model_len = max_model_len
48
self.enable_caching = enable_caching
49
+ # The scheduling granularity (LCM of all group block sizes), must be a multiple
50
+ # of the hash_block_size and the block size of each group.
51
+ assert scheduler_block_size % hash_block_size == 0 and all(
52
+ scheduler_block_size % g.kv_cache_spec.block_size == 0
53
+ for g in kv_cache_config.kv_cache_groups
54
+ )
55
+ self.scheduler_block_size = scheduler_block_size
56
57
self.block_pool = BlockPool(
58
num_gpu_blocks=kv_cache_config.num_blocks,
@@ -73,6 +80,7 @@ class KVCacheCoordinator(ABC):
80
kv_cache_group_id=i,
81
dcp_world_size=dcp_world_size,
82
pcp_world_size=pcp_world_size,
83
+ scheduler_block_size=self.scheduler_block_size,
84
)
85
for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
86
)
@@ -290,6 +298,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
298
enable_kv_cache_events: bool,
299
dcp_world_size: int,
300
pcp_world_size: int,
301
+ scheduler_block_size: int,
302
hash_block_size: int,
303
metrics_collector: KVCacheMetricsCollector | None = None,
304
):
@@ -302,6 +311,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
311
enable_kv_cache_events,
312
dcp_world_size=dcp_world_size,
313
pcp_world_size=pcp_world_size,
314
+ scheduler_block_size=scheduler_block_size,
315
hash_block_size=hash_block_size,
316
metrics_collector=metrics_collector,
317
)
@@ -338,6 +348,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
348
enable_kv_cache_events: bool,
349
dcp_world_size: int,
350
pcp_world_size: int,
351
+ scheduler_block_size: int,
352
hash_block_size: int,
353
metrics_collector: KVCacheMetricsCollector | None = None,
354
):
@@ -350,6 +361,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
361
enable_kv_cache_events,
362
dcp_world_size=dcp_world_size,
363
pcp_world_size=pcp_world_size,
364
+ scheduler_block_size=scheduler_block_size,
365
hash_block_size=hash_block_size,
366
metrics_collector=metrics_collector,
367
)
@@ -405,6 +417,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
417
enable_kv_cache_events: bool,
418
dcp_world_size: int,
419
pcp_world_size: int,
420
+ scheduler_block_size: int,
421
hash_block_size: int,
422
metrics_collector: KVCacheMetricsCollector | None = None,
423
):
@@ -417,6 +430,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
430
enable_kv_cache_events,
431
dcp_world_size=dcp_world_size,
432
pcp_world_size=pcp_world_size,
433
+ scheduler_block_size=scheduler_block_size,
434
hash_block_size=hash_block_size,
435
metrics_collector=metrics_collector,
436
)
@@ -468,14 +482,6 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
482
key=lambda x: not isinstance(x[0], FullAttentionSpec),
483
)
484
471
- # The LCM of the block sizes of all attention types.
472
- # The cache hit length must be a multiple of the LCM of the block sizes
473
- # to make sure the cache hit length is a multiple of the block size of
474
- # each attention type. Requiring this because we don't support partial
475
- # block cache hit yet.
476
- block_sizes = [spec.block_size for spec, _, _ in attention_groups]
477
- self.lcm_block_size = lcm(*block_sizes)
478
-
485
# Attention-group indices (into ``self.attention_groups``) that
486
# contain at least one EAGLE/MTP KV cache group.
487
self.eagle_attn_group_indices: set[int] = {
@@ -486,18 +492,18 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
492
493
def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
494
# Cache hits in this coordinator are always a multiple of
489
- # ``lcm_block_size`` tokens (see ``find_longest_cache_hit``). Within an
490
- # aligned region, SWA groups only consult a subset of blocks per
491
- # ``lcm_block_size``-segment so the unused blocks also stay out of the
492
- # prefix-cache hash map.
495
+ # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``).
496
+ # Within an aligned region, SWA groups only consult a subset of blocks
497
+ # per ``scheduler_block_size``-segment so the unused blocks also stay
498
+ # out of the prefix-cache hash map.
499
num_computed_tokens = (
494
- num_computed_tokens // self.lcm_block_size * self.lcm_block_size
500
+ num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size
501
)
502
for manager in self.single_type_managers:
503
manager.cache_blocks(
504
request,
505
num_computed_tokens,
500
- alignment_tokens=self.lcm_block_size,
506
+ alignment_tokens=self.scheduler_block_size,
507
)
508
509
def find_longest_cache_hit(
@@ -576,7 +582,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
582
block_pool=self.block_pool,
583
kv_cache_spec=spec,
584
use_eagle=use_eagle,
579
- alignment_tokens=self.lcm_block_size,
585
+ alignment_tokens=self.scheduler_block_size,
586
)
587
_new_hit_length = len(hit_blocks[0]) * spec.block_size
588
if use_eagle:
@@ -616,6 +622,7 @@ def get_kv_cache_coordinator(
622
enable_kv_cache_events: bool,
623
dcp_world_size: int,
624
pcp_world_size: int,
625
+ scheduler_block_size: int,
626
hash_block_size: int,
627
metrics_collector: KVCacheMetricsCollector | None = None,
628
) -> KVCacheCoordinator:
@@ -628,6 +635,7 @@ def get_kv_cache_coordinator(
635
enable_kv_cache_events,
636
dcp_world_size=dcp_world_size,
637
pcp_world_size=pcp_world_size,
638
+ scheduler_block_size=scheduler_block_size,
639
hash_block_size=hash_block_size,
640
metrics_collector=metrics_collector,
641
)
@@ -641,6 +649,7 @@ def get_kv_cache_coordinator(
649
enable_kv_cache_events,
650
dcp_world_size=dcp_world_size,
651
pcp_world_size=pcp_world_size,
652
+ scheduler_block_size=scheduler_block_size,
653
hash_block_size=hash_block_size,
654
metrics_collector=metrics_collector,
655
)
@@ -653,6 +662,7 @@ def get_kv_cache_coordinator(
662
enable_kv_cache_events,
663
dcp_world_size=dcp_world_size,
664
pcp_world_size=pcp_world_size,
665
+ scheduler_block_size=scheduler_block_size,
666
hash_block_size=hash_block_size,
667
metrics_collector=metrics_collector,
668
)
vllm/v1/core/kv_cache_manager.py
+2
@@ -112,6 +112,7 @@ class KVCacheManager:
112
self,
113
kv_cache_config: KVCacheConfig,
114
max_model_len: int,
115
+ scheduler_block_size: int,
116
hash_block_size: int,
117
max_num_batched_tokens: int | None = None,
118
enable_caching: bool = True,
@@ -147,6 +148,7 @@ class KVCacheManager:
148
enable_kv_cache_events=enable_kv_cache_events,
149
dcp_world_size=dcp_world_size,
150
pcp_world_size=pcp_world_size,
151
+ scheduler_block_size=scheduler_block_size,
152
hash_block_size=hash_block_size,
153
metrics_collector=self.metrics_collector,
154
)
vllm/v1/core/sched/scheduler.py
+1
@@ -237,6 +237,7 @@ class Scheduler(SchedulerInterface):
237
enable_kv_cache_events=self.enable_kv_cache_events,
238
dcp_world_size=self.dcp_world_size,
239
pcp_world_size=self.pcp_world_size,
240
+ scheduler_block_size=self.block_size,
241
hash_block_size=hash_block_size,
242
metrics_collector=self.kv_metrics_collector,
243
)
vllm/v1/core/single_type_kv_cache_manager.py
+4
@@ -40,6 +40,7 @@ class SingleTypeKVCacheManager(ABC):
40
block_pool: BlockPool,
41
enable_caching: bool,
42
kv_cache_group_id: int,
43
+ scheduler_block_size: int,
44
dcp_world_size: int = 1,
45
pcp_world_size: int = 1,
46
max_admission_blocks_per_request: int | None = None,
@@ -50,6 +51,8 @@ class SingleTypeKVCacheManager(ABC):
51
kv_cache_spec: The kv_cache_spec for this manager.
52
block_pool: The block pool.
53
kv_cache_group_id: The id of the kv cache group of this manager.
54
+ scheduler_block_size: The scheduling granularity (LCM of all group
55
+ block sizes); a multiple of this manager's ``block_size``.
56
max_admission_blocks_per_request: Recycling-aware per-request
57
block cap used by `get_num_blocks_to_allocate`. Only set for
58
spec types that recycle blocks across chunks (SWA,
@@ -57,6 +60,7 @@ class SingleTypeKVCacheManager(ABC):
60
correct for full-attention-style specs that hold every
61
block until the request finishes.
62
"""
63
+ self.scheduler_block_size = scheduler_block_size
64
self.block_size = kv_cache_spec.block_size
65
self.dcp_world_size = dcp_world_size
66
self.pcp_world_size = pcp_world_size
vllm/v1/simple_kv_offload/manager.py
+1
@@ -127,6 +127,7 @@ class SimpleCPUOffloadScheduler:
127
enable_kv_cache_events=self.enable_kv_cache_events,
128
dcp_world_size=dcp_world_size,
129
pcp_world_size=pcp_world_size,
130
+ scheduler_block_size=self.block_size,
131
hash_block_size=self.hash_block_size,
132
)
133
self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool