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 4, 16:15 UTC
[kv_offload+HMA][9/N]: Support lookup with multiple KV groups (#39401)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Apr 24, 2026 at 15:32 UTC
51adca74e6be951c86e920046a83bfc061193ba2
1 file changed
+72
-41
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+72
-41
@@ -120,6 +120,13 @@ class OffloadingConnectorScheduler:
120
self.config = SchedulerOffloadConfig.from_spec(spec)
121
self.manager: OffloadingManager = spec.get_manager()
122
123
+ attention_groups: list[int] = []
124
+ for idx, _ in enumerate(spec.kv_cache_config.kv_cache_groups):
125
+ # currently treat all groups as full attention
126
+ attention_groups.append(idx)
127
+
128
+ self.lookup_groups = attention_groups
129
+
130
self._req_status: dict[ReqId, RequestOffloadState] = {}
131
# requests to load for the current scheduler step
132
self._reqs_to_load: dict[ReqId, TransferSpec] = {}
@@ -204,64 +211,88 @@ class OffloadingConnectorScheduler:
211
group_state.block_ids.clear()
212
else:
213
req_status = RequestOffloadState(config=self.config, req=request)
207
- req_status.update_offload_keys()
214
self._req_status[request.request_id] = req_status
215
216
+ req_status.update_offload_keys()
217
req_status.num_locally_computed_tokens = num_computed_tokens
218
212
- # Below assertions will be removed once this function supports HMA
213
- assert len(self.config.kv_group_configs) == 1
214
- assert len(req_status.group_states) == 1
215
- group_config = self.config.kv_group_configs[0]
216
- group_state = req_status.group_states[0]
219
+ for gs in req_status.group_states:
220
+ self.manager.touch(gs.offload_keys)
221
218
- num_blocks = request.num_tokens // group_config.offloaded_block_size
222
+ # Start with the full request size as the maximum loadable
223
+ max_hit_size_tokens: int = req_status.req.num_tokens
224
+ num_hit_tokens: int = 0
225
+ defer_lookup = False
226
+ delay_request = False
227
+ for group_idx in self.lookup_groups:
228
+ group_config: GroupOffloadConfig = self.config.kv_group_configs[group_idx]
229
+ offloaded_block_size = group_config.offloaded_block_size
230
+ offload_keys = req_status.group_states[group_idx].offload_keys
231
220
- assert len(request.block_hashes) // self.config.block_size_factor == num_blocks
221
- offload_keys = group_state.offload_keys
232
+ num_blocks = max_hit_size_tokens // offloaded_block_size
233
+ assert len(offload_keys) >= num_blocks
234
223
- self.manager.touch(offload_keys)
235
+ # Constrain to block-aligned boundary for this group
236
+ max_hit_size_tokens = num_blocks * offloaded_block_size
237
+ num_hit_tokens = max_hit_size_tokens - num_computed_tokens
238
+ if num_hit_tokens < offloaded_block_size:
239
+ # we can only load less than a block, better skip
240
+ return 0, False
241
+
242
+ start_block_idx = num_computed_tokens // offloaded_block_size
243
+ offload_keys = offload_keys[start_block_idx:num_blocks]
244
+ # Full attention relies on all previous KV cache blocks.
245
+ # Thus, we search for a maximal prefix of KV cache which are all cached.
246
+ block_hits = self._maximal_prefix_lookup(
247
+ offload_keys, req_status.req_context
248
+ )
249
+ if block_hits == 0:
250
+ return 0, False
251
225
- full_block_tokens = group_config.offloaded_block_size * num_blocks
226
- if full_block_tokens - num_computed_tokens < group_config.offloaded_block_size:
227
- # we can load less than a block, skip
228
- return 0, False
252
+ if block_hits is None:
253
+ defer_lookup = True
254
+ else:
255
+ # Further constrain based on what's actually available by backend
256
+ max_hit_size_tokens = offloaded_block_size * (
257
+ start_block_idx + block_hits
258
+ )
259
230
- start_block_idx = num_computed_tokens // group_config.offloaded_block_size
231
- # Full attention relays on all previous KV cache blocks.
232
- # Thus, we search for a maximal prefix of KV cache which are all cached.
233
- hits = self._maximal_prefix_lookup(
234
- offload_keys[start_block_idx:], req_status.req_context
235
- )
236
- if hits is None:
237
- # indicates a lookup that should be tried later
260
+ num_hit_tokens = max_hit_size_tokens - num_computed_tokens
261
+ if num_hit_tokens < offloaded_block_size:
262
+ # we can only load less than a block, better skip
263
+ return 0, False
264
+
265
+ if (
266
+ block_hits
267
+ and self._blocks_being_loaded
268
+ and any(
269
+ key in self._blocks_being_loaded
270
+ for key in offload_keys[:block_hits]
271
+ )
272
+ ):
273
+ # hit blocks are being loaded, delay request
274
+ delay_request = True
275
+
276
+ if defer_lookup:
277
+ logger.debug(
278
+ "Offloading manager delayed request %s as backend requested",
279
+ req_status.req.request_id,
280
+ )
281
+ return None, False
282
+
283
+ if delay_request:
284
+ logger.debug(
285
+ "Delaying request %s since some of its blocks are already being loaded",
286
+ req_status.req.request_id,
287
+ )
288
return None, False
239
- if hits == 0:
240
- return 0, False
289
242
- num_hit_tokens = (
243
- group_config.offloaded_block_size * (start_block_idx + hits)
244
- - num_computed_tokens
245
- )
290
logger.debug(
291
"Request %s hit %s offloaded tokens after %s GPU hit tokens",
292
request.request_id,
293
num_hit_tokens,
294
num_computed_tokens,
295
)
252
- if num_hit_tokens < group_config.offloaded_block_size:
253
- return 0, False
254
-
255
- if self._blocks_being_loaded and any(
256
- key in self._blocks_being_loaded
257
- for key in offload_keys[start_block_idx : start_block_idx + hits]
258
- ):
259
- # hit blocks are being loaded, delay request
260
- logger.debug(
261
- "Delaying request %s since some of its blocks are already being loaded",
262
- request.request_id,
263
- )
264
- return None, False
296
297
return num_hit_tokens, True
298