@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b029de9902

Fork 0

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 14, 05:30 UTC

[Optimization] Make new_block_ids None if empty (#23262)

Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>

Woosuk Kwon committed Aug 20, 2025 at 18:25 UTC b029de9902aa3ac58806c8c17776c7074175b6db

5 files changed +57 -27

vllm/v1/core/kv_cache_manager.py

+26 -4

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from dataclasses import dataclass

5 -from typing import Optional

5 +from typing import Literal, Optional, overload

7 from vllm.distributed.kv_events import KVCacheEvent

8 from vllm.logger import init_logger

@@ -37,7 +37,24 @@ class KVCacheBlocks:

37 tuple(blk1 + blk2

38 for blk1, blk2 in zip(self.blocks, other.blocks)))

40 - def get_block_ids(self) -> tuple[list[int], ...]:

40 + @overload

41 + def get_block_ids(

42 + self,

43 + allow_none: Literal[False] = False,

44 + ) -> tuple[list[int], ...]:

45 + ...

46 +

47 + @overload

48 + def get_block_ids(

49 + self,

50 + allow_none: Literal[True] = True,

51 + ) -> Optional[tuple[list[int], ...]]:

52 + ...

53 +

54 + def get_block_ids(

55 + self,

56 + allow_none: bool = False,

57 + ):

58 """

59 Converts the KVCacheBlocks instance to block_ids.

@@ -46,6 +63,8 @@ class KVCacheBlocks:

63 * the outer tuple corresponds to KV cache groups

64 * each inner list contains the block_ids of the blocks in that group

65 """

66 + if allow_none and all(len(group) == 0 for group in self.blocks):

67 + return None

68 return tuple([blk.block_id for blk in group] for group in self.blocks)

70 def get_unhashed_block_ids(self) -> list[int]:

@@ -348,10 +367,13 @@ class KVCacheManager:

