941fb50835 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -18,6 +18,12 @@ from vllm.logger import init_logger

18 from vllm.utils.math_utils import cdiv

19 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

20 from vllm.v1.core.sched.output import SchedulerOutput

21 +from vllm.v1.kv_cache_interface import (

22 + FullAttentionSpec,

23 + KVCacheSpec,

24 + MambaSpec,

25 + SlidingWindowSpec,

26 +)

27 from vllm.v1.kv_offload.base import (

28 GPULoadStoreSpec,

29 OffloadingManager,

@@ -44,8 +50,12 @@ class TransferJobStatus:

50 # Offload keys this job covers; passed to manager.complete_*().

51 keys: set[OffloadKey]

52 is_store: bool

47 - # GPU blocks the fence tracks. Store src blocks; None for loads.

48 - gpu_block_ids: list[int] | None = None

53 + # Store src block IDs whose ref_cnt protects them while the request

54 + # runs. Only registered in _block_id_to_pending_jobs on request_finished.

55 + non_sliding_window_block_ids: list[int] | None = None

56 + # Store src block IDs that may be freed before the request finishes.

57 + # Registered in _block_id_to_pending_jobs at store creation time.

58 + sliding_window_block_ids: list[int] | None = None

59

60

61 class GroupOffloadConfig(NamedTuple):

@@ -53,6 +63,23 @@ class GroupOffloadConfig(NamedTuple):

63 gpu_block_size: int

64 offloaded_block_size: int

65 hash_block_size_factor: int

66 + # None below means full attention

67 + sliding_window_size_in_blocks: int | None

68 +

69 +

70 +def get_sliding_window_size_in_blocks(

71 + kv_cache_spec: KVCacheSpec, offloaded_block_size: int

72 +) -> int | None:

73 + if isinstance(kv_cache_spec, SlidingWindowSpec):

74 + assert kv_cache_spec.sliding_window > 0

75 + return cdiv(kv_cache_spec.sliding_window, offloaded_block_size)

76 +

77 + if isinstance(kv_cache_spec, MambaSpec):

78 + # Mamba depends on a single state

79 + return 1

80 +

81 + assert isinstance(kv_cache_spec, FullAttentionSpec)

82 + return None

83

84

85 class SchedulerOffloadConfig(NamedTuple):

@@ -73,6 +100,10 @@ class SchedulerOffloadConfig(NamedTuple):

100 (gpu_block_size * spec.block_size_factor)

101 // spec.hash_block_size

102 ),

103 + sliding_window_size_in_blocks=get_sliding_window_size_in_blocks(

104 + spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec,

105 + gpu_block_size * spec.block_size_factor,

106 + ),

107 )

108 for idx, gpu_block_size in enumerate(spec.gpu_block_size)

109 ),

@@ -86,6 +117,9 @@ class RequestGroupState:

117 block_ids: list[int] = field(default_factory=list)

118 # index of next block (of size offloaded_block_size) to offload

119 next_stored_block_idx: int = 0

120 + # number of offloaded blocks hit (including GPU prefix cache)

121 + # when the request first started

122 + num_hit_blocks: int = 0

123

124

125 @dataclass(slots=True)

@@ -139,6 +173,14 @@ class RequestOffloadState:

173 num_blocks = num_offloadable_tokens // group_config.offloaded_block_size

174 group_state.next_stored_block_idx = num_blocks

175

176 + def update_num_hit_blocks(self, num_cached_tokens: int) -> None:

177 + for group_config, group_state in zip(

178 + self.config.kv_group_configs, self.group_states

179 + ):

180 + group_state.num_hit_blocks = (

181 + num_cached_tokens // group_config.offloaded_block_size

182 + )

183 +

184

185 class OffloadingConnectorScheduler:

186 """Implementation of Scheduler side methods"""

@@ -147,12 +189,25 @@ class OffloadingConnectorScheduler:

189 self.config = SchedulerOffloadConfig.from_spec(spec)

190 self.manager: OffloadingManager = spec.get_manager()

191

150 - attention_groups: list[int] = []

151 - for idx, _ in enumerate(spec.kv_cache_config.kv_cache_groups):

152 - # currently treat all groups as full attention

153 - attention_groups.append(idx)

192 + full_attention_groups: list[int] = []

193 + sliding_window_groups: list[int] = []

194 + for group_config in self.config.kv_group_configs:

195 + if group_config.sliding_window_size_in_blocks is None:

196 + full_attention_groups.append(group_config.group_idx)

197 + else:

