1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
| ```python
def allocate_slots(
self,
request: Request,
num_new_tokens: int,
num_new_computed_tokens: int = 0,
new_computed_blocks: KVCacheBlocks | None = None,
num_lookahead_tokens: int = 0,
num_external_computed_tokens: int = 0,
delay_cache_blocks: bool = False,
num_encoder_tokens: int = 0,
full_sequence_must_fit: bool = False,
) -> KVCacheBlocks | None:
# """Add slots for a request with new tokens to append.
# Args:
# request: The request to allocate slots.
# num_new_tokens: The number of new tokens to be allocated and computed.
# num_new_computed_tokens: The number of new computed tokens just
# hitting the prefix caching, excluding external tokens.
# new_computed_blocks: The cached blocks for the above new computed
# tokens, grouped as a tuple by kv cache groups.
# num_lookahead_tokens: The number of speculative tokens to allocate.
# This is used by spec decode proposers with kv-cache such
# as eagle.
# num_external_computed_tokens: The number of tokens that their
# KV caches are not cached by vLLM but cached by the connector.
# delay_cache_blocks: Whether to skip caching the blocks. This is
# used by P/D when allocating blocks used in a KV transfer
# which will complete in a future step.
# num_encoder_tokens: The number of encoder tokens to allocate for
# cross-attention in encoder-decoder models(e.g., Whisper).
# For decoder-only models, this should be 0.
# full_sequence_must_fit: Only allocate blocks if the KV cache has enough
# free blocks to hold the full sequence, accounting for prefix cache hits
# and sliding window. Used as an admission gate to prevent over-admitting
# requests when chunked prefill would otherwise only check the first chunk
# Blocks layout:
# ```
# ----------------------------------------------------------------------
# | < comp > | < new_comp > | < ext_comp > | < new > | < lookahead > |
# ----------------------------------------------------------------------
# | < to be computed > |
# ----------------------------------------------------------------------
# | < to be allocated > |
# ----------------------------------------------------------------------
# | < to be cached (roughly, |
# | details below)> |
# ----------------------------------------------------------------------
# | Prefix-cached tokens from either vLLM |
# | or connector. Can be safely removed if |
# | they are outside sliding window. |
# ----------------------------------------------------------------------
# | < cached by vLLM > | not cached by |
# | vLLM, but |
# | ref_cnt | ref_cnt not | cached by |
# | increased| increased yet| connector |
# ----------------------------------------------------------------------
# ```
# Abbrivations:
# ```
# comp = request.num_computed_tokens
# new_comp = num_new_computed_tokens
# = len(new_computed_blocks) * block_size
# ext_comp = num_external_computed_tokens, cached by the connector
# new = num_new_tokens, including unverified draft tokens
# lookahead = num_lookahead_tokens
# ```
# NOTE: for new tokens which include both verified and unverified draft
# tokens, we only cache the verified tokens (by capping the number at
# `request.num_tokens`).
# The allocation has three stages:
# - Free unnecessary blocks in `comp` and check
# if we have sufficient free blocks (return None if not).
# - Handle prefix tokens (`comp + new_comp + ext_comp`):
# - Free unnecessary blocks (e.g. outside sliding window)
# - Allocate new blocks for `ext_comp` tokens inside
# sliding window
# - Allocate new blocks for tokens to be computed (`new + lookahead`)
# Returns:
# A list of new allocated blocks.
# """
# # When loading KV data asynchronously, we may have zero new tokens to
# # compute while still allocating slots for externally computed tokens.
# if num_new_tokens == 0 and num_external_computed_tokens == 0:
# raise ValueError(
# "num_new_tokens must be greater than 0 when there are no "
# "external computed tokens"
# )
# if new_computed_blocks is not None:
# new_computed_block_list = new_computed_blocks.blocks
# else:
# new_computed_block_list = self.empty_kv_cache_blocks.blocks
# # The number of computed tokens is the number of computed tokens plus
# # the new prefix caching hits
# num_local_computed_tokens = (
# request.num_computed_tokens + num_new_computed_tokens
# )
# total_computed_tokens = min(
# num_local_computed_tokens + num_external_computed_tokens,
# self.max_model_len,
# )
# if full_sequence_must_fit:
# # First check and fail if the full request sequence won't fit.
# full_num_tokens = min(request.num_tokens, self.max_model_len)
# num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
# request_id=request.request_id,
# num_tokens=full_num_tokens,
# new_computed_blocks=new_computed_block_list,
# num_encoder_tokens=num_encoder_tokens,
# total_computed_tokens=total_computed_tokens,
# num_tokens_main_model=full_num_tokens,
# apply_admission_cap=True,
# )
# if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
# return None
# num_tokens_main_model = total_computed_tokens + num_new_tokens
# num_tokens_need_slot = min(
# num_tokens_main_model + num_lookahead_tokens, self.max_model_len
# )
# # Free the blocks that are skipped during the attention computation
# # (e.g., tokens outside the sliding window).
# # We can do this even if we cannot schedule this request due to
# # insufficient free blocks.
# # Should call this function before allocating new blocks to reduce
# # the number of evicted blocks.
self.coordinator.remove_skipped_blocks(
request.request_id, total_computed_tokens
)
num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
request_id=request.request_id,
num_tokens=num_tokens_need_slot,
new_computed_blocks=new_computed_block_list,
num_encoder_tokens=num_encoder_tokens,
total_computed_tokens=num_local_computed_tokens
+ num_external_computed_tokens,
num_tokens_main_model=num_tokens_main_model,
)
if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
# Cannot allocate new blocks
return None
if (
new_computed_block_list is not self.empty_kv_cache_blocks.blocks
or num_external_computed_tokens > 0
):
# Append the new computed blocks to the request blocks until now to
# avoid the case where the new blocks cannot be allocated.
self.coordinator.allocate_new_computed_blocks(
request_id=request.request_id,
new_computed_blocks=new_computed_block_list,
num_local_computed_tokens=num_local_computed_tokens,
num_external_computed_tokens=num_external_computed_tokens,
)
new_blocks = self.coordinator.allocate_new_blocks(
request.request_id,
num_tokens_need_slot,
num_tokens_main_model,
num_encoder_tokens,
)
# P/D: delay caching blocks if we have to recv from
# remote. Update state for locally cached blocks.
if not self.enable_caching or delay_cache_blocks:
return self.create_kv_cache_blocks(new_blocks)
# NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
# + num_external_computed_tokens + num_new_tokens, but must exclude
# "non-committable" tokens (e.g., draft tokens that could be rejected).
# Therefore, we cap the number at `request.num_tokens`, ensuring only
# "finalized" tokens are cached.
num_tokens_to_cache = min(
total_computed_tokens + num_new_tokens,
request.num_tokens,
)
self.coordinator.cache_blocks(request, num_tokens_to_cache)
return self.create_kv_cache_blocks(new_blocks)
def free(self, request: Request) -> None:
"""Free the blocks allocated for the request.
We free the blocks in reverse order so that the tail blocks are evicted
first when caching is enabled.
Args:
request: The request to free the blocks.
"""
self.coordinator.free(request.request_id)
def remove_skipped_blocks(
self, request_id: str, total_computed_tokens: int
) -> None:
"""Remove the blocks that are no longer needed from `blocks` and replace
the removed blocks with null_block.
Args:
request_id: The request ID.
total_computed_tokens: The total number of computed tokens, including
local computed tokens and external computed tokens.
"""
self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens)
def evict_blocks(self, block_ids: set[int]) -> None:
"""evict blocks from the prefix cache by their block IDs.
Args:
block_ids: Set of block IDs to evict from cache.
"""
self.block_pool.evict_blocks(block_ids)
def reset_prefix_cache(self) -> bool:
"""Reset prefix cache. This function may be used in RLHF
flows to invalidate prefix caching after the weights are updated,
or used for resetting prefix caching status for benchmarking.
Returns:
bool: True if the prefix cache is successfully reset,
False otherwise.
"""
if not self.block_pool.reset_prefix_cache():
return False
if self.log_stats:
assert self.prefix_cache_stats is not None
self.prefix_cache_stats.reset = True
return True
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
"""Calculate the number of common prefix blocks for each kv cache group.
The function selects a running request and iterates through its blocks.
A block is considered a common prefix block if ALL requests with
allocated KV cache share it (i.e., ref_cnt equals the number of entries
in req_to_blocks).
NOTE(woosuk): The number of requests with allocated KV cache is **greater
than or equal to** the number of requests scheduled in the current step.
This is because having allocated KV cache only indicates that:
1. The request has not yet finished, and
2. The request holds its blocks unfreed.
While all scheduled requests must have allocated KV cache, the inverse
is not necessarily true. There may be requests with allocated KV cache
that are not scheduled in the current step.
This can result in an edge case where the number of common prefix blocks
is 0, even though all scheduled requests share a common prefix. This
occurs because there may be unscheduled requests that do not share the
common prefix. Currently, this case cannot be easily detected, so the
function returns 0 in such cases.
Args:
running_request_id: The request ID of any running request, used to
identify the common prefix blocks.
Returns:
list[int]: The number of common prefix blocks for each kv cache
group.
"""
return self.coordinator.get_num_common_prefix_blocks(running_request_id)
def take_events(self) -> list[KVCacheEvent]:
"""Take the KV cache events from the block pool.
Returns:
A list of KV cache events.
"""
return self.block_pool.take_events()
def get_blocks(self, request_id: str) -> KVCacheBlocks:
"""Get the blocks of a request."""
return self.create_kv_cache_blocks(self.coordinator.get_blocks(request_id))
def get_block_ids(self, request_id: str) -> tuple[list[int], ...]:
"""Get the block ids of a request."""
return self.get_blocks(request_id).get_block_ids()
def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
"""Cache the blocks for the request, if enabled.
Args:
request: The request to cache the blocks.
num_computed_tokens: The number of computed tokens, including tokens
that are already cached and tokens to be cached.
"""
if self.enable_caching:
self.coordinator.cache_blocks(request, num_computed_tokens)
def create_kv_cache_blocks(
self, blocks: tuple[list[KVCacheBlock], ...]
) -> KVCacheBlocks:
# Only create new KVCacheBlocks for non-empty blocks
return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
def take_new_block_ids(self) -> list[int]:
"""Drain and return new attention block IDs for zeroing."""
ids: list[int] = []
for mgr in self.coordinator.single_type_managers:
ids.extend(mgr.take_new_block_ids())
return ids
def new_step_starts(self) -> None:
"""Called when a new step is started."""
self.coordinator.new_step_starts()
|
