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 8, 15:15 UTC
[BugFix] Avoid prefix cache hit in the same schedule step for mamba layers (#29387)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Feb 9, 2026 at 23:41 UTC
97fa8f65909d4d8f2eb0edc2137fb22f576a5b25
6 files changed
+78
-3
tests/models/language/generation/test_hybrid.py
+28
@@ -7,6 +7,7 @@ import pytest
7
8
from tests.models.registry import HF_EXAMPLE_MODELS
9
from tests.utils import multi_gpu_test
10
+from vllm import LLM
11
from vllm.engine.arg_utils import EngineArgs
12
from vllm.platforms import current_platform
13
from vllm.sampling_params import SamplingParams
@@ -769,3 +770,30 @@ def test_apc_multiple_prompts_partial_cached_outputs(
770
name_0="vllm_no_cache",
771
name_1=f"vllm_cache_it_{r_idx + 1}",
772
)
773
+
774
+
775
+# we have to use a real large model to get reasonable results
776
+# the model can't be a hybrid model as we need block_size 16
777
+@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
778
+def test_apc_common_prefix_same_batch(
779
+ model: str,
780
+ monkeypatch,
781
+) -> None:
782
+ # Required to put the two requests in the same batch
783
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
784
+ llm = LLM(
785
+ model=model,
786
+ enforce_eager=True,
787
+ block_size=16,
788
+ mamba_block_size=16,
789
+ enable_prefix_caching=True,
790
+ seed=42,
791
+ )
792
+ prompts = [
793
+ "hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
794
+ "hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
795
+ ]
796
+ sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
797
+ outputs = llm.generate(prompts, sampling_params)
798
+ for output in outputs:
799
+ assert "two" in output.outputs[0].text
tests/v1/core/test_prefix_caching.py
+2
@@ -857,6 +857,8 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
857
# Should have blocks for all groups
858
assert len(blocks.get_block_ids()) == num_groups
859
860
+ manager.new_step_starts()
861
+
862
# Second request: should hit cached blocks for common prefix
863
req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn)
864
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
vllm/v1/core/kv_cache_coordinator.py
+5
@@ -247,6 +247,11 @@ class KVCacheCoordinator(ABC):
247
) -> tuple[tuple[list[KVCacheBlock], ...], int]:
248
pass
249
250
+ def new_step_starts(self) -> None:
251
+ """Called when a new step is started."""
252
+ for manager in self.single_type_managers:
253
+ manager.new_step_starts()
254
+
255
256
class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
257
"""
vllm/v1/core/kv_cache_manager.py
+4
@@ -488,3 +488,7 @@ class KVCacheManager:
488
) -> KVCacheBlocks:
489
# Only create new KVCacheBlocks for non-empty blocks
490
return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
491
+
492
+ def new_step_starts(self) -> None:
493
+ """Called when a new step is started."""
494
+ self.coordinator.new_step_starts()
vllm/v1/core/sched/scheduler.py
+2
@@ -347,6 +347,8 @@ class Scheduler(SchedulerInterface):
347
# For logging.
348
scheduled_timestamp = time.monotonic()
349
350
+ self.kv_cache_manager.new_step_starts()
351
+
352
# First, schedule the RUNNING requests.
353
req_index = 0
354
while req_index < len(self.running) and token_budget > 0:
vllm/v1/core/single_type_kv_cache_manager.py
+37
-3
@@ -7,7 +7,11 @@ from collections.abc import Sequence
7
8
from vllm.utils.math_utils import cdiv
9
from vllm.v1.core.block_pool import BlockPool
10
-from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock
10
+from vllm.v1.core.kv_cache_utils import (
11
+ BlockHashList,
12
+ BlockHashWithGroupId,
13
+ KVCacheBlock,
14
+)
15
from vllm.v1.kv_cache_interface import (
16
ChunkedLocalAttentionSpec,
17
CrossAttentionSpec,
@@ -396,6 +400,10 @@ class SingleTypeKVCacheManager(ABC):
400
# The default behavior is to not skip any tokens.
401
return 0
402
403
+ def new_step_starts(self) -> None:
404
+ # do nothing by default
405
+ return None
406
+
407
408
class FullAttentionManager(SingleTypeKVCacheManager):
409
@classmethod
@@ -742,8 +750,11 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
750
751
752
class MambaManager(SingleTypeKVCacheManager):
745
- def __init__(self, kv_cache_spec: MambaSpec, **kwargs) -> None:
746
- super().__init__(kv_cache_spec, **kwargs)
753
+ def __init__(
754
+ self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs
755
+ ) -> None:
756
+ super().__init__(kv_cache_spec, block_pool, **kwargs)
757
+ self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()
758
self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
759
self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks
760
if self.mamba_cache_mode == "align":
@@ -838,6 +849,15 @@ class MambaManager(SingleTypeKVCacheManager):
849
num_tokens_main_model: int,
850
) -> int:
851
assert isinstance(self.kv_cache_spec, MambaSpec)
852
+ if (
853
+ len(new_computed_blocks) > 0
854
+ and new_computed_blocks[-1].block_hash in self.cached_blocks_this_step
855
+ ):
856
+ # Mamba can't rely on blocks generated by other requests in the current step
857
+ # To put it in the next step, we return num_gpu_blocks + 1 so
858
+ # that kv_cache_manager will think there is no enough blocks to allocte now
859
+ # and don't schedule it in the current step.
860
+ return self.block_pool.num_gpu_blocks + 1
861
if self.mamba_cache_mode != "align":
862
# Allocate extra `num_speculative_blocks` blocks for
863
# speculative decoding (MTP/EAGLE) with linear attention.
@@ -972,6 +992,20 @@ class MambaManager(SingleTypeKVCacheManager):
992
"""
993
return num_computed_tokens - 1
994
995
+ def cache_blocks(self, request: Request, num_tokens: int) -> None:
996
+ num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)
997
+ super().cache_blocks(request, num_tokens)
998
+ num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
999
+ if num_cached_blocks_after > num_cached_blocks_before:
1000
+ for block in self.req_to_blocks[request.request_id][
1001
+ num_cached_blocks_before:num_cached_blocks_after
1002
+ ]:
1003
+ assert block.block_hash is not None
1004
+ self.cached_blocks_this_step.add(block.block_hash)
1005
+
1006
+ def new_step_starts(self) -> None:
1007
+ self.cached_blocks_this_step.clear()
1008
+
1009
1010
class CrossAttentionManager(SingleTypeKVCacheManager):
1011
"""Manager for cross-attention KV cache in encoder-decoder models."""