@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7b1bc0a3eb

Fork 0

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 4, 15:00 UTC

[Bugfix] Cap SWA/chunked-local runtime admission to startup pool-sizing bound (#40946)

Signed-off-by: Dao Le <Dao007forever@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Dao007forever committed Apr 26, 2026 at 21:33 UTC 7b1bc0a3eb01a6bc2650eda9970049f7825240d7

8 files changed +213 -23

tests/v1/core/test_prefix_caching.py

+108

@@ -2512,3 +2512,111 @@ def test_block_lookup_cache_multi_blocks_per_key():

2512 assert cache.pop(key1, 11) is block11

2513 assert cache.get_one_block(key1) is None

2514 assert cache.pop(key1, 12) is None

2515 +

2516 +

2517 +def test_can_fit_full_sequence_swa_cap_admits_long_prompt():

2518 + """Hybrid full+SWA model with a pool sized at the startup minimum should

2519 + admit a prompt longer than the SWA cap, because SlidingWindowManager

2520 + recycles blocks during chunked prefill (issue #39734)."""

2521 + block_size = 16

2522 + sliding_window = 4 * block_size # 64 tokens

2523 + max_num_batched_tokens = 8 * block_size # 128 tokens

2524 + max_model_len = 64 * block_size # 1024 tokens — much larger than the SWA cap

2525 + # Startup pool sizing: full demands cdiv(max_model_len, bs) = 64 blocks,

2526 + # SWA demands cdiv(SW-1+max_batched, bs) + 1 = cdiv(191, 16) + 1 = 13.

2527 + # Pool minimum = 64 + 13 = 77; +1 for the null block.

2528 + num_blocks = 64 + 13 + 1

2529 +

2530 + config = KVCacheConfig(

2531 + num_blocks=num_blocks,

2532 + kv_cache_tensors=[],

2533 + kv_cache_groups=[

2534 + KVCacheGroupSpec(

2535 + ["layer_full"],

2536 + FullAttentionSpec(

2537 + block_size=block_size,

2538 + num_kv_heads=1,

2539 + head_size=1,

2540 + dtype=torch.float32,

2541 + ),

2542 + ),

2543 + KVCacheGroupSpec(

2544 + ["layer_swa"],

2545 + SlidingWindowSpec(

2546 + block_size=block_size,

2547 + num_kv_heads=1,

2548 + head_size=1,

2549 + dtype=torch.float32,

2550 + sliding_window=sliding_window,

2551 + ),

2552 + ),

2553 + ],

2554 + )

2555 +

2556 + manager = KVCacheManager(

2557 + config,

2558 + max_model_len=max_model_len,

2559 + max_num_batched_tokens=max_num_batched_tokens,

2560 + enable_caching=True,

2561 + hash_block_size=block_size,

2562 + )

2563 +

2564 + # A prompt that is shorter than max_model_len but longer than SW + chunk:

2565 + # cdiv(prompt_len, bs) = 32 blocks. Without the cap, admission would

2566 + # demand 32 (full) + 32 (SWA) = 64 blocks. With the cap, SWA contributes

2567 + # only 13, so total = 32 + 13 = 45 ≤ pool size.

2568 + prompt_len = 32 * block_size

2569 + req = make_request("long", list(range(prompt_len)), block_size, sha256)

2570 +

2571 + assert manager.can_fit_full_sequence(req)

2572 +

2573 +

2574 +def test_can_fit_full_sequence_full_attention_still_gates_oversized():

2575 + """The cap only loosens the SWA group; a prompt that exceeds the

2576 + full-attention pool capacity must still be rejected."""

2577 + block_size = 16

2578 + sliding_window = 4 * block_size

2579 + max_num_batched_tokens = 8 * block_size

2580 + max_model_len = 64 * block_size

2581 + # Provide a tiny pool — even a small prompt should be rejected.

2582 + num_blocks = 5

2583 +

2584 + config = KVCacheConfig(

2585 + num_blocks=num_blocks,

2586 + kv_cache_tensors=[],

2587 + kv_cache_groups=[

2588 + KVCacheGroupSpec(

2589 + ["layer_full"],

2590 + FullAttentionSpec(

2591 + block_size=block_size,

2592 + num_kv_heads=1,

2593 + head_size=1,

2594 + dtype=torch.float32,

2595 + ),

2596 + ),

2597 + KVCacheGroupSpec(

2598 + ["layer_swa"],

2599 + SlidingWindowSpec(

2600 + block_size=block_size,

2601 + num_kv_heads=1,

2602 + head_size=1,

2603 + dtype=torch.float32,

2604 + sliding_window=sliding_window,

2605 + ),

2606 + ),

2607 + ],

2608 + )

2609 +

2610 + manager = KVCacheManager(

2611 + config,

2612 + max_model_len=max_model_len,

2613 + max_num_batched_tokens=max_num_batched_tokens,

2614 + enable_caching=True,

2615 + hash_block_size=block_size,

2616 + )

2617 +

2618 + # 16 blocks of full attention demand alone exceeds the 5-block pool.

2619 + prompt_len = 16 * block_size

2620 + req = make_request("oversized", list(range(prompt_len)), block_size, sha256)

2621 +

2622 + assert not manager.can_fit_full_sequence(req)

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -22,11 +22,13 @@ pytestmark = pytest.mark.cpu_test

24 def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):

