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, 12:15 UTC
[v1][Spec Decode] Make sliding window compatible with eagle prefix caching (#17398)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
May 1, 2025 at 02:25 UTC
81ecf425f0a645e6d6f1c339ba79ef37e95a9569
4 files changed
+96
-23
tests/v1/core/test_prefix_caching.py
+60
-4
@@ -15,7 +15,7 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
15
from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
16
hash_block_tokens)
17
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
18
- KVCacheGroupSpec)
18
+ KVCacheGroupSpec, SlidingWindowSpec)
19
20
21
def make_request(request_id,
@@ -863,11 +863,11 @@ def test_eagle_enabled_removes_last_block():
863
req_eagle = make_request("eagle_divisible", token_ids)
864
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
865
866
- # Should retain 2 blocks:
866
+ # Should retain 1 block:
867
# 1. Original 3 blocks → pop last hash → 2 matched blocks
868
- # 2. last_block_hash is not None → Eagle pop is not SKIPPED
868
+ # 2. drop last matched block → 1 remaining block
869
assert len(computed_blocks) == 1
870
- assert num_tokens == 1 * block_size # 32 tokens
870
+ assert num_tokens == 1 * block_size # 16 tokens
871
872
873
def test_eagle_with_partial_blocks():
@@ -894,3 +894,59 @@ def test_eagle_with_partial_blocks():
894
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
895
assert len(computed_blocks) == 1
896
assert num_tokens == 1 * block_size
897
+
898
+
899
+def test_eagle_with_sliding_window():
900
+ """Test Eagle behavior with sliding window."""
901
+ block_size = 16
902
+ sliding_window_spec = SlidingWindowSpec(
903
+ block_size=block_size,
904
+ num_kv_heads=1,
905
+ head_size=1,
906
+ dtype=torch.float32,
907
+ sliding_window=block_size,
908
+ use_mla=False,
909
+ )
910
+ manager = KVCacheManager(
911
+ KVCacheConfig(
912
+ num_blocks=10,
913
+ tensors={},
914
+ kv_cache_groups=[KVCacheGroupSpec(['layer'], sliding_window_spec)],
915
+ ),
916
+ max_model_len=8192,
917
+ enable_caching=True,
918
+ use_eagle=True,
919
+ )
920
+
921
+ # 2 full blocks + 5 tokens (non-divisible length)
922
+ token_ids = [0] * (2 * block_size + 5)
923
+ req = make_request("partial_block_test", token_ids)
924
+
925
+ # Prime the cache
926
+ computed_blocks, _ = manager.get_computed_blocks(req)
927
+ manager.allocate_slots(req, len(token_ids), computed_blocks)
928
+ # record the block hash of the first block in the request for later use
929
+ block_hash_first_block = manager.req_to_block_hashes[req.request_id][0]
930
+ assert block_hash_first_block is not None
931
+ manager.free(req)
932
+
933
+ # New request with Eagle enabled
934
+ req_eagle = make_request("partial_eagle", token_ids)
935
+ computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
936
+ # Original match: 2 full blocks → Eagle removes 1 → 1 remaining
937
+ assert len(computed_blocks) == 1
938
+ assert num_tokens == 1 * block_size
939
+
940
+ # Evict the first block in the request
941
+ assert manager.block_pool.get_cached_block(
942
+ block_hash_first_block) is not None
943
+ manager.block_pool.cached_block_hash_to_block.pop(block_hash_first_block)
944
+
945
+ # New request
946
+ req_after_evict = make_request("partial_eagle_after_evict", token_ids)
947
+ computed_blocks, num_tokens = manager.get_computed_blocks(req_after_evict)
948
+ # Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is
949
+ # not considered. But after dropping the last matched block due to eagle,
950
+ # there will be no matched prefix.
951
+ assert len(computed_blocks) == 0
952
+ assert num_tokens == 0
tests/v1/core/test_specialized_manager.py
+6
-2
@@ -19,7 +19,9 @@ def test_sliding_window_possible_cached_prefix():
19
)
20
21
block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
22
- manager = SlidingWindowManager(sliding_window_spec, block_pool)
22
+ manager = SlidingWindowManager(sliding_window_spec,
23
+ block_pool,
24
+ use_eagle=False)
25
26
def run_one_case(block_is_cached, expect_length):
27
block_hash_list = [
@@ -79,7 +81,9 @@ def test_sliding_window_remove_skipped_blocks():
81
82
block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
83
82
- manager = SlidingWindowManager(sliding_window_spec, block_pool)
84
+ manager = SlidingWindowManager(sliding_window_spec,
85
+ block_pool,
86
+ use_eagle=False)
87
88
null_block_id = block_pool.null_block.block_id
89
vllm/v1/core/kv_cache_manager.py
+1
-7
@@ -52,6 +52,7 @@ class KVCacheManager:
52
self.specialized_manager = get_specialized_manager(
53
kv_cache_spec=kv_cache_spec,
54
block_pool=self.block_pool,
55
+ use_eagle=self.use_eagle,
56
)
57
58
# Mapping from request ID to blocks to track the blocks allocated
@@ -141,13 +142,6 @@ class KVCacheManager:
142
computed_blocks = (
143
self.specialized_manager.find_longest_cache_hit(block_hashes))
144
144
- if self.use_eagle and len(computed_blocks) > 0:
145
- # Drop the last matched block if (1) eagle is enabled and
146
- # (2) there is a cache hit.
147
- # This is to recompute the last block to get the required
148
- # hidden states for eagle drafting head.
149
- computed_blocks.pop()
150
-
145
if self.log_stats:
146
assert self.prefix_cache_stats is not None
147
self.prefix_cache_stats.queries += len(block_hashes)
vllm/v1/core/specialized_manager.py
+29
-10
@@ -18,6 +18,7 @@ class SpecializedManager(ABC):
18
self,
19
kv_cache_spec: KVCacheSpec,
20
block_pool: BlockPool,
21
+ use_eagle: bool,
22
) -> None:
23
"""
24
Initializes the SpecializedManager.
@@ -30,12 +31,17 @@ class SpecializedManager(ABC):
31
self.kv_cache_spec = kv_cache_spec
32
self.block_pool = block_pool
33
34
+ # Needs special handling for find_longest_cache_hit if eagle is enabled
35
+ self.use_eagle = use_eagle
36
+
37
@abstractmethod
38
def find_longest_cache_hit(
39
self, block_hashes: list[BlockHashType]) -> list[KVCacheBlock]:
40
"""
41
Get the longest cache hit prefix of the blocks. If no cache hit is
38
- found, return an empty list.
42
+ found, return an empty list. if eagle is enabled, drop the last matched
43
+ block to force recompute the last block to get the required hidden
44
+ states for eagle drafting head.
45
46
Args:
47
block_hashes: The block hashes of the request.
@@ -79,6 +85,8 @@ class FullAttentionManager(SpecializedManager):
85
computed_blocks.append(cached_block)
86
else:
87
break
88
+ if self.use_eagle and len(computed_blocks) > 0:
89
+ computed_blocks.pop()
90
return computed_blocks
91
92
def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
@@ -89,14 +97,20 @@ class FullAttentionManager(SpecializedManager):
97
98
class SlidingWindowManager(SpecializedManager):
99
92
- def __init__(self, kv_cache_spec: SlidingWindowSpec,
93
- block_pool: BlockPool):
94
- super().__init__(kv_cache_spec, block_pool)
100
+ def __init__(self, kv_cache_spec: SlidingWindowSpec, block_pool: BlockPool,
101
+ use_eagle: bool):
102
+ super().__init__(kv_cache_spec, block_pool, use_eagle)
103
self.sliding_window = kv_cache_spec.sliding_window
104
# The number of contiguous blocks needed for prefix cache hit.
105
# -1 since the input token itself is also included in the window
106
self.sliding_window_contiguous_blocks = cdiv(
107
(kv_cache_spec.sliding_window - 1), self.block_size)
108
+ if self.use_eagle:
109
+ # Need to drop the last matched block if eagle is enabled. For
110
+ # sliding window layer, we achieve this by increasing the number of
111
+ # contiguous blocks needed for prefix cache hit by one and dropping
112
+ # the last matched block.
113
+ self.sliding_window_contiguous_blocks += 1
114
self._null_block = block_pool.null_block
115
116
def find_longest_cache_hit(
@@ -109,6 +123,7 @@ class SlidingWindowManager(SpecializedManager):
123
computed_blocks = [self._null_block] * len(block_hashes)
124
num_contiguous_blocks = 0
125
126
+ match_found = False
127
# Search from right to left and early stop when a match is found.
128
for i in range(len(block_hashes) - 1, -1, -1):
129
if cached_block := self.block_pool.get_cached_block(
@@ -121,12 +136,16 @@ class SlidingWindowManager(SpecializedManager):
136
# E.g., [NULL, NULL, 8, 3, NULL, 9] -> [NULL, NULL, 8, 3]
137
# when sliding_window_contiguous_blocks=2.
138
del computed_blocks[i + num_contiguous_blocks:]
124
- return computed_blocks
139
+ match_found = True
140
+ break
141
else:
142
num_contiguous_blocks = 0
127
- # The first `num_contiguous_blocks` is a cache hit even if
128
- # `num_contiguous_blocks < sliding_window_contiguous_blocks`.
129
- del computed_blocks[num_contiguous_blocks:]
143
+ if not match_found:
144
+ # The first `num_contiguous_blocks` is a cache hit even if
145
+ # `num_contiguous_blocks < sliding_window_contiguous_blocks`.
146
+ del computed_blocks[num_contiguous_blocks:]
147
+ if self.use_eagle and len(computed_blocks) > 0:
148
+ computed_blocks.pop()
149
return computed_blocks
150
151
def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
@@ -155,7 +174,7 @@ spec_manager_map: dict[type[KVCacheSpec], type[SpecializedManager]] = {
174
175
176
def get_specialized_manager(kv_cache_spec: KVCacheSpec,
158
- block_pool: BlockPool) -> SpecializedManager:
177
+ **kwargs) -> SpecializedManager:
178
manager_class = spec_manager_map[type(kv_cache_spec)]
160
- manager = manager_class(kv_cache_spec, block_pool)
179
+ manager = manager_class(kv_cache_spec, **kwargs)
180
return manager