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, 14:45 UTC
[V1][Spec Decode] KV cache slots for eagle heads (#16370)
Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>
Lily Liu committed
Apr 12, 2025 at 19:42 UTC
f49e5aff11c986ed4d45202b1716c5d74786efa9
4 files changed
+98
-18
tests/v1/core/test_kv_cache_utils.py
+74
-12
@@ -7,6 +7,7 @@ from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
7
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
8
from vllm.sampling_params import SamplingParams
9
from vllm.utils import GiB_bytes, sha256
10
+from vllm.v1.core.kv_cache_manager import KVCacheManager
11
# disable yapf here as it formats differently than isort such that both fail
12
# yapf: disable
13
from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType,
@@ -48,6 +49,18 @@ def make_request(request_id,
49
)
50
51
52
+def new_kv_cache_spec(block_size=16,
53
+ num_kv_heads=2,
54
+ head_size=64,
55
+ dtype=torch.float32,
56
+ use_mla=False):
57
+ return FullAttentionSpec(block_size=block_size,
58
+ num_kv_heads=num_kv_heads,
59
+ head_size=head_size,
60
+ dtype=dtype,
61
+ use_mla=use_mla)
62
+
63
+
64
def test_none_hash():
65
assert NONE_HASH is not None
66
assert isinstance(NONE_HASH, int)
@@ -327,18 +340,6 @@ def test_metrics():
340
341
342
def test_unify_kv_cache_configs():
330
-
331
- def new_kv_cache_spec(block_size=16,
332
- num_kv_heads=2,
333
- head_size=64,
334
- dtype=torch.float32,
335
- use_mla=False):
336
- return FullAttentionSpec(block_size=block_size,
337
- num_kv_heads=num_kv_heads,
338
- head_size=head_size,
339
- dtype=dtype,
340
- use_mla=use_mla)
341
-
343
same_kv_cache_config = [
344
KVCacheConfig(
345
num_blocks=10,
@@ -470,3 +471,64 @@ def test_estimate_max_model_len(model_id, max_model_len,
471
estimated_max_len = estimate_max_model_len(vllm_config, kv_cache_spec,
472
8 * GiB_bytes)
473
assert estimated_max_len == want_estimated_max_len
474
+
475
+
476
+def test_allocate_with_lookahead():
477
+ """Verify that lookahead tokens correctly affect block allocation"""
478
+ block_size = 4
479
+ config = KVCacheConfig(
480
+ num_blocks=10,
481
+ tensors={
482
+ "layer1": KVCacheTensor(100),
483
+ },
484
+ kv_cache_groups=[
485
+ KVCacheGroupSpec(["layer1"],
486
+ new_kv_cache_spec(block_size=block_size)),
487
+ ],
488
+ )
489
+
490
+ request = make_request(
491
+ request_id=0,
492
+ prompt_token_ids=[],
493
+ mm_positions=None,
494
+ mm_hashes=None,
495
+ )
496
+
497
+ # Test case 1: Requires additional lookahead tokens
498
+ kv_cache_manager = KVCacheManager(kv_cache_config=config,
499
+ max_model_len=100,
500
+ num_preallocate_tokens=0)
501
+ blocks = kv_cache_manager.allocate_slots(
502
+ request,
503
+ num_tokens=3,
504
+ num_lookahead_tokens=2, # Total required: 3+2=5 tokens
505
+ )
506
+ assert len(blocks) == 2 # ceil(5/4)=2 blocks
507
+
508
+ # Test case 2: With precomputed blocks
509
+ kv_cache_manager = KVCacheManager(kv_cache_config=config,
510
+ max_model_len=100,
511
+ num_preallocate_tokens=4)
512
+ # num_preallocate_blocks = 4 // 4 - 2 // 4 = 1
513
+ # required_blocks = ceil((3 + 2) /4) = 2
514
+ # total_blocks = 1 + 2 = 3
515
+ blocks = kv_cache_manager.allocate_slots(
516
+ request,
517
+ num_tokens=3,
518
+ num_lookahead_tokens=2,
519
+ )
520
+ assert len(blocks) == 3
521
+
522
+ # Test case 3: With precomputed blocks
523
+ # num_preallocate_blocks = 4 // 4 - 4 // 4 = 0
524
+ # required_blocks = ceil((3 + 4) / 4) = 2
525
+ # total_blocks = 0 + 2 = 2
526
+ kv_cache_manager = KVCacheManager(kv_cache_config=config,
527
+ max_model_len=100,
528
+ num_preallocate_tokens=4)
529
+ blocks = kv_cache_manager.allocate_slots(
530
+ request,
531
+ num_tokens=3,
532
+ num_lookahead_tokens=4,
533
+ )
534
+ assert len(blocks) == 2
vllm/v1/core/kv_cache_manager.py
+12
-4
@@ -164,7 +164,8 @@ class KVCacheManager:
164
self,
165
request: Request,
166
num_tokens: int,
167
- new_computed_blocks: Optional[list[KVCacheBlock]] = None
167
+ new_computed_blocks: Optional[list[KVCacheBlock]] = None,
168
+ num_lookahead_tokens: int = 0,
169
) -> Optional[list[KVCacheBlock]]:
170
"""Add slots for a request with new tokens to append.
171
@@ -174,6 +175,9 @@ class KVCacheManager:
175
not include the tokens that have already been computed.
176
new_computed_blocks: A list of new computed blocks just hitting the
177
prefix caching.
178
+ num_lookahead_tokens: The number of speculative tokens to allocate.
179
+ This is used by spec decode proposers with kv-cache such
180
+ as eagle.
181
182
Blocks layout:
183
-----------------------------------------------------------------------
@@ -211,8 +215,9 @@ class KVCacheManager:
215
# the new prefix caching hits
216
num_computed_tokens = (request.num_computed_tokens +
217
len(new_computed_blocks) * self.block_size)
214
- num_required_blocks = cdiv(num_computed_tokens + num_tokens,
215
- self.block_size)
218
+ num_required_blocks = cdiv(
219
+ num_computed_tokens + num_tokens + num_lookahead_tokens,
220
+ self.block_size)
221
num_new_blocks = (num_required_blocks - len(req_blocks) -
222
len(new_computed_blocks))
223
@@ -246,8 +251,11 @@ class KVCacheManager:
251
else:
252
# Get new blocks from the free block pool considering
253
# preallocated blocks.
254
+ num_preallocate_blocks = max(
255
+ 0, self.num_preallocate_blocks -
256
+ num_lookahead_tokens // self.block_size)
257
num_new_blocks = min(
250
- num_new_blocks + self.num_preallocate_blocks,
258
+ num_new_blocks + num_preallocate_blocks,
259
self.block_pool.get_num_free_blocks(),
260
# Should not exceed the maximum number of blocks per request.
261
# This is especially because the block table has the shape
vllm/v1/core/sched/scheduler.py
+11
-2
@@ -7,7 +7,8 @@ from collections import deque
7
from collections.abc import Iterable
8
from typing import Optional, Union
9
10
-from vllm.config import CacheConfig, LoRAConfig, ModelConfig, SchedulerConfig
10
+from vllm.config import (CacheConfig, LoRAConfig, ModelConfig, SchedulerConfig,
11
+ SpeculativeConfig)
12
from vllm.logger import init_logger
13
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
14
from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
@@ -39,6 +40,7 @@ class Scheduler(SchedulerInterface):
40
lora_config: Optional[LoRAConfig],
41
kv_cache_config: KVCacheConfig,
42
structured_output_manager: StructuredOutputManager,
43
+ speculative_config: SpeculativeConfig = None,
44
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
45
include_finished_set: bool = False,
46
log_stats: bool = False,
@@ -112,6 +114,11 @@ class Scheduler(SchedulerInterface):
114
self.encoder_cache_manager = EncoderCacheManager(
115
cache_size=encoder_cache_size)
116
117
+ self.num_lookahead_tokens = 0
118
+ if speculative_config and speculative_config.method == "eagle":
119
+ self.num_lookahead_tokens = \
120
+ speculative_config.num_speculative_tokens
121
+
122
def schedule(self) -> SchedulerOutput:
123
# NOTE(woosuk) on the scheduling algorithm:
124
# There's no "decoding phase" nor "prefill phase" in the scheduler.
@@ -188,7 +195,9 @@ class Scheduler(SchedulerInterface):
195
196
while True:
197
new_blocks = self.kv_cache_manager.allocate_slots(
191
- request, num_new_tokens)
198
+ request,
199
+ num_new_tokens,
200
+ num_lookahead_tokens=self.num_lookahead_tokens)
201
if new_blocks is None:
202
# The request cannot be scheduled.
203
# Preempt the lowest-priority request.
vllm/v1/engine/core.py
+1
@@ -98,6 +98,7 @@ class EngineCore:
98
cache_config=vllm_config.cache_config,
99
lora_config=vllm_config.lora_config,
100
kv_cache_config=kv_cache_config,
101
+ speculative_config=vllm_config.speculative_config,
102
structured_output_manager=self.structured_output_manager,
103
include_finished_set=vllm_config.parallel_config.data_parallel_size
104
> 1,