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

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Hybrid] Marconi-style admission policy for hybrid cache (#37898)

Signed-off-by: Stanislaw Wozniak <stw@zurich.ibm.com>

Stan Wozniak committed Jun 10, 2026 at 19:03 UTC dc66e01a7040fdfa8287fbfbbf73f9ac46a1b679

3 files changed +89

tests/v1/core/test_prefix_caching.py

+62

@@ -5,6 +5,7 @@

5 import copy

6 from collections.abc import Callable

7 from math import lcm

8 +from types import SimpleNamespace

10 import pytest

11 import torch

@@ -33,6 +34,7 @@ from vllm.v1.core.kv_cache_utils import (

34 init_none_hash,

35 make_block_hash_with_group_id,

36 )

37 +from vllm.v1.core.sched.scheduler import Scheduler

38 from vllm.v1.kv_cache_interface import (

39 FullAttentionSpec,

40 KVCacheConfig,

@@ -1023,6 +1025,66 @@ def test_prefill_hybrid_model_mamba_align():

1025 manager.free(req0)

1026

1027

1028 +def test_hybrid_cache_mamba_align_shared_prefix_detection():

1029 + """Test shared prefix detection heuristic for mamba align cache mode

1030 +

1031 + HybridKVCacheCoordinator returns num_uncached_common > 0 when a shared

1032 + uncached prefix is detected. With mamba_align cache, _mamba_block_aligned_split

1033 + enforces scheduling aligned with the common prefix.

1034 + """

1035 + block_size = 16

1036 + manager = make_kv_cache_manager(

1037 + _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba_align"]),

1038 + max_model_len=8192,

1039 + enable_caching=True,

1040 + hash_block_size=block_size,

1041 + )

1042 + hash_fn = sha256

1043 +

1044 + # Request: 3 blocks

1045 + prefix = [i for i in range(3) for _ in range(block_size)]

1046 + req_0 = make_request("0", prefix, block_size, hash_fn)

1047 + computed_blocks, num_computed = manager.get_computed_blocks(req_0)

1048 + num_uncached_common = manager.coordinator.num_uncached_common_prefix_tokens

1049 + assert num_computed == 0 # nothing cached yet

1050 + assert num_uncached_common == 0

1051 + manager.allocate_slots(req_0, 3 * block_size, 0, computed_blocks)

1052 +

1053 + # Request: 3 blocks (shared with above) + 7 different tokens

1054 + req_1 = make_request("1", prefix + [100] * 7, block_size, hash_fn)

1055 + computed_blocks, num_computed = manager.get_computed_blocks(req_1)

1056 + num_uncached_common = manager.coordinator.num_uncached_common_prefix_tokens

1057 + assert num_computed == 3 * block_size # we should observe a 3-block cache hit

1058 + assert num_uncached_common == 0

1059 + manager.allocate_slots(req_1, 7, 3 * block_size, computed_blocks)

1060 +

1061 + # Request: 3 blocks, but only 2 blocks shared (replace the last token in 3rd block):

1062 + req_2 = make_request("2", prefix[:-1] + [101], block_size, hash_fn)

1063 + computed_blocks, num_computed = manager.get_computed_blocks(req_2)

1064 + num_uncached_common = manager.coordinator.num_uncached_common_prefix_tokens

1065 + assert num_computed == 0 # mamba_align doesn't cache intermediate blocks

1066 + assert num_uncached_common == 2 * block_size # heuristic detects a shared prefix

1067 +

1068 + # Next, validate scheduler logic for num_uncached_common_prefix_tokens > 0

1069 + # Create minimal mock with just the needed attributes

1070 + mock = SimpleNamespace(

1071 + cache_config=SimpleNamespace(block_size=block_size), use_eagle=False

1072 + )

1073 + num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(

1074 + self=mock,

1075 + request=req_2,

1076 + num_new_tokens=3 * block_size,

1077 + num_uncached_common_prefix_tokens=num_uncached_common,

1078 + )

1079 + assert num_new_tokens_adjusted == 2 * block_size # adjust to the common prefix

1080 +

1081 + manager.allocate_slots(req_2, 3 * block_size, 0, computed_blocks)

1082 + # Cleanup

1083 + manager.free(req_0)

1084 + manager.free(req_1)

1085 + manager.free(req_2)

1086 +

1087 +

1088 def test_hybrid_model_mamba_align_with_dynamic_draft_tokens():

1089 """Regression test for https://github.com/vllm-project/vllm/issues/39271.

1090

vllm/v1/core/kv_cache_coordinator.py

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

611

612 num_groups = len(self.kv_cache_config.kv_cache_groups)

613 hit_length = max_cache_hit_length

614 + longest_hit_length = 0

615 hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups

616

617 # Simple hybrid (1 full attn + 1 other): one iteration suffices.

@@ -667,6 +668,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

668 for group_id, blocks in zip(group_ids, hit_blocks):

669 hit_blocks_by_group[group_id] = blocks

670

671 + longest_hit_length = max(longest_hit_length, curr_hit_length)

672 +

673 if curr_hit_length >= hit_length:

674 break

675 hit_length = curr_hit_length

@@ -681,6 +684,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

684 if (blks := hit_blocks_by_group[group_id]) is not None:

685 del blks[num_blocks:]

686

687 + # Uncached shared prefix detection: If any attn. group cached a longer prefix

688 + # than the current prefix, it is an uncached common prefix across requests:

689 + self.num_uncached_common_prefix_tokens = longest_hit_length - hit_length

690 return tuple(

691 blocks if blocks is not None else [] for blocks in hit_blocks_by_group

692 ), hit_length

vllm/v1/core/sched/scheduler.py

+21

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

296 num_new_tokens: int,

297 num_new_local_computed_tokens: int = 0,

298 num_external_computed_tokens: int = 0,

299 + num_uncached_common_prefix_tokens: int = 0,

300 ) -> int:

301 num_computed_tokens = (

302 request.num_computed_tokens

@@ -335,6 +336,16 @@ class Scheduler(SchedulerInterface):

336 else:

337 # prefill the last few tokens

338 pass

339 +

340 + # Marconi cache admission optimization:

341 + # cache common prefixes by scheduling num_new_tokens = common prefix length

342 + if (

343 + num_uncached_common_prefix_tokens >= block_size

344 + and num_new_tokens > num_uncached_common_prefix_tokens

345 + ):

346 + num_new_tokens = num_uncached_common_prefix_tokens

347 + # keep alignment to block_size

348 + num_new_tokens = num_new_tokens // block_size * block_size

349 return num_new_tokens

350

351 def schedule(self) -> SchedulerOutput:

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

615 num_external_computed_tokens = 0

616 load_kv_async = False

617 connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0

618 + num_uncached_common_prefix_tokens = 0

619

620 # Get already-cached tokens.

621 if request.num_computed_tokens == 0:

@@ -612,6 +624,14 @@ class Scheduler(SchedulerInterface):

624 self.kv_cache_manager.get_computed_blocks(request)

625 )

626

627 + # In case of hybrid models, obtain hint for Marconi-style APC logic

628 + if self.has_mamba_layers:

629 + num_uncached_common_prefix_tokens = getattr(

630 + self.kv_cache_manager.coordinator,

631 + "num_uncached_common_prefix_tokens",

632 + 0,

633 + )

634 +

635 # Get externally-cached tokens if using a KVConnector.

636 if self.connector is not None:

637 ext_tokens, load_kv_async = (

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

744 num_new_tokens,

745 num_new_local_computed_tokens,

746 num_external_computed_tokens,

747 + num_uncached_common_prefix_tokens,

748 )

749 if num_new_tokens == 0:

750 break