25 + # Tests don't exercise admission gating; pass a large cap that is a no-op.

26 return SlidingWindowManager(

27 sliding_window_spec,

28 block_pool=block_pool,

29 enable_caching=enable_caching,

30 kv_cache_group_id=0,

31 + max_admission_blocks_per_request=10**9,

32 )

@@ -38,6 +40,7 @@ def get_chunked_local_attention_manager(

40 block_pool=block_pool,

41 enable_caching=enable_caching,

42 kv_cache_group_id=0,

43 + max_admission_blocks_per_request=10**9,

44 )

vllm/v1/core/kv_cache_coordinator.py

+13

@@ -34,6 +34,7 @@ class KVCacheCoordinator(ABC):

34 self,

35 kv_cache_config: KVCacheConfig,

36 max_model_len: int,

37 + max_num_batched_tokens: int,

38 use_eagle: bool,

39 enable_caching: bool,

40 enable_kv_cache_events: bool,

@@ -65,6 +66,8 @@ class KVCacheCoordinator(ABC):

66 self.single_type_managers = tuple(

67 get_manager_for_kv_cache_spec(

68 kv_cache_spec=kv_cache_group.kv_cache_spec,

69 + max_num_batched_tokens=max_num_batched_tokens,

70 + max_model_len=max_model_len,

71 block_pool=self.block_pool,

72 enable_caching=enable_caching,

73 kv_cache_group_id=i,

@@ -271,6 +274,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

274 self,

275 kv_cache_config: KVCacheConfig,

276 max_model_len: int,

277 + max_num_batched_tokens: int,

278 use_eagle: bool,

279 enable_kv_cache_events: bool,

280 dcp_world_size: int,

@@ -281,6 +285,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

285 super().__init__(

286 kv_cache_config,

287 max_model_len,

288 + max_num_batched_tokens,

289 use_eagle,

290 False,

291 enable_kv_cache_events,

@@ -316,6 +321,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

321 self,

322 kv_cache_config: KVCacheConfig,

323 max_model_len: int,

324 + max_num_batched_tokens: int,

325 use_eagle: bool,

326 enable_caching: bool,

327 enable_kv_cache_events: bool,

@@ -327,6 +333,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

333 super().__init__(

334 kv_cache_config,

335 max_model_len,

336 + max_num_batched_tokens,

337 use_eagle,

338 enable_caching,

339 enable_kv_cache_events,

@@ -381,6 +388,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

388 self,

389 kv_cache_config: KVCacheConfig,

390 max_model_len: int,

391 + max_num_batched_tokens: int,

392 use_eagle: bool,

393 enable_caching: bool,

394 enable_kv_cache_events: bool,

@@ -392,6 +400,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

400 super().__init__(

401 kv_cache_config,

402 max_model_len,

403 + max_num_batched_tokens,

404 use_eagle,

405 enable_caching,

406 enable_kv_cache_events,

@@ -574,6 +583,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

583 def get_kv_cache_coordinator(

584 kv_cache_config: KVCacheConfig,

585 max_model_len: int,

586 + max_num_batched_tokens: int,

587 use_eagle: bool,

588 enable_caching: bool,

589 enable_kv_cache_events: bool,

@@ -586,6 +596,7 @@ def get_kv_cache_coordinator(

596 return KVCacheCoordinatorNoPrefixCache(

597 kv_cache_config,

598 max_model_len,

599 + max_num_batched_tokens,

600 use_eagle,

601 enable_kv_cache_events,

602 dcp_world_size=dcp_world_size,

@@ -597,6 +608,7 @@ def get_kv_cache_coordinator(

608 return UnitaryKVCacheCoordinator(

609 kv_cache_config,

610 max_model_len,

611 + max_num_batched_tokens,

612 use_eagle,

613 enable_caching,

614 enable_kv_cache_events,

@@ -608,6 +620,7 @@ def get_kv_cache_coordinator(

620 return HybridKVCacheCoordinator(

621 kv_cache_config,

622 max_model_len,

623 + max_num_batched_tokens,

624 use_eagle,

625 enable_caching,

626 enable_kv_cache_events,

vllm/v1/core/kv_cache_manager.py

@@ -109,6 +109,7 @@ class KVCacheManager:

109 kv_cache_config: KVCacheConfig,

110 max_model_len: int,

111 hash_block_size: int,

112 + max_num_batched_tokens: int | None = None,

113 enable_caching: bool = True,

114 use_eagle: bool = False,

115 log_stats: bool = False,

@@ -118,6 +119,11 @@ class KVCacheManager:

119 metrics_collector: KVCacheMetricsCollector | None = None,

120 ) -> None:

121 self.max_model_len = max_model_len

122 + # When unset, fall back to `max_model_len` so the recycling-aware cap

123 + # collapses to the prior (uncapped) admission behavior. The scheduler

124 + # always supplies the real value at runtime.

125 + if max_num_batched_tokens is None:

126 + max_num_batched_tokens = max_model_len

127

128 self.enable_caching = enable_caching

129 self.use_eagle = use_eagle

@@ -131,6 +137,7 @@ class KVCacheManager:

137 self.coordinator = get_kv_cache_coordinator(

138 kv_cache_config=kv_cache_config,

139 max_model_len=self.max_model_len,

140 + max_num_batched_tokens=max_num_batched_tokens,

141 use_eagle=self.use_eagle,

142 enable_caching=self.enable_caching,

143 enable_kv_cache_events=enable_kv_cache_events,

vllm/v1/core/sched/scheduler.py

@@ -228,6 +228,7 @@ class Scheduler(SchedulerInterface):

228 self.kv_cache_manager = KVCacheManager(

229 kv_cache_config=kv_cache_config,

230 max_model_len=self.max_model_len,

231 + max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,

232 enable_caching=self.cache_config.enable_prefix_caching,

233 use_eagle=self.use_eagle,

234 log_stats=self.log_stats,

vllm/v1/core/single_type_kv_cache_manager.py

+35 -1

@@ -41,6 +41,7 @@ class SingleTypeKVCacheManager(ABC):

41 kv_cache_group_id: int,

42 dcp_world_size: int = 1,

43 pcp_world_size: int = 1,

44 + max_admission_blocks_per_request: int | None = None,

45 ) -> None:

46 """

47 Initializes the SingleTypeKVCacheManager.

@@ -48,6 +49,12 @@ class SingleTypeKVCacheManager(ABC):

49 kv_cache_spec: The kv_cache_spec for this manager.

50 block_pool: The block pool.

51 kv_cache_group_id: The id of the kv cache group of this manager.

52 + max_admission_blocks_per_request: Recycling-aware per-request

53 + block cap used by `get_num_blocks_to_allocate`. Only set for

54 + spec types that recycle blocks across chunks (SWA,

55 + chunked-local); `None` (the default) means no cap, which is

56 + correct for full-attention-style specs that hold every

57 + block until the request finishes.

58 """

59 self.block_size = kv_cache_spec.block_size

60 self.dcp_world_size = dcp_world_size

@@ -57,6 +64,7 @@ class SingleTypeKVCacheManager(ABC):

64 self.kv_cache_spec = kv_cache_spec

65 self.block_pool = block_pool

66 self.enable_caching = enable_caching

67 + self._max_admission_blocks_per_request = max_admission_blocks_per_request

68 self.new_block_ids: list[int] = []

70 # Mapping from request ID to blocks to track the blocks allocated

@@ -105,6 +113,19 @@ class SingleTypeKVCacheManager(ABC):

113 """

114

115 num_required_blocks = cdiv(num_tokens, self.block_size)

116 + if self._max_admission_blocks_per_request is not None:

117 + # Recycling-aware specs (SWA, chunked-local) cap the per-request

118 + # reservation here so admission matches the startup pool sizer

119 + # (`SlidingWindowSpec.max_admission_blocks_per_request` / its

120 + # chunked-local counterpart). `remove_skipped_blocks` runs from

121 + # `allocate_slots` before each chunk's `get_num_blocks_to_allocate`,

122 + # so per-request peak real-held blocks <= this cap, which keeps

123 + # `sum(reservations) <= pool` <=> `sum(peak_real_held) <= pool`.

124 + # Drift between the two would re-introduce the deadlock from

125 + # issue #39734 or, worse, mid-prefill OOM.

126 + num_required_blocks = min(

127 + num_required_blocks, self._max_admission_blocks_per_request

128 + )

129 num_req_blocks = len(self.req_to_blocks.get(request_id, ()))

130

131 if request_id in self.num_cached_block:

@@ -1126,8 +1147,21 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

1147

1148

1149 def get_manager_for_kv_cache_spec(

1129 - kv_cache_spec: KVCacheSpec, **kwargs

1150 + kv_cache_spec: KVCacheSpec,

1151 + max_num_batched_tokens: int,

1152 + max_model_len: int,

1153 + **kwargs,

1154 ) -> SingleTypeKVCacheManager:

1155 manager_class = spec_manager_map[type(kv_cache_spec)]

1156 + # SlidingWindow / ChunkedLocalAttention managers recycle blocks across

1157 + # chunks; the runtime admission cap must match the recycling-aware bound

1158 + # the startup pool sizer uses (single source of truth: the spec method).

1159 + if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):

1160 + kwargs["max_admission_blocks_per_request"] = (

1161 + kv_cache_spec.max_admission_blocks_per_request(

1162 + max_num_batched_tokens=max_num_batched_tokens,

1163 + max_model_len=max_model_len,

1164 + )

1165 + )

1166 manager = manager_class(kv_cache_spec, **kwargs)

1167 return manager

vllm/v1/kv_cache_interface.py

+43 -22

@@ -376,19 +376,28 @@ class MLAAttentionSpec(FullAttentionSpec):

376 class ChunkedLocalAttentionSpec(AttentionSpec):

377 attention_chunk_size: int

378

379 - def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

380 - max_model_len = vllm_config.model_config.max_model_len

381 - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

382 -

383 - # During chunked prefill, we allocate KV cache for at most

384 - # `self.attention_chunk_size` computed tokens plus the newly scheduled

385 - # tokens. And we won't allocate KV cache for more than `max_model_len`

386 - # tokens.

379 + def max_admission_blocks_per_request(

380 + self, max_num_batched_tokens: int, max_model_len: int

381 + ) -> int:

382 + """Per-request admission cap, in blocks.

383 +

384 + Single source of truth for both startup pool sizing

385 + (`max_memory_usage_bytes`) and the runtime admission gate, so requests

386 + admitted by startup can also be admitted at runtime.

387 + """

388 + # During chunked prefill, we hold KV for at most one chunk window.

389 num_tokens = min(

390 self.attention_chunk_size + max_num_batched_tokens, max_model_len

391 )

392 + return cdiv(num_tokens, self.block_size)

393

391 - return cdiv(num_tokens, self.block_size) * self.page_size_bytes

394 + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

395 + max_model_len = vllm_config.model_config.max_model_len

396 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

397 + max_blocks = self.max_admission_blocks_per_request(

398 + max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len

399 + )

400 + return max_blocks * self.page_size_bytes

401

402

403 @dataclass(frozen=True, kw_only=True)

@@ -409,26 +418,38 @@ class SlidingWindowSpec(AttentionSpec):

418 * get_dtype_size(self.dtype)

419 )

420

421 + def max_admission_blocks_per_request(

422 + self, max_num_batched_tokens: int, max_model_len: int

423 + ) -> int:

424 + """Per-request admission cap, in blocks.

425 +

426 + Single source of truth for both startup pool sizing

427 + (`max_memory_usage_bytes`) and the runtime admission gate. Per-request

428 + real-held blocks plateau at this bound because

429 + `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`

430 + before each chunk's `get_num_blocks_to_allocate`.

431 + """

432 + # During chunked prefill, we hold KV for the last `sliding_window-1`

433 + # computed tokens plus the newly scheduled tokens, and never more

434 + # than `max_model_len`.

435 + num_tokens = min(

436 + self.sliding_window - 1 + max_num_batched_tokens, max_model_len

437 + )

438 + # +1 because the sliding window may not start from the beginning of

439 + # the block. E.g. block size 4 and num_token 4 needs two blocks

440 + # [XXCD][EF] to store the 6-token window [CDEF].

441 + return cdiv(num_tokens, self.block_size) + 1

442 +

443 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

444 assert vllm_config.parallel_config.decode_context_parallel_size == 1, (

445 "DCP not support sliding window."

446 )

447 max_model_len = vllm_config.model_config.max_model_len

448 max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

418 -

419 - # During chunked prefill, we allocate KV cache for the last

420 - # `self.sliding_window-1` computed tokens plus the newly scheduled

421 - # tokens. And we won't allocate KV cache for more than `max_model_len`

422 - # tokens.

423 - num_tokens = min(

424 - self.sliding_window - 1 + max_num_batched_tokens, max_model_len

449 + max_blocks = self.max_admission_blocks_per_request(

450 + max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len

451 )

426 -

427 - # +1 here because the sliding window may not start from the beginning

428 - # of the block. For example, if the block size is 4 and num_token

429 - # is 4, we need two blocks [XXCD] [EF] to store the sliding

430 - # window [CDEF] of 6 tokens.

431 - return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes

452 + return max_blocks * self.page_size_bytes

453

454

455 @dataclass(frozen=True, kw_only=True)

vllm/v1/simple_kv_offload/manager.py

@@ -110,6 +110,9 @@ class SimpleCPUOffloadScheduler:

110 self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(

111 kv_cache_config=self.cpu_kv_cache_config,

112 max_model_len=vllm_config.model_config.max_model_len,

113 + max_num_batched_tokens=(

114 + vllm_config.scheduler_config.max_num_batched_tokens

115 + ),

116 use_eagle=False,

117 enable_caching=True,

118 enable_kv_cache_events=self.enable_kv_cache_events,