198 + sliding_window_groups.append(group_config.group_idx)

199 +

200 + # sort sliding window groups by window size in decreasing order

201 + def _sliding_window_sort_key(i: int) -> int:

202 + val = self.config.kv_group_configs[i].sliding_window_size_in_blocks

203 + assert val is not None

204 + return val

205 +

206 + sliding_window_groups.sort(key=_sliding_window_sort_key, reverse=True)

207

155 - self.lookup_groups = attention_groups

208 + # used by _lookup

209 + self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups)

210 + self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups

211

212 self._req_status: dict[ReqId, RequestOffloadState] = {}

213 self._current_batch_load_jobs: dict[int, TransferJob] = {}

@@ -167,8 +222,11 @@ class OffloadingConnectorScheduler:

222 self._job_counter: int = 0

223 self._jobs: dict[int, TransferJobStatus] = {}

224

170 - # block_id -> pending store job_ids. Populated only for finished

171 - # requests (running-request blocks are protected by their ref_cnt).

225 + # block_id -> pending store job_ids. Used to track jobs that needs

226 + # flushing in case a block is re-allocated by the KV cache manager.

227 + # Populated only for finished requests (running-request blocks are

228 + # protected by their ref_cnt) and for sliding window blocks (which can

229 + # be freed before a request finishes).

230 self._block_id_to_pending_jobs: dict[int, set[int]] = {}

231

232 def _generate_job_id(self) -> int:

@@ -176,10 +234,18 @@ class OffloadingConnectorScheduler:

234 self._job_counter += 1

235 return job_id

236

237 + def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None:

238 + for bid in block_ids or ():

239 + pending = self._block_id_to_pending_jobs[bid]

240 + pending.remove(job_id)

241 + if not pending:

242 + del self._block_id_to_pending_jobs[bid]

243 +

244 def _maximal_prefix_lookup(

245 self, keys: Iterable[OffloadKey], req_context: ReqContext

246 ) -> int | None:

182 - """Find the length of the maximal prefix of offloaded blocks."""

247 + """Return the number of consecutive offloaded blocks from the start,

248 + or None if the backend deferred a lookup."""

249 hit_count = 0

250 defer_lookup = False

251 for key in keys:

@@ -200,8 +266,9 @@ class OffloadingConnectorScheduler:

266 sliding_window_size: int,

267 req_context: ReqContext,

268 ) -> int | None:

203 - """Find the maximal ending position of consecutive offloaded blocks

204 - within a sliding window."""

269 + """Return the end index (in `keys`) of the last run of

270 + `sliding_window_size` consecutive hits, scanning from the end.

271 + Returns 0 on miss, None if the backend deferred a lookup."""

272 defer_lookup = False

273 consecutive_hits = 0

274 for idx in range(len(keys) - 1, -1, -1):

@@ -219,6 +286,160 @@ class OffloadingConnectorScheduler:

286 return idx + sliding_window_size if not defer_lookup else None

287 return consecutive_hits if not defer_lookup else None

288

289 + def _touch(self, req_status: RequestOffloadState):

290 + for group_config, group_state in zip(

291 + self.config.kv_group_configs, req_status.group_states

292 + ):

293 + if group_config.sliding_window_size_in_blocks is None:

294 + self.manager.touch(group_state.offload_keys)

295 + else:

296 + # we aim to keep just blocks that are necessary to hit

297 + # the original request (+ decoded blocks)

298 + blocks_to_skip = max(

299 + 0,

300 + group_state.num_hit_blocks

301 + - group_config.sliding_window_size_in_blocks,

302 + )

303 + self.manager.touch(group_state.offload_keys[blocks_to_skip:])

304 +

305 + def _lookup(self, req_status: RequestOffloadState) -> int | None:

306 + """

307 + Find how many tokens beyond num_locally_computed_tokens can be loaded.

308 +

309 + Iterates full-attention groups first (prefix lookup), then sliding-window

310 + groups (suffix lookup). Each group may tighten max_hit_size_tokens, which

311 + can invalidate an earlier group's result, so the loop re-runs when that

312 + happens until num_hit_tokens converges.

313 + """

314 + num_computed_tokens = req_status.num_locally_computed_tokens

315 + max_hit_size_tokens: int = req_status.req.num_tokens

316 + if self._sliding_window_groups:

317 + # the last prompt token has to be recomputed to get the logprobs

318 + # for sliding window attention, we must reduce by 1 to make sure

319 + # we still have a hit after reduction

320 + max_hit_size_tokens -= 1