367 """

368 return self.block_pool.take_events()

369

370 + def get_blocks(self, request_id: str) -> KVCacheBlocks:

371 + """Get the blocks of a request."""

372 + return KVCacheBlocks(self.coordinator.get_blocks(request_id))

373 +

374 def get_block_ids(self, request_id: str) -> tuple[list[int], ...]:

375 """Get the block ids of a request."""

353 - return KVCacheBlocks(

354 - self.coordinator.get_blocks(request_id)).get_block_ids()

376 + return self.get_blocks(request_id).get_block_ids()

377

378 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:

379 """Cache the blocks for the request, if enabled."""

vllm/v1/core/sched/output.py

+1 -1

@@ -91,7 +91,7 @@ class CachedRequestData:

91 # NOTE(woosuk): new_token_ids is only used for pipeline parallelism.

92 # When PP is not used, new_token_ids will be empty.

93 new_token_ids: list[list[int]]

94 - new_block_ids: list[tuple[list[int], ...]]

94 + new_block_ids: list[Optional[tuple[list[int], ...]]]

95 num_computed_tokens: list[int]

97 @property

vllm/v1/core/sched/scheduler.py

+12 -12

@@ -19,7 +19,7 @@ from vllm.logger import init_logger

19 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

20 from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,

21 compute_encoder_budget)

22 -from vllm.v1.core.kv_cache_manager import KVCacheManager

22 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager

23 from vllm.v1.core.sched.interface import SchedulerInterface

24 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,

25 SchedulerOutput)

@@ -185,7 +185,7 @@ class Scheduler(SchedulerInterface):

185 # uses structured decoding.

186 structured_output_request_ids: dict[str, int] = {}

187

188 - req_to_new_block_ids: dict[str, tuple[list[int], ...]] = {}

188 + req_to_new_blocks: dict[str, KVCacheBlocks] = {}

189 num_scheduled_tokens: dict[str, int] = {}

190 token_budget = self.max_num_scheduled_tokens

191 # Encoder-related.

@@ -288,8 +288,7 @@ class Scheduler(SchedulerInterface):

288 # Therefore, we might introduce some additional

289 # cycle to fill in the bitmask, which could be a big no-op.

290 structured_output_request_ids[request.request_id] = req_index

291 - req_to_new_block_ids[request.request_id] = (

292 - new_blocks.get_block_ids())

291 + req_to_new_blocks[request.request_id] = new_blocks

292 num_scheduled_tokens[request.request_id] = num_new_tokens

293 token_budget -= num_new_tokens

294 req_index += 1

@@ -496,8 +495,8 @@ class Scheduler(SchedulerInterface):

495

496 if self.lora_config and request.lora_request:

497 scheduled_loras.add(request.lora_request.lora_int_id)

499 - req_to_new_block_ids[request.request_id] = (

500 - self.kv_cache_manager.get_block_ids(request.request_id))

498 + req_to_new_blocks[request.request_id] = (

499 + self.kv_cache_manager.get_blocks(request.request_id))

500 num_scheduled_tokens[request.request_id] = num_new_tokens

501 token_budget -= num_new_tokens

502 request.status = RequestStatus.RUNNING

@@ -546,8 +545,8 @@ class Scheduler(SchedulerInterface):

545 )

546 # Construct the scheduler output.

547 new_reqs_data = [

549 - NewRequestData.from_request(req,

550 - req_to_new_block_ids[req.request_id])

548 + NewRequestData.from_request(

549 + req, req_to_new_blocks[req.request_id].get_block_ids())

550 for req in scheduled_new_reqs

551 ]

552 cached_reqs_data = self._make_cached_request_data(

@@ -555,7 +554,7 @@ class Scheduler(SchedulerInterface):

554 scheduled_resumed_reqs,

555 num_scheduled_tokens,

556 scheduled_spec_decode_tokens,

558 - req_to_new_block_ids,

557 + req_to_new_blocks,

558 )

559 scheduler_output = SchedulerOutput(

560 scheduled_new_reqs=new_reqs_data,

@@ -628,11 +627,11 @@ class Scheduler(SchedulerInterface):

627 resumed_reqs: list[Request],

628 num_scheduled_tokens: dict[str, int],

629 spec_decode_tokens: dict[str, list[int]],

631 - req_to_new_block_ids: dict[str, tuple[list[int], ...]],

630 + req_to_new_blocks: dict[str, KVCacheBlocks],

631 ) -> CachedRequestData:

632 req_ids: list[str] = []

633 new_token_ids: list[list[int]] = []

635 - new_block_ids: list[tuple[list[int], ...]] = []

634 + new_block_ids: list[Optional[tuple[list[int], ...]]] = []

635 num_computed_tokens: list[int] = []

636

637 use_connector = self.connector is not None

@@ -655,7 +654,8 @@ class Scheduler(SchedulerInterface):

654 # out of bounds errors. TODO: Remove this once the KVConnector

655 # is updated to handle token IDs properly.

656 new_token_ids.append([])

658 - new_block_ids.append(req_to_new_block_ids[req_id])

657 + new_block_ids.append(

658 + req_to_new_blocks[req_id].get_block_ids(allow_none=True))

659 num_computed_tokens.append(req.num_computed_tokens)

660 # Because resumed_reqs is usually empty, it is more efficient to do

661 # in-place appending so that we don't need to allocate a new list.

vllm/v1/worker/gpu_model_runner.py

+9 -5

@@ -574,11 +574,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

574

575 # Update the block IDs.

576 if not resumed_from_preemption:

577 - # Append the new blocks to the existing block IDs.

578 - for block_ids, new_ids in zip(req_state.block_ids,

579 - new_block_ids):

580 - block_ids.extend(new_ids)

577 + if new_block_ids is not None:

578 + # Append the new blocks to the existing block IDs.

579 + for block_ids, new_ids in zip(req_state.block_ids,

580 + new_block_ids):

581 + block_ids.extend(new_ids)

582 else:

583 + assert new_block_ids is not None

584 # The request is resumed from preemption.

585 # Replace the existing block IDs with the new ones.

586 req_state.block_ids = new_block_ids

@@ -594,7 +596,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

596 # Update the persistent batch.

597 self.input_batch.num_computed_tokens_cpu[req_index] = (

598 num_computed_tokens)

597 - self.input_batch.block_table.append_row(new_block_ids, req_index)

599 + if new_block_ids is not None:

600 + self.input_batch.block_table.append_row(

601 + new_block_ids, req_index)

602

603 # For the last rank, we don't need to update the token_ids_cpu

604 # because the sampled tokens are already cached.

vllm/v1/worker/tpu_model_runner.py

+9 -5

@@ -418,11 +418,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

418 # Update the cached states.

419 req_state.num_computed_tokens = num_computed_tokens

420 if not resumed_from_preemption:

421 - # Append the new blocks to the existing block IDs.

422 - for block_ids, new_ids in zip(req_state.block_ids,

423 - new_block_ids):

424 - block_ids.extend(new_ids)

421 + if new_block_ids is not None:

422 + # Append the new blocks to the existing block IDs.

423 + for block_ids, new_ids in zip(req_state.block_ids,

424 + new_block_ids):

425 + block_ids.extend(new_ids)

426 else:

427 + assert new_block_ids is not None

428 # The request is resumed from preemption.

429 # Replace the existing block IDs with the new ones.

430 req_state.block_ids = new_block_ids

@@ -438,7 +440,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

440 # Update the persistent batch.

441 self.input_batch.num_computed_tokens_cpu[req_index] = (

442 num_computed_tokens)

441 - self.input_batch.block_table.append_row(new_block_ids, req_index)

443 + if new_block_ids is not None:

444 + self.input_batch.block_table.append_row(

445 + new_block_ids, req_index)

446

447 # Add the new or resumed requests to the persistent batch.

448 # The smaller empty indices are filled first.