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 15, 13:30 UTC
[V1][Spec Decode] Make eagle compatible with prefix caching. (#17137)
Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>
Lily Liu committed
Apr 27, 2025 at 09:29 UTC
20e489eaa1d1de7ca05297e6d46797432e0f7c3d
4 files changed
+81
-9
tests/v1/core/test_prefix_caching.py
+57
@@ -719,3 +719,60 @@ def test_prefix_cache_stats_disabled():
719
720
# Ensure prefix_cache_stats remains None
721
assert manager.prefix_cache_stats is None
722
+
723
+
724
+def test_eagle_enabled_removes_last_block():
725
+ """Verify Eagle does NOT remove blocks when request
726
+ length is divisible by block size."""
727
+ block_size = 16
728
+ manager = KVCacheManager(
729
+ make_kv_cache_config(block_size, num_blocks=10),
730
+ max_model_len=8192,
731
+ enable_caching=True,
732
+ use_eagle=True,
733
+ )
734
+
735
+ # Request with 3 full blocks (48 tokens)
736
+ token_ids = [0] * (3 * block_size)
737
+ req = make_request("divisible_request", token_ids)
738
+
739
+ # Prime the cache
740
+ computed_blocks, _ = manager.get_computed_blocks(req)
741
+ manager.allocate_slots(req, len(token_ids), computed_blocks)
742
+ manager.free(req)
743
+
744
+ # New request with same tokens + Eagle enabled
745
+ req_eagle = make_request("eagle_divisible", token_ids)
746
+ computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
747
+
748
+ # Should retain 2 blocks:
749
+ # 1. Original 3 blocks → pop last hash → 2 matched blocks
750
+ # 2. last_block_hash is not None → Eagle pop is not SKIPPED
751
+ assert len(computed_blocks) == 1
752
+ assert num_tokens == 1 * block_size # 32 tokens
753
+
754
+
755
+def test_eagle_with_partial_blocks():
756
+ """Test Eagle behavior with requests containing partial blocks."""
757
+ block_size = 16
758
+ manager = KVCacheManager(
759
+ make_kv_cache_config(block_size, num_blocks=10),
760
+ max_model_len=8192,
761
+ enable_caching=True,
762
+ use_eagle=True,
763
+ )
764
+ # 2 full blocks + 5 tokens (non-divisible length)
765
+ token_ids = [0] * (2 * block_size + 5)
766
+ req = make_request("partial_block_test", token_ids)
767
+
768
+ # Prime the cache
769
+ computed_blocks, _ = manager.get_computed_blocks(req)
770
+ manager.allocate_slots(req, len(token_ids), computed_blocks)
771
+ manager.free(req)
772
+
773
+ # New request with Eagle enabled
774
+ req_eagle = make_request("partial_eagle", token_ids)
775
+ computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
776
+ # Original match: 2 full blocks → Eagle removes 1 → 1 remaining
777
+ assert len(computed_blocks) == 1
778
+ assert num_tokens == 1 * block_size
tests/v1/e2e/test_spec_decode.py
-1
@@ -44,7 +44,6 @@ def test_prompts():
44
45
@pytest.fixture
46
def sampling_config():
47
- # Only support greedy for now
47
return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False)
48
49
vllm/v1/core/kv_cache_manager.py
+10
@@ -25,6 +25,7 @@ class KVCacheManager:
25
max_model_len: int,
26
enable_caching: bool = True,
27
caching_hash_algo: str = "builtin",
28
+ use_eagle: bool = False,
29
log_stats: bool = False,
30
) -> None:
31
assert len(kv_cache_config.kv_cache_groups) == 1, (
@@ -38,6 +39,7 @@ class KVCacheManager:
39
40
self.enable_caching = enable_caching
41
self.caching_hash_fn = sha256 if caching_hash_algo == "sha256" else hash
42
+ self.use_eagle = use_eagle
43
self.log_stats = log_stats
44
# FIXME: make prefix cache stats conditional on log_stats
45
self.prefix_cache_stats = PrefixCacheStats() if log_stats else None
@@ -134,6 +136,14 @@ class KVCacheManager:
136
137
computed_blocks = (
138
self.specialized_manager.find_longest_cache_hit(block_hashes))
139
+
140
+ if self.use_eagle and len(computed_blocks) > 0:
141
+ # Drop the last matched block if (1) eagle is enabled and
142
+ # (2) there is a cache hit.
143
+ # This is to recompute the last block to get the required
144
+ # hidden states for eagle drafting head.
145
+ computed_blocks.pop()
146
+
147
if self.log_stats:
148
assert self.prefix_cache_stats is not None
149
self.prefix_cache_stats.queries += len(block_hashes)
vllm/v1/core/sched/scheduler.py
+14
-8
@@ -74,13 +74,6 @@ class Scheduler(SchedulerInterface):
74
num_gpu_blocks = self.cache_config.num_gpu_blocks
75
assert num_gpu_blocks is not None and num_gpu_blocks > 0
76
77
- # Create the KV cache manager.
78
- self.kv_cache_manager = KVCacheManager(
79
- kv_cache_config=kv_cache_config,
80
- max_model_len=self.max_model_len,
81
- enable_caching=self.cache_config.enable_prefix_caching,
82
- caching_hash_algo=self.cache_config.prefix_caching_hash_algo,
83
- log_stats=self.log_stats)
77
self.block_size = self.cache_config.block_size
78
79
# req_id -> Request
@@ -123,12 +116,24 @@ class Scheduler(SchedulerInterface):
116
cache_size=encoder_cache_size)
117
118
speculative_config = vllm_config.speculative_config
119
+
120
+ self.use_eagle = False
121
self.num_spec_tokens = self.num_lookahead_tokens = 0
122
if speculative_config:
123
self.num_spec_tokens = speculative_config.num_speculative_tokens
124
if speculative_config.use_eagle():
125
+ self.use_eagle = True
126
self.num_lookahead_tokens = self.num_spec_tokens
127
128
+ # Create the KV cache manager.
129
+ self.kv_cache_manager = KVCacheManager(
130
+ kv_cache_config=kv_cache_config,
131
+ max_model_len=self.max_model_len,
132
+ enable_caching=self.cache_config.enable_prefix_caching,
133
+ caching_hash_algo=self.cache_config.prefix_caching_hash_algo,
134
+ use_eagle=self.use_eagle,
135
+ log_stats=self.log_stats)
136
+
137
def schedule(self) -> SchedulerOutput:
138
# NOTE(woosuk) on the scheduling algorithm:
139
# There's no "decoding phase" nor "prefill phase" in the scheduler.
@@ -317,7 +322,8 @@ class Scheduler(SchedulerInterface):
322
323
# Get already-cached tokens.
324
computed_blocks, num_computed_tokens = \
320
- self.kv_cache_manager.get_computed_blocks(request)
325
+ self.kv_cache_manager.get_computed_blocks(
326
+ request)
327
328
# Get externally-cached tokens if using a KVConnector.
329
num_external_tokens = (