321 + num_hit_tokens: int = 0

322 + defer_lookup = False

323 + lookup_groups = self._lookup_groups

324 + while lookup_groups:

325 + looked_up_sliding_window: bool = False

326 + groups_iter = iter(lookup_groups)

327 + lookup_groups = ()

328 + for group_idx in groups_iter:

329 + group_config: GroupOffloadConfig = self.config.kv_group_configs[

330 + group_idx

331 + ]

332 + group_state: RequestGroupState = req_status.group_states[group_idx]

333 + offloaded_block_size = group_config.offloaded_block_size

334 + offload_keys = group_state.offload_keys

335 +

336 + assert (

337 + len(offload_keys)

338 + >= req_status.req.num_tokens // offloaded_block_size

339 + )

340 +

341 + # Constrain to block-aligned boundary for this group

342 + max_hit_size_tokens = min(

343 + max_hit_size_tokens, len(offload_keys) * offloaded_block_size

344 + )

345 + if max_hit_size_tokens - num_computed_tokens < offloaded_block_size:

346 + # we can only load less than a block, better skip

347 + return 0

348 +

349 + num_blocks = min(

350 + cdiv(max_hit_size_tokens, offloaded_block_size), len(offload_keys)

351 + )

352 + start_block_idx = num_computed_tokens // offloaded_block_size

353 + offload_keys = offload_keys[start_block_idx:num_blocks]

354 + sliding_window_size_in_blocks = (

355 + group_config.sliding_window_size_in_blocks

356 + )

357 +

358 + # end index (in the sliced offload_keys) up to which we

359 + # have backend-confirmed hits

360 + num_hit_blocks: int | None

361 + if sliding_window_size_in_blocks is None:

362 + num_hit_blocks = self._maximal_prefix_lookup(

363 + offload_keys, req_status.req_context

364 + )

365 + else:

366 + num_hit_blocks = self._sliding_window_lookup(

367 + offload_keys,

368 + sliding_window_size_in_blocks,

369 + req_status.req_context,

370 + )

371 + if num_hit_blocks == 0:

372 + return 0

373 +

374 + if num_hit_blocks is None:

375 + defer_lookup = True

376 + else:

377 + max_hit_size_tokens = min(

378 + max_hit_size_tokens,

379 + offloaded_block_size * (start_block_idx + num_hit_blocks),

380 + )

381 +

382 + new_num_hit_tokens = max_hit_size_tokens - num_computed_tokens

383 + if new_num_hit_tokens < offloaded_block_size:

384 + # we can only load less than a block, better skip

385 + return 0

386 +

387 + if new_num_hit_tokens < num_hit_tokens:

388 + if defer_lookup:

389 + # make another iteration on all groups to check

390 + # if we still need to defer lookup

391 + defer_lookup = False

392 + lookup_groups = self._lookup_groups

393 + elif looked_up_sliding_window and not lookup_groups:

394 + # we need another iteration to confirm previously looked up

395 + # sliding window works with the new_num_hit_tokens

396 + lookup_groups = self._sliding_window_groups

397 +

398 + looked_up_sliding_window |= sliding_window_size_in_blocks is not None

399 + num_hit_tokens = new_num_hit_tokens

400 +

401 + if defer_lookup:

402 + logger.debug(

403 + "Offloading manager delayed request %s as backend requested",

404 + req_status.req.request_id,

405 + )

406 + return None

407 +

408 + # possibly delay request if any of the hit blocks is already being loaded

409 + if self._blocks_being_loaded:

410 + for group_config, group_state in zip(

411 + self.config.kv_group_configs, req_status.group_states

412 + ):

413 + offloaded_block_size = group_config.offloaded_block_size

414 + sliding_window_size_in_blocks = (

415 + group_config.sliding_window_size_in_blocks

416 + )

417 + offload_keys = group_state.offload_keys

418 + num_blocks = cdiv(

419 + num_computed_tokens + num_hit_tokens, offloaded_block_size

420 + )

421 + start_block_idx = num_computed_tokens // offloaded_block_size

422 + offload_keys = offload_keys[start_block_idx:num_blocks]

423 + if sliding_window_size_in_blocks is not None:

424 + offload_keys = offload_keys[-sliding_window_size_in_blocks:]

425 + if any(key in self._blocks_being_loaded for key in offload_keys):

426 + # hit blocks are being loaded, delay request

427 + logger.debug(

428 + "Delaying request %s since some of its"

429 + " blocks are already being loaded",

430 + req_status.req.request_id,

431 + )

432 + return None

433 +

434 + logger.debug(

435 + "Request %s hit %s offloaded tokens after %s GPU hit tokens",

436 + req_status.req.request_id,

437 + num_hit_tokens,

438 + num_computed_tokens,

439 + )

440 +

441 + return num_hit_tokens

442 +

443 def get_num_new_matched_tokens(

444 self, request: Request, num_computed_tokens: int

445 ) -> tuple[int | None, bool]:

@@ -241,96 +462,28 @@ class OffloadingConnectorScheduler:

462 - `True` if tokens will be loaded asynchronously

463 (between scheduler steps).

464 """

465 + is_new_request = False

466 if req_status := self._req_status.get(request.request_id):

467 # make sure block IDs are cleared

468 for group_state in req_status.group_states:

469 group_state.block_ids.clear()

470 else:

471 + is_new_request = True

472 req_status = RequestOffloadState(config=self.config, req=request)

473 self._req_status[request.request_id] = req_status

474

475 req_status.update_offload_keys()

476 req_status.num_locally_computed_tokens = num_computed_tokens

477

255 - for gs in req_status.group_states:

256 - self.manager.touch(gs.offload_keys)

257 -

258 - # Start with the full request size as the maximum loadable

259 - max_hit_size_tokens: int = req_status.req.num_tokens

260 - num_hit_tokens: int = 0

261 - defer_lookup = False

262 - delay_request = False

263 - for group_idx in self.lookup_groups:

264 - group_config: GroupOffloadConfig = self.config.kv_group_configs[group_idx]

265 - offloaded_block_size = group_config.offloaded_block_size

266 - offload_keys = req_status.group_states[group_idx].offload_keys

267 -

268 - num_blocks = max_hit_size_tokens // offloaded_block_size

269 - assert len(offload_keys) >= num_blocks

270 -

271 - # Constrain to block-aligned boundary for this group

272 - max_hit_size_tokens = num_blocks * offloaded_block_size

273 - num_hit_tokens = max_hit_size_tokens - num_computed_tokens

274 - if num_hit_tokens < offloaded_block_size:

275 - # we can only load less than a block, better skip

276 - return 0, False

277 -

278 - start_block_idx = num_computed_tokens // offloaded_block_size

279 - offload_keys = offload_keys[start_block_idx:num_blocks]

280 - # Full attention relies on all previous KV cache blocks.

281 - # Thus, we search for a maximal prefix of KV cache which are all cached.

282 - block_hits = self._maximal_prefix_lookup(

283 - offload_keys, req_status.req_context

478 + num_hit_tokens = self._lookup(req_status)

479 + if is_new_request:

480 + req_status.update_num_hit_blocks(

481 + num_computed_tokens + (num_hit_tokens or 0)

482 )

285 - if block_hits == 0:

286 - return 0, False

483

288 - if block_hits is None:

289 - defer_lookup = True

290 - else:

291 - # Further constrain based on what's actually available by backend

292 - max_hit_size_tokens = offloaded_block_size * (

293 - start_block_idx + block_hits

294 - )

484 + self._touch(req_status)

485

296 - num_hit_tokens = max_hit_size_tokens - num_computed_tokens

297 - if num_hit_tokens < offloaded_block_size:

298 - # we can only load less than a block, better skip

299 - return 0, False

300 -

301 - if (

302 - block_hits

303 - and self._blocks_being_loaded

304 - and any(

305 - key in self._blocks_being_loaded

306 - for key in offload_keys[:block_hits]

307 - )

308 - ):

309 - # hit blocks are being loaded, delay request

310 - delay_request = True

311 -

312 - if defer_lookup:

313 - logger.debug(

314 - "Offloading manager delayed request %s as backend requested",

315 - req_status.req.request_id,

316 - )

317 - return None, False

318 -

319 - if delay_request:

320 - logger.debug(

321 - "Delaying request %s since some of its blocks are already being loaded",

322 - req_status.req.request_id,

323 - )

324 - return None, False

325 -

326 - logger.debug(

327 - "Request %s hit %s offloaded tokens after %s GPU hit tokens",

328 - request.request_id,

329 - num_hit_tokens,

330 - num_computed_tokens,

331 - )

332 -

333 - return num_hit_tokens, True

486 + return num_hit_tokens, bool(num_hit_tokens)

487

488 def update_state_after_alloc(

489 self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int

@@ -375,6 +528,13 @@ class OffloadingConnectorScheduler:

528 )

529 num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks

530

531 + if group_config.sliding_window_size_in_blocks is not None:

532 + assert (

533 + num_pending_gpu_blocks

534 + <= group_config.sliding_window_size_in_blocks

535 + * self.config.block_size_factor

536 + )

537 +

538 num_blocks = cdiv(num_cached_tokens, offloaded_block_size)

539 assert len(offload_keys) >= num_blocks

540 if num_pending_gpu_blocks:

@@ -513,17 +673,21 @@ class OffloadingConnectorScheduler:

673 req_status.advance_stored_idx(num_offloadable_tokens)

674 continue

675

516 - for group_state in req_status.group_states:

517 - self.manager.touch(group_state.offload_keys)

676 + self._touch(req_status)

677

678 keys_to_store = set(store_output.keys_to_store)

679

680 group_sizes: list[int] = []

681 block_indices: list[int] = []

682 src_block_ids: list[int] = []

683 + sliding_window_block_ids: list[int] = []

684 + non_sliding_window_block_ids: list[int] = []

685 for group_config, group_state in zip(

686 self.config.kv_group_configs, req_status.group_states

687 ):

688 + is_sliding_window = (

689 + group_config.sliding_window_size_in_blocks is not None

690 + )

691 num_blocks = num_offloadable_tokens // group_config.offloaded_block_size

692 start_block_idx = group_state.next_stored_block_idx

693 block_ids = group_state.block_ids

@@ -547,6 +711,11 @@ class OffloadingConnectorScheduler:

711 elif start_gpu_block_idx is None:

712 start_gpu_block_idx = gpu_block_idx + i

713 src_block_ids.append(block_id)

714 + if is_sliding_window:

715 + sliding_window_block_ids.append(block_id)

716 + else:

717 + non_sliding_window_block_ids.append(block_id)

718 +

719 group_sizes.append(num_group_blocks)

720 block_indices.append(start_gpu_block_idx or 0)

721 group_state.next_stored_block_idx = num_blocks

@@ -562,12 +731,21 @@ class OffloadingConnectorScheduler:

731 any_jid = next(iter(req_status.transfer_jobs))

732 assert self._jobs[any_jid].is_store

733 req_status.transfer_jobs.add(job_id)

734 +

735 + # Watch sliding window blocks as they may get evicted

736 + # before the request finishes

737 + for bid in sliding_window_block_ids or ():

738 + self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)

739 +

740 + # the non-sliding window blocks will be watched only

741 + # when the request finishes

742 self._jobs[job_id] = TransferJobStatus(

743 req_id=req_id,

744 pending_count=self.config.num_workers,

745 keys=set(keys_to_store),

746 is_store=True,

570 - gpu_block_ids=src_block_ids,

747 + non_sliding_window_block_ids=non_sliding_window_block_ids,

748 + sliding_window_block_ids=sliding_window_block_ids or None,

749 )

750

751 store_jobs[job_id] = TransferJob(

@@ -632,12 +810,16 @@ class OffloadingConnectorScheduler:

810 self._blocks_being_loaded.difference_update(job_status.keys)

811

812 req_status = self._req_status[job_status.req_id]

635 - if self._block_id_to_pending_jobs and req_status.req.is_finished():

636 - for bid in job_status.gpu_block_ids or ():

637 - pending = self._block_id_to_pending_jobs[bid]

638 - pending.remove(job_id)

639 - if not pending:

640 - del self._block_id_to_pending_jobs[bid]

813 + if self._block_id_to_pending_jobs:

814 + # Sliding window blocks are tracked from store creation

815 + # and must be cleaned up unconditionally.

816 + self._remove_pending_job(job_id, job_status.sliding_window_block_ids)

817 + # Non-sliding-window blocks are only tracked after

818 + # request_finished, so only clean up for finished requests.

819 + if req_status.req.is_finished():

820 + self._remove_pending_job(

821 + job_id, job_status.non_sliding_window_block_ids

822 + )

823

824 del self._jobs[job_id]

825 req_status.transfer_jobs.remove(job_id)

@@ -671,7 +853,7 @@ class OffloadingConnectorScheduler:

853 # Register them so future block reuse triggers a flush.

854 for job_id in req_status.transfer_jobs:

855 job_status = self._jobs[job_id]

674 - for bid in job_status.gpu_block_ids or ():

856 + for bid in job_status.non_sliding_window_block_ids or ():

857 self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)

858 return False, None

859

[kv_offload+HMA][12/N]: Scheduler-side support for sliding window groups (#